diff --git a/agent/context_compressor.py b/agent/context_compressor.py index 76c4260b69f8..2e0f66260d33 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -1856,6 +1856,16 @@ class ContextCompressor(ContextEngine): self._verify_compaction_cleared_threshold = False self.awaiting_real_usage_after_compression = False + def snapshot_preflight_display_tokens(self) -> int: + """Capture the display token count before a speculative preflight seed.""" + return self.last_prompt_tokens + + def rollback_interrupted_preflight_display_tokens(self, snapshot: int) -> None: + """Restore a speculative display seed without touching compaction state.""" + if self.awaiting_real_usage_after_compression and self.last_prompt_tokens == -1: + return + self.last_prompt_tokens = snapshot + def should_defer_preflight_to_real_usage(self, rough_tokens: int) -> bool: """Return True when a high rough preflight estimate is known-noisy. diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 7bd958612b6d..a3deadd10488 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -2065,6 +2065,8 @@ def run_conversation( ) continue # Retry the API call + agent._turn_received_provider_response = True + # Check finish_reason before proceeding if agent.api_mode == "codex_responses": status = getattr(response, "status", None) diff --git a/agent/turn_context.py b/agent/turn_context.py index 6b2d05858824..f3f5e46b21fe 100644 --- a/agent/turn_context.py +++ b/agent/turn_context.py @@ -706,6 +706,8 @@ def build_turn_context( # issue #27405 (a few very large messages slipping past the count gate). _preflight_compressed = False _preflight_compression_blocked = False + agent._turn_received_provider_response = False + agent._turn_preflight_display_snapshot = None if agent.compression_enabled and _should_run_preflight_estimate( messages, agent.context_compressor.protect_first_n, @@ -718,6 +720,9 @@ def build_turn_context( tools=agent.tools or None, ) _compressor = agent.context_compressor + agent._turn_preflight_display_snapshot = ( + _compressor.snapshot_preflight_display_tokens() + ) _defer_preflight = getattr( _compressor, "should_defer_preflight_to_real_usage", diff --git a/agent/turn_finalizer.py b/agent/turn_finalizer.py index 1a7b52ff516c..17b49479a88b 100644 --- a/agent/turn_finalizer.py +++ b/agent/turn_finalizer.py @@ -201,6 +201,23 @@ def finalize_turn( ) ) + # Preflight can seed the display count before the provider receives the + # request. Roll that estimate back only when an interrupt wins the race + # before any successful provider response. Compaction state remains owned + # by the real-usage/post-compaction path, including its ``-1`` sentinel. + _preflight_snapshot = getattr( + agent, "_turn_preflight_display_snapshot", None + ) + if ( + interrupted + and _preflight_snapshot is not None + and not getattr(agent, "_turn_received_provider_response", False) + and getattr(agent, "context_compressor", None) is not None + ): + agent.context_compressor.rollback_interrupted_preflight_display_tokens( + _preflight_snapshot + ) + # Post-loop cleanup must never lose the response. Trajectory save, # resource teardown, and session persistence all touch fallible # surfaces — file I/O / JSON serialization (_save_trajectory), remote @@ -625,4 +642,7 @@ def finalize_turn( except Exception as exc: logger.warning("on_session_end hook failed: %s", exc) + agent._turn_preflight_display_snapshot = None + agent._turn_received_provider_response = False + return result diff --git a/tests/run_agent/test_413_compression.py b/tests/run_agent/test_413_compression.py index 1d259e23901b..cf05a9a90546 100644 --- a/tests/run_agent/test_413_compression.py +++ b/tests/run_agent/test_413_compression.py @@ -1465,6 +1465,127 @@ class TestPreflightCompression: assert mock_compress.call_count == 2 + def test_interrupt_before_first_provider_call_restores_preflight_display_seed(self, agent): + """Interrupted turns must not keep a speculative preflight display seed. + + Preflight runs before the main loop checks ``_interrupt_requested``. + If the user interrupts during that window, no provider usage ever + arrives to validate the rough estimate, so the old display token count + must be restored instead of leaking the speculative value forward. + """ + agent.compression_enabled = True + agent._interrupt_requested = True + agent.context_compressor.context_length = 200_000 + agent.context_compressor.threshold_tokens = 130_000 + agent.context_compressor.last_prompt_tokens = 74_400 + + big_history = [] + for i in range(20): + big_history.append({"role": "user", "content": f"Message {i} padded text"}) + big_history.append({"role": "assistant", "content": f"Response {i} padded text"}) + + with ( + patch("agent.turn_context.estimate_request_tokens_rough", return_value=144_669), + patch.object(agent.context_compressor, "should_compress", return_value=False), + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + ): + result = agent.run_conversation("hello", conversation_history=big_history) + + assert result["interrupted"] is True + assert agent.client.chat.completions.create.call_count == 0 + assert agent.context_compressor.last_prompt_tokens == 74_400 + + def test_usage_less_provider_response_prevents_display_seed_rollback(self, agent): + """A successful response counts even when the provider omits usage.""" + agent.compression_enabled = True + agent.context_compressor.context_length = 200_000 + agent.context_compressor.threshold_tokens = 130_000 + agent.context_compressor.last_prompt_tokens = 74_400 + + big_history = [] + for i in range(20): + big_history.append({"role": "user", "content": f"Message {i} padded text"}) + big_history.append({"role": "assistant", "content": f"Response {i} padded text"}) + + tool_call = SimpleNamespace( + id="tc1", + type="function", + function=SimpleNamespace(name="web_search", arguments='{"query":"test"}'), + ) + agent.client.chat.completions.create.side_effect = [ + _mock_response( + content=None, + finish_reason="tool_calls", + tool_calls=[tool_call], + usage=None, + ) + ] + + def _interrupt_after_tool(*_args, **_kwargs): + agent._interrupt_requested = True + + with ( + patch("agent.turn_context.estimate_request_tokens_rough", return_value=144_669), + patch.object(agent.context_compressor, "should_compress", return_value=False), + patch.object(agent, "_execute_tool_calls", side_effect=_interrupt_after_tool), + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + ): + result = agent.run_conversation("hello", conversation_history=big_history) + + assert result["interrupted"] is True + assert agent.client.chat.completions.create.call_count == 1 + assert agent.context_compressor.last_prompt_tokens == 144_669 + + def test_interrupt_keeps_post_compression_state(self, agent): + """Display rollback must not restore real post-compaction state. + + A completed preflight compaction still leaves the conversation in the + post-compression ``-1`` sentinel state. Its anti-thrashing verdict also + remains owned by the completed compaction boundary rather than the + speculative display snapshot. + """ + agent.compression_enabled = True + agent._interrupt_requested = True + agent.context_compressor.context_length = 200_000 + agent.context_compressor.threshold_tokens = 130_000 + agent.context_compressor.last_prompt_tokens = 74_400 + agent.context_compressor._ineffective_compression_count = 1 + + big_history = [] + for i in range(20): + big_history.append({"role": "user", "content": f"Message {i} padded text"}) + big_history.append({"role": "assistant", "content": f"Response {i} padded text"}) + + def _fake_preflight_compress(msgs, *_args, **_kwargs): + agent.context_compressor.last_prompt_tokens = -1 + agent.context_compressor.awaiting_real_usage_after_compression = True + agent.context_compressor.compression_count += 1 + agent.context_compressor._ineffective_compression_count = 2 + agent.context_compressor._last_compression_savings_pct = 0.0 + return msgs, agent._cached_system_prompt + + with ( + patch("agent.turn_context.estimate_request_tokens_rough", return_value=144_669), + patch.object(agent.context_compressor, "should_compress", return_value=True), + patch.object(agent, "_compress_context", side_effect=_fake_preflight_compress), + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + ): + result = agent.run_conversation("hello", conversation_history=big_history) + + assert result["interrupted"] is True + assert agent.client.chat.completions.create.call_count == 0 + assert agent.context_compressor.last_prompt_tokens == -1 + assert agent.context_compressor.awaiting_real_usage_after_compression is True + assert agent.context_compressor._ineffective_compression_count == 2 + assert agent.context_compressor._last_compression_savings_pct == 0.0 + + class TestToolResultPreflightCompression: """Compression should trigger when tool results push context past the threshold."""