From 8b2adead78c25142f98da7d7163e7fe5e24e09b2 Mon Sep 17 00:00:00 2001
From: 0z1-ghb <162235745+0z1-ghb@users.noreply.github.com>
Date: Wed, 20 May 2026 01:26:35 +0300
Subject: [PATCH] =?UTF-8?q?fix(compressor):=20ABC=20compliance=20=E2=80=94?=
 =?UTF-8?q?=20total=5Ftokens,=20api=5Fmode,=20logger=20consistency?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 agent/agent_init.py                           |  1 +
 agent/context_compressor.py                   |  7 +++---
 agent/context_engine.py                       |  1 +
 tests/agent/test_last_total_tokens.py         | 22 +++++++++++++++++++
 .../test_plugin_context_engine_init.py        |  3 +--
 5 files changed, 29 insertions(+), 5 deletions(-)
 create mode 100644 tests/agent/test_last_total_tokens.py

diff --git a/agent/agent_init.py b/agent/agent_init.py
index 00e90edd295..0c3cff066e2 100644
--- a/agent/agent_init.py
+++ b/agent/agent_init.py
@@ -1429,6 +1429,7 @@ def init_agent(
             base_url=agent.base_url,
             api_key=getattr(agent, "api_key", ""),
             provider=agent.provider,
+            api_mode=agent.api_mode,
         )
         if not agent.quiet_mode:
             _ra().logger.info("Using context engine: %s", _selected_engine.name)
diff --git a/agent/context_compressor.py b/agent/context_compressor.py
index 62636809094..49907e2c331 100644
--- a/agent/context_compressor.py
+++ b/agent/context_compressor.py
@@ -609,6 +609,7 @@ class ContextCompressor(ContextEngine):
         """Update tracked token usage from API response."""
         self.last_prompt_tokens = usage.get("prompt_tokens", 0)
         self.last_completion_tokens = usage.get("completion_tokens", 0)
+        self.last_total_tokens = usage.get("total_tokens", self.last_prompt_tokens + self.last_completion_tokens)
 
     def should_compress(self, prompt_tokens: int = None) -> bool:
         """Check if context exceeds the compression threshold.
@@ -897,7 +898,7 @@ class ContextCompressor(ContextEngine):
         into the warning log.
         """
         self._summary_model_fallen_back = True
-        logging.warning(
+        logger.warning(
             "Summary model '%s' %s (%s). "
             "Falling back to main model '%s' for compression.",
             self.summary_model, reason, e, self.model,
@@ -1086,7 +1087,7 @@ The user has requested that this compaction PRIORITISE preserving all informatio
             # No provider configured — long cooldown, unlikely to self-resolve
             self._summary_failure_cooldown_until = time.monotonic() + _SUMMARY_FAILURE_COOLDOWN_SECONDS
             self._last_summary_error = "no auxiliary LLM provider configured"
-            logging.warning("Context compression: no provider available for "
+            logger.warning("Context compression: no provider available for "
                             "summary. Middle turns will be dropped without summary "
                             "for %d seconds.",
                             _SUMMARY_FAILURE_COOLDOWN_SECONDS)
@@ -1182,7 +1183,7 @@ The user has requested that this compaction PRIORITISE preserving all informatio
             if len(err_text) > 220:
                 err_text = err_text[:217].rstrip() + "..."
             self._last_summary_error = err_text
-            logging.warning(
+            logger.warning(
                 "Failed to generate context summary: %s. "
                 "Further summary attempts paused for %d seconds.",
                 e,
diff --git a/agent/context_engine.py b/agent/context_engine.py
index 2947da54d8c..c30a7a84752 100644
--- a/agent/context_engine.py
+++ b/agent/context_engine.py
@@ -200,6 +200,7 @@ class ContextEngine(ABC):
         base_url: str = "",
         api_key: str = "",
         provider: str = "",
+        api_mode: str = "",
     ) -> None:
         """Called when the user switches models or on fallback activation.
 
diff --git a/tests/agent/test_last_total_tokens.py b/tests/agent/test_last_total_tokens.py
new file mode 100644
index 00000000000..ed4735ae253
--- /dev/null
+++ b/tests/agent/test_last_total_tokens.py
@@ -0,0 +1,22 @@
+"""Test that last_total_tokens is correctly set by ContextCompressor."""
+
+from agent.context_compressor import ContextCompressor
+
+
+def test_update_from_response_sets_total_tokens():
+    """ABC contract: last_total_tokens must be set from API response."""
+    c = ContextCompressor(model="test", quiet_mode=True, config_context_length=200000)
+
+    c.update_from_response({"prompt_tokens": 100, "completion_tokens": 30, "total_tokens": 130})
+    assert c.last_total_tokens == 130
+
+    c.update_from_response({"prompt_tokens": 100, "completion_tokens": 30})
+    assert c.last_total_tokens == 130
+
+
+def test_session_reset_clears_total_tokens():
+    """on_session_reset must zero total_tokens."""
+    c = ContextCompressor(model="test", quiet_mode=True, config_context_length=200000)
+    c.update_from_response({"prompt_tokens": 100, "completion_tokens": 30, "total_tokens": 130})
+    c.on_session_reset()
+    assert c.last_total_tokens == 0
diff --git a/tests/run_agent/test_plugin_context_engine_init.py b/tests/run_agent/test_plugin_context_engine_init.py
index 60e89889088..83895ac6dce 100644
--- a/tests/run_agent/test_plugin_context_engine_init.py
+++ b/tests/run_agent/test_plugin_context_engine_init.py
@@ -87,5 +87,4 @@ def test_plugin_engine_update_model_args():
     assert kw["context_length"] == 131_072
     assert "model" in kw
     assert "provider" in kw
-    # Should NOT pass api_mode — the ABC doesn't accept it
-    assert "api_mode" not in kw
+    assert "api_mode" in kw