From d6ef4190cc44d301041ca6e4f2d37fed1351e383 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:31:28 -0700 Subject: [PATCH] refactor(turn): collapse preflight estimator branch, build_api_messages pops, docstrings; add turn_context_compaction tests --- agent/turn_context.py | 38 ++++------ tests/agent/test_turn_context_compaction.py | 80 +++++++++++++++++++++ 2 files changed, 92 insertions(+), 26 deletions(-) create mode 100644 tests/agent/test_turn_context_compaction.py diff --git a/agent/turn_context.py b/agent/turn_context.py index 507670380f..edfcb4d317 100644 --- a/agent/turn_context.py +++ b/agent/turn_context.py @@ -1,9 +1,10 @@ """Per-turn setup for ``run_conversation`` (the turn prologue). ``build_turn_context`` runs the once-per-turn setup (stdio guard, sanitization, prompt -restore-or-build, session row, preflight compression, pre_llm_call hook, prefetch, -persistence) mutating ``agent`` exactly as the inline code did, and returns a -``TurnContext`` carrying only the locals the loop reads back.""" +restore-or-build, session row, idle/preflight compaction — see +``turn_context_compaction`` — pre_llm_call hook, prefetch, persistence), mutating +``agent`` as the loop expects, and returns a ``TurnContext`` carrying only the locals +the loop reads back. ``build_api_messages`` builds the wire copy for one API call.""" from __future__ import annotations @@ -64,17 +65,11 @@ def _preflight_request_tokens( "using generic transcript estimate", exc_info=True, ) - if _agent_stale_thinking_on_wire(agent): - return estimate_request_tokens_rough( - messages, - system_prompt=system_prompt or "", - tools=tools, - ) return estimate_request_tokens_rough( messages, system_prompt=system_prompt or "", tools=tools, - charge_stale_thinking=False, + charge_stale_thinking=_agent_stale_thinking_on_wire(agent), ) @@ -148,10 +143,8 @@ def extract_api_content_sidecar(msg: Mapping[str, Any]) -> Optional[str]: def consume_gateway_turn_context_notes(agent: Any) -> str: - """Pop the gateway's per-turn must-deliver notes off the agent (one-shot). - - Staged on ``agent._gateway_turn_context_notes``; consuming them keeps the system - prompt byte-stable and prevents a cached agent replaying a stale note.""" + """Pop the gateway's per-turn must-deliver notes off the agent (one-shot, so the + system prompt stays byte-stable and a cached agent never replays a stale note).""" notes = getattr(agent, "_gateway_turn_context_notes", "") or "" if hasattr(agent, "_gateway_turn_context_notes"): try: @@ -1066,8 +1059,6 @@ def build_api_messages( from agent.agent_runtime_helpers import fill_empty_non_final_wire_payload from agent.conversation_loop import _clone_message_for_send - _ext_prefetch_cache = ext_prefetch_cache - _plugin_user_context = plugin_user_context api_messages = [] for idx, msg in enumerate(messages): @@ -1098,9 +1089,7 @@ def build_api_messages( else: # Callers that bypass the prologue stamping: compose live. _composed = compose_user_api_content( - api_msg.get("content", ""), - _ext_prefetch_cache, - _plugin_user_context, + api_msg.get("content", ""), ext_prefetch_cache, plugin_user_context ) if _composed is not None: api_msg["content"] = _composed @@ -1118,13 +1107,10 @@ def build_api_messages( # This ensures multi-turn reasoning context is preserved agent._copy_reasoning_content_for_api(msg, api_msg) - # Remove 'reasoning' field - it's for trajectory storage only - # We've copied it to 'reasoning_content' for the API above - if "reasoning" in api_msg: - api_msg.pop("reasoning") - # Remove finish_reason - not accepted by strict APIs (e.g. Mistral) - if "finish_reason" in api_msg: - api_msg.pop("finish_reason") + # 'reasoning' is trajectory-only (copied to 'reasoning_content' above); + # finish_reason is rejected by strict APIs (e.g. Mistral). + api_msg.pop("reasoning", None) + api_msg.pop("finish_reason", None) # Fill empty non-final user/assistant wire copies so the pre-call sanitizer # stops re-healing and flooding errors.log; durable history is untouched. # After the reasoning copy so thinking-only turns keep payload (#96870). diff --git a/tests/agent/test_turn_context_compaction.py b/tests/agent/test_turn_context_compaction.py new file mode 100644 index 0000000000..2bc6c9d04c --- /dev/null +++ b/tests/agent/test_turn_context_compaction.py @@ -0,0 +1,80 @@ +"""Unit tests for ``agent.turn_context_compaction`` (turn-start compaction extracted +from ``build_turn_context``).""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from agent.turn_context_compaction import ( + CompactionOutcome, + _codex_native_auto_compaction, + _rearm_uncompressed_overflow_warn, + run_turn_start_compaction, +) + + +def _agent(**kw): + compressor = SimpleNamespace( + protect_first_n=3, protect_last_n=3, threshold_tokens=1_000, context_length=8_000, + summary_target_ratio=0.5, + ) + base = dict( + compression_enabled=False, context_compressor=compressor, session_id="s1", + model="m", _clear_context_overflow_warn=MagicMock(), + ) + base.update(kw) + return SimpleNamespace(**base) + + +def test_codex_native_auto_compaction_gate(): + assert _codex_native_auto_compaction( + SimpleNamespace(api_mode="codex_app_server", codex_app_server_auto_compaction="native") + ) + assert _codex_native_auto_compaction( + SimpleNamespace(api_mode="codex_app_server", codex_app_server_auto_compaction="OFF") + ) + assert not _codex_native_auto_compaction( + SimpleNamespace(api_mode="codex_app_server", codex_app_server_auto_compaction="hermes") + ) + assert not _codex_native_auto_compaction(SimpleNamespace(api_mode="chat_completions")) + + +def test_disabled_compression_rearms_overflow_warn_when_under_window(): + agent = _agent() + msgs = [{"role": "user", "content": "hi"}] + out = run_turn_start_compaction( + agent, messages=msgs, system_message=None, active_system_prompt="sys", + conversation_history=None, current_turn_user_idx=0, user_message="hi", + effective_task_id="t", + ) + assert isinstance(out, CompactionOutcome) + assert out.messages is msgs and out.current_turn_user_idx == 0 + assert out.compressed is False and out.blocked is False + agent._clear_context_overflow_warn.assert_called_once() + assert agent._turn_received_provider_response is False + assert agent._turn_preflight_display_snapshot is None + + +def test_multimodal_content_forces_real_estimate(): + agent = _agent() + msgs = [{"role": "user", "content": [{"type": "text", "text": "x"}]}] + with patch( + "agent.turn_context._preflight_request_tokens", return_value=9_999 + ) as est: + _rearm_uncompressed_overflow_warn(agent, msgs, "sys") + est.assert_called_once() + agent._clear_context_overflow_warn.assert_not_called() + + +def test_preflight_gate_skips_small_transcripts(): + agent = _agent(compression_enabled=True) + agent.context_compressor.should_compress = MagicMock() + msgs = [{"role": "user", "content": "hi"}] + with patch("agent.turn_context._preflight_request_tokens") as est: + out = run_turn_start_compaction( + agent, messages=msgs, system_message=None, active_system_prompt="sys", + conversation_history=None, current_turn_user_idx=0, user_message="hi", + effective_task_id="t", + ) + est.assert_not_called() + agent.context_compressor.should_compress.assert_not_called() + assert out.messages is msgs