fix(compression): separate provider-exact vs projected token state

agent/context_compressor.py

@@ -27,6 +27,7 @@
 from agent.context_engine import ContextEngine
 from agent.model_metadata import (
     MINIMUM_CONTEXT_LENGTH,
+    estimate_request_tokens_rough,
     get_model_context_length,
     estimate_messages_tokens_rough,
 )
@@ -569,7 +570,10 @@ def __init__(
             )
         self._context_probed = False  # True after a step-down from context error
 
-        self.last_prompt_tokens = 0
+        self.last_provider_prompt_tokens = 0
+        self.projected_prompt_tokens = 0
+        self.projected_prompt_tokens_source = "none"
+        self._transcript_mutated_since_api = False
         self.last_completion_tokens = 0
 
         self.summary_model = summary_model_override or ""
@@ -593,11 +597,56 @@ def __init__(
         self._last_aux_model_failure_error: Optional[str] = None
         self._last_aux_model_failure_model: Optional[str] = None
 
+    @property
+    def last_prompt_tokens(self) -> int:
+        """Backward-compatible alias for projected context pressure."""
+        return self.projected_prompt_tokens
+
+    @last_prompt_tokens.setter
+    def last_prompt_tokens(self, value: int) -> None:
+        self.projected_prompt_tokens = int(value or 0)
+
     def update_from_response(self, usage: Dict[str, Any]):
         """Update tracked token usage from API response."""
-        self.last_prompt_tokens = usage.get("prompt_tokens", 0)
+        prompt_tokens = usage.get("prompt_tokens", 0)
+        self.last_provider_prompt_tokens = prompt_tokens
+        self.projected_prompt_tokens = prompt_tokens
+        self.projected_prompt_tokens_source = "provider_exact"
+        self._transcript_mutated_since_api = False
         self.last_completion_tokens = usage.get("completion_tokens", 0)
 
+    def _mark_transcript_dirty(self, reason: str = "") -> None:
+        """Mark the tracked request projection as stale after local mutation."""
+        self._transcript_mutated_since_api = True
+        if reason:
+            logger.debug("Context projection marked dirty: %s", reason)
+
+    def get_current_request_pressure(
+        self,
+        messages: List[Dict[str, Any]],
+        system_prompt: str,
+        tools: List[Dict[str, Any]],
+    ) -> tuple[int, str]:
+        """Return current request pressure using exact state when still valid."""
+        if self.last_provider_prompt_tokens > 0 and not self._transcript_mutated_since_api:
+            self.projected_prompt_tokens = self.last_provider_prompt_tokens
+            self.projected_prompt_tokens_source = "provider_exact"
+            return self.last_provider_prompt_tokens, "provider_exact"
+
+        est = estimate_request_tokens_rough(
+            messages,
+            system_prompt=system_prompt,
+            tools=tools,
+        )
+        source = (
+            self.projected_prompt_tokens_source
+            if self.projected_prompt_tokens_source == "estimated_post_compression"
+            else "estimated"
+        )
+        self.projected_prompt_tokens = est
+        self.projected_prompt_tokens_source = source
+        return est, source
+
     def should_compress(self, prompt_tokens: int = None) -> bool:
         """Check if context exceeds the compression threshold.
 
@@ -620,6 +669,17 @@ def should_compress(self, prompt_tokens: int = None) -> bool:
             return False
         return True
 
+    def on_session_reset(self) -> None:
+        """Reset token tracking and compression metadata for a fresh session."""
+        super().on_session_reset()
+        self.last_provider_prompt_tokens = 0
+        self.projected_prompt_tokens = 0
+        self.projected_prompt_tokens_source = "none"
+        self._transcript_mutated_since_api = False
+        self._context_probed = False
+        self._context_probe_persistable = False
+        self._previous_summary = None
+
     # ------------------------------------------------------------------
     # Tool output pruning (cheap pre-pass, no LLM call)
     # ------------------------------------------------------------------

agent/conversation_compression.py

@@ -415,7 +415,9 @@ def compress_context(
         system_prompt=new_system_prompt or "",
         tools=agent.tools or None,
     )
-    agent.context_compressor.last_prompt_tokens = _compressed_est
+    agent.context_compressor.projected_prompt_tokens = _compressed_est
+    agent.context_compressor.projected_prompt_tokens_source = "estimated_post_compression"
+    agent.context_compressor._mark_transcript_dirty("post_compression")
     agent.context_compressor.last_completion_tokens = 0
 
     # Clear the file-read dedup cache.  After compression the original

agent/conversation_loop.py

@@ -396,6 +396,7 @@ def run_conversation(
     # Add user message
     user_msg = {"role": "user", "content": user_message}
     messages.append(user_msg)
+    agent._mark_context_pressure_dirty("user_message_appended")
     current_turn_user_idx = len(messages) - 1
     agent._persist_user_message_idx = current_turn_user_idx
 
@@ -1784,6 +1785,7 @@ def _stop_spinner():
                                 active_system_prompt = _sanitized_system
                                 agent._cached_system_prompt = _sanitized_system
                                 _system_sanitized = True
+                                agent._mark_context_pressure_dirty("active_system_prompt_swapped")
                         if isinstance(getattr(agent, "ephemeral_system_prompt", None), str):
                             _sanitized_ephemeral = _strip_non_ascii(agent.ephemeral_system_prompt)
                             if _sanitized_ephemeral != agent.ephemeral_system_prompt:
@@ -3133,6 +3135,7 @@ def _stop_spinner():
 
                     assistant_msg = agent._build_assistant_message(assistant_message, finish_reason)
                     messages.append(assistant_msg)
+                    agent._mark_context_pressure_dirty("assistant_turn_committed")
                     for tc in assistant_message.tool_calls:
                         if tc.function.name not in agent.valid_tool_names:
                             content = f"Tool '{tc.function.name}' does not exist. Available tools: {available}"
@@ -3217,6 +3220,7 @@ def _stop_spinner():
                         # Append the assistant message with its (broken) tool_calls
                         recovery_assistant = agent._build_assistant_message(assistant_message, finish_reason)
                         messages.append(recovery_assistant)
+                        agent._mark_context_pressure_dirty("assistant_turn_committed")
 
                         # Respond with tool error results for each tool call
                         invalid_names = {name for name, _ in invalid_json_args}
@@ -3305,6 +3309,7 @@ def _stop_spinner():
                 agent._post_tool_empty_retried = False
 
                 messages.append(assistant_msg)
+                agent._mark_context_pressure_dirty("assistant_turn_committed")
                 agent._emit_interim_assistant_message(assistant_msg)
 
                 # Close any open streaming display (response box, reasoning
@@ -3320,6 +3325,7 @@ def _stop_spinner():
                         pass
 
                 agent._execute_tool_calls(assistant_message, messages, effective_task_id, api_call_count)
+                agent._mark_context_pressure_dirty("tool_execution_batch_completed")
 
                 if agent._tool_guardrail_halt_decision is not None:
                     decision = agent._tool_guardrail_halt_decision
@@ -3366,21 +3372,11 @@ def _stop_spinner():
                 # a session can grow unbounded after disconnects because
                 # should_compress(0) never fires.  (#2153)
                 _compressor = agent.context_compressor
-                if _compressor.last_prompt_tokens > 0:
-                    # Only use prompt_tokens — completion/reasoning
-                    # tokens don't consume context window space.
-                    # Thinking models (GLM-5.1, QwQ, DeepSeek R1)
-                    # inflate completion_tokens with reasoning,
-                    # causing premature compression.  (#12026)
-                    _real_tokens = _compressor.last_prompt_tokens
-                else:
-                    # Include tool schemas — with 50+ tools enabled
-                    # these add 20-30K tokens the messages-only
-                    # estimate misses, which can skip compression
-                    # past the configured threshold (#14695).
-                    _real_tokens = estimate_request_tokens_rough(
-                        messages, tools=agent.tools or None
-                    )
+                _real_tokens, _tokens_source = _compressor.get_current_request_pressure(
+                    messages=messages,
+                    system_prompt=active_system_prompt or "",
+                    tools=agent.tools or [],
+                )
 
                 if agent.compression_enabled and _compressor.should_compress(_real_tokens):
                     agent._safe_print("  ⟳ compacting context…")
@@ -3727,6 +3723,7 @@ def _stop_spinner():
                     messages.pop()
 
                 messages.append(final_msg)
+                agent._mark_context_pressure_dirty("assistant_turn_committed")
 
                 _turn_exit_reason = f"text_response(finish_reason={finish_reason})"
                 if not agent.quiet_mode:

run_agent.py

@@ -3710,6 +3710,12 @@ def _should_sanitize_tool_calls(self) -> bool:
         """
         return self.api_mode != "codex_responses"
 
+    def _mark_context_pressure_dirty(self, reason: str) -> None:
+        """Invalidate provider-exact prompt-token state after local mutation."""
+        _compressor = getattr(self, "context_compressor", None)
+        if _compressor and hasattr(_compressor, "_mark_transcript_dirty"):
+            _compressor._mark_transcript_dirty(reason)
+
     def _compress_context(self, messages: list, system_message: str, *, approx_tokens: int = None, task_id: str = "default", focus_topic: str = None) -> tuple:
         """Forwarder — see ``agent.conversation_compression.compress_context``."""
         from agent.conversation_compression import compress_context

tests/agent/test_context_compressor.py

@@ -47,11 +47,68 @@ def test_updates_fields(self, compressor):
             "total_tokens": 6000,
         })
         assert compressor.last_prompt_tokens == 5000
+        assert compressor.last_provider_prompt_tokens == 5000
+        assert compressor.projected_prompt_tokens == 5000
+        assert compressor.projected_prompt_tokens_source == "provider_exact"
+        assert compressor._transcript_mutated_since_api is False
         assert compressor.last_completion_tokens == 1000
 
     def test_missing_fields_default_zero(self, compressor):
         compressor.update_from_response({})
         assert compressor.last_prompt_tokens == 0
+        assert compressor.projected_prompt_tokens_source == "provider_exact"
+        assert compressor._transcript_mutated_since_api is False
+
+
+class TestProjectedPressure:
+    def test_last_prompt_tokens_aliases_projected_prompt_tokens(self, compressor):
+        compressor.last_prompt_tokens = 1234
+        assert compressor.projected_prompt_tokens == 1234
+        assert compressor.last_prompt_tokens == 1234
+
+    def test_get_current_request_pressure_uses_provider_exact_when_clean(self, compressor):
+        compressor.update_from_response({"prompt_tokens": 777, "completion_tokens": 0})
+
+        tokens, source = compressor.get_current_request_pressure(
+            messages=[{"role": "user", "content": "hello"}],
+            system_prompt="system",
+            tools=[],
+        )
+
+        assert (tokens, source) == (777, "provider_exact")
+        assert compressor.projected_prompt_tokens == 777
+        assert compressor.projected_prompt_tokens_source == "provider_exact"
+
+    def test_get_current_request_pressure_reestimates_when_dirty(self, compressor):
+        compressor.update_from_response({"prompt_tokens": 100, "completion_tokens": 0})
+        compressor.projected_prompt_tokens_source = "estimated_post_compression"
+        compressor._mark_transcript_dirty("tool_batch")
+
+        tokens, source = compressor.get_current_request_pressure(
+            messages=[{"role": "user", "content": "hello"}],
+            system_prompt="system prompt",
+            tools=[{"type": "function", "function": {"name": "tool"}}],
+        )
+
+        assert source == "estimated_post_compression"
+        assert tokens == compressor.projected_prompt_tokens
+        assert tokens != compressor.last_provider_prompt_tokens
+
+    def test_get_current_request_pressure_includes_tools_schema(self, compressor):
+        compressor._mark_transcript_dirty("post_compression")
+
+        tokens_without_tools, _ = compressor.get_current_request_pressure(
+            messages=[{"role": "user", "content": "hello"}],
+            system_prompt="system",
+            tools=[],
+        )
+        tokens_with_tools, _ = compressor.get_current_request_pressure(
+            messages=[{"role": "user", "content": "hello"}],
+            system_prompt="system",
+            tools=[{"type": "function", "function": {"name": "x", "description": "tool schema padding" * 50}}],
+        )
+
+        assert tokens_with_tools > tokens_without_tools
 
 
 

tests/agent/test_context_engine.py

@@ -174,6 +174,9 @@ class TestCompressorSessionReset:
     def test_reset_clears_state(self):
         c = ContextCompressor(model="test", quiet_mode=True, config_context_length=200000)
         c.last_prompt_tokens = 50000
+        c.last_provider_prompt_tokens = 49000
+        c.projected_prompt_tokens_source = "estimated"
+        c._transcript_mutated_since_api = True
         c.compression_count = 3
         c._previous_summary = "some old summary"
         c._context_probed = True
@@ -182,6 +185,9 @@ def test_reset_clears_state(self):
         c.on_session_reset()
 
         assert c.last_prompt_tokens == 0
+        assert c.last_provider_prompt_tokens == 0
+        assert c.projected_prompt_tokens_source == "none"
+        assert c._transcript_mutated_since_api is False
         assert c.last_completion_tokens == 0
         assert c.last_total_tokens == 0
         assert c.compression_count == 0

tests/run_agent/test_projected_token_state.py

@@ -0,0 +1,92 @@
+"""Regression tests for exact vs projected compression token tracking."""
+
+import sys
+import types
+from pathlib import Path
+
+sys.modules.setdefault("fire", types.SimpleNamespace(Fire=lambda *a, **k: None))
+sys.modules.setdefault("firecrawl", types.SimpleNamespace(Firecrawl=object))
+sys.modules.setdefault("fal_client", types.SimpleNamespace())
+
+from unittest.mock import patch
+
+import run_agent
+from agent.context_compressor import ContextCompressor
+
+
+def _make_agent_and_compressor():
+    with patch("agent.context_compressor.get_model_context_length", return_value=100000):
+        compressor = ContextCompressor(model="test/model", quiet_mode=True)
+
+    compressor.compress = lambda messages, current_tokens=None, focus_topic=None: [
+        {"role": "assistant", "content": "compressed"}
+    ]
+
+    agent = run_agent.AIAgent.__new__(run_agent.AIAgent)
+    agent.session_id = "session-1"
+    agent.model = "test/model"
+    agent.platform = ""
+    agent.context_compressor = compressor
+    agent._memory_manager = None
+    agent._todo_store = types.SimpleNamespace(format_for_injection=lambda: "")
+    agent._invalidate_system_prompt = lambda: None
+    agent._build_system_prompt = lambda system_message: f"sys::{system_message}"
+    agent._cached_system_prompt = None
+    agent._session_db = None
+    agent._session_db_created = False
+    agent._last_flushed_db_idx = 0
+    agent._last_compression_summary_warning = None
+    agent._last_aux_fallback_warning_key = None
+    agent._emit_warning = lambda *a, **k: None
+    agent._vprint = lambda *a, **k: None
+    agent.status_callback = None
+    agent.commit_memory_session = lambda messages: None
+    agent.logs_dir = Path("/tmp")
+    agent.session_log_file = agent.logs_dir / "session.json"
+    agent.tools = [{"type": "function", "function": {"name": "tool", "description": "schema padding" * 40}}]
+    agent.log_prefix = ""
+    return agent, compressor
+
+
+def test_compress_context_preserves_last_provider_prompt_tokens():
+    agent, compressor = _make_agent_and_compressor()
+    compressor.last_provider_prompt_tokens = 800
+    compressor.projected_prompt_tokens = 800
+    compressor.projected_prompt_tokens_source = "provider_exact"
+    compressor._transcript_mutated_since_api = False
+
+    messages, new_system = agent._compress_context(
+        [{"role": "user", "content": "hello"}],
+        "base-system",
+        approx_tokens=900,
+    )
+
+    assert new_system == "sys::base-system"
+    assert messages == [{"role": "assistant", "content": "compressed"}]
+    assert compressor.last_provider_prompt_tokens == 800
+    assert compressor.projected_prompt_tokens_source == "estimated_post_compression"
+    assert compressor._transcript_mutated_since_api is True
+
+
+def test_post_compression_pressure_does_not_reuse_stale_provider_exact():
+    agent, compressor = _make_agent_and_compressor()
+    compressor.last_provider_prompt_tokens = 800
+    compressor.projected_prompt_tokens = 800
+    compressor.projected_prompt_tokens_source = "provider_exact"
+    compressor._transcript_mutated_since_api = False
+
+    messages, new_system = agent._compress_context(
+        [{"role": "user", "content": "hello"}],
+        "base-system",
+        approx_tokens=900,
+    )
+
+    tokens, source = compressor.get_current_request_pressure(
+        messages=messages,
+        system_prompt=new_system,
+        tools=agent.tools,
+    )
+
+    assert source == "estimated_post_compression"
+    assert tokens == compressor.projected_prompt_tokens
+    assert tokens != 800