diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 85c72bf377e47..53a867d2f7d3c 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -168,33 +168,29 @@ _HANDOFF_SKIP_FINAL_RESPONSE = ( # to treat it as cancellation metadata rather than assistant prose. INTERRUPT_WAITING_FOR_MODEL_PREFIX = "Operation interrupted: waiting for model response (" -# Refund the per-turn compression budget only when the assembled request has -# dropped this far below threshold_tokens. The gap between this margin and -# 1.0 is the anti-thrash belt: a compaction that lands barely under threshold -# keeps its burnt attempt, so borderline shrink/regrow cycles cannot recycle -# the budget indefinitely. -_COMPRESSION_BUDGET_REFUND_MARGIN = 0.8 - -def _should_refund_compression_budget( +def _should_rearm_compression_budget( compression_attempts: int, - request_pressure_tokens: int, + *, + completed_compaction_pending: bool, + prompt_tokens: int, threshold_tokens: int, ) -> bool: - """True when burnt compression attempts should be returned to the turn. + """Return True after a provider proves a completed compaction worked. - A compaction that brought the assembled request comfortably back under - threshold made real progress, so it must not permanently consume the - per-turn overflow-recovery backstop. No-progress passes never get here - below the margin, so they still exhaust the budget (#11529). + Rough estimates cannot safely rearm the anti-thrash budget: they can dip + below the threshold while the provider-visible prompt remains too large. + Require the completed-compaction latch plus a positive, normalized prompt + count below the threshold from the next successful provider response. """ return bool( compression_attempts + and completed_compaction_pending and threshold_tokens > 0 - and request_pressure_tokens - < threshold_tokens * _COMPRESSION_BUDGET_REFUND_MARGIN + and 0 < prompt_tokens < threshold_tokens ) + # Modules that indicate a deterministic local processing error when they # appear in an exception traceback WITHOUT any API-call module. Used by the # outer-loop error classifier to avoid retrying bugs that will fail @@ -1655,7 +1651,10 @@ def run_conversation( compression_attempts = 0 # One resolved per-turn compression attempt cap, shared by every site that # consumes ``compression_attempts``: the pre-API pressure gate, the - # overflow/413 retry handlers, and the post-tool compaction gate. + # overflow/413 retry handlers, and the post-tool compaction gate. The + # counter is a consecutive unverified/ineffective-attempt backstop: a + # completed compaction rearms it only after a successful provider response + # reports a prompt below the threshold. # Config-driven via compression.max_attempts (parsed + validated in # agent_init); default 3 preserves the prior hardcoded behavior for # objects without the attribute (older pickles / minimal stubs). @@ -2316,31 +2315,6 @@ def run_conversation( f"{_previous_preflight_pressure:,}", f"{request_pressure_tokens:,}", ) - # Refund the shared overflow-recovery budget once the assembled - # request is comfortably back under threshold: a compaction that got - # us here made real progress, so it must not permanently consume the - # per-turn backstop. Without this, a long tool-heavy turn burns all - # attempts on *successful* pre-API compactions, the gate below goes - # permanently dark, and the context grows unchecked until the - # provider rejects the request terminally (root cause of the - # max-compression-attempts dead end on marathon turns). The - # compressor's own should_compress() must agree the pressure is gone — - # when it and the local estimate disagree (#36718 noise, runaway- - # compaction stubs), the budget stays burnt and the hard per-turn cap - # keeps its original meaning. - if _should_refund_compression_budget( - compression_attempts, request_pressure_tokens, _preflight_threshold - ) and not _compressor.should_compress(request_pressure_tokens): - logger.info( - "Compression budget refunded: ~%s request tokens < %s%% of " - "%s threshold (attempts were %s/%s)", - f"{request_pressure_tokens:,}", - int(_COMPRESSION_BUDGET_REFUND_MARGIN * 100), - f"{_preflight_threshold:,}", - compression_attempts, - max_compression_attempts, - ) - compression_attempts = 0 _defer_preflight = getattr( _compressor, "should_defer_preflight_to_real_usage", lambda _t: False ) @@ -3736,7 +3710,37 @@ def run_conversation( "cache_write_tokens": canonical_usage.cache_write_tokens, "reasoning_tokens": canonical_usage.reasoning_tokens, } + # Capture the boundary latch before update_from_response() + # consumes it. Only a real provider prompt count for the + # request immediately following a completed compaction can + # prove that attempt effective and rearm the shared budget. + _completed_compaction_pending = bool( + getattr( + agent.context_compressor, + "_verify_compaction_cleared_threshold", + False, + ) + ) agent.context_compressor.update_from_response(usage_dict) + _compression_threshold = int( + getattr(agent.context_compressor, "threshold_tokens", 0) + or 0 + ) + if _should_rearm_compression_budget( + compression_attempts, + completed_compaction_pending=_completed_compaction_pending, + prompt_tokens=prompt_tokens, + threshold_tokens=_compression_threshold, + ): + logger.info( + "Compression budget rearmed after provider-confirmed " + "recovery: prompt=%s < threshold=%s (attempts were %s/%s)", + f"{prompt_tokens:,}", + f"{_compression_threshold:,}", + compression_attempts, + max_compression_attempts, + ) + compression_attempts = 0 # Stash this response's canonical usage so the post-turn # on_turn_complete() observation hook can forward it (the diff --git a/tests/run_agent/test_compression_budget_refund.py b/tests/run_agent/test_compression_budget_refund.py index 5de0878480312..b3d464d2e4a49 100644 --- a/tests/run_agent/test_compression_budget_refund.py +++ b/tests/run_agent/test_compression_budget_refund.py @@ -1,4 +1,4 @@ -"""Behavioral tests for the per-turn compression budget refund. +"""Behavioral tests for provider-confirmed compression-budget rearming. ``compression_attempts`` is a shared per-turn backstop (pre-API gate, overflow/413 handlers, post-tool gate). Before the refund fix, *successful* @@ -7,15 +7,9 @@ attempts on compactions that worked, the pre-API gate went dark for the rest of the turn, and the context grew unchecked until the provider rejected the request terminally ("max compression attempts (N) reached"). -The refund returns the budget when BOTH hold at the top of a loop pass: - -* the assembled request sits below ``threshold_tokens * - _COMPRESSION_BUDGET_REFUND_MARGIN`` (real progress, not a borderline - shrink), and -* the compressor's own ``should_compress()`` agrees there is no pressure - (divergent-signal guard — a compressor that still demands compression - keeps the hard cap's original meaning, see - test_post_tool_compression_attempt_cap.py). +The budget is rearmed only when a completed compaction is followed by a real +provider prompt count below the configured threshold. Rough estimates and +usage-less responses cannot reopen the anti-thrash cap. These tests drive ``run_conversation()`` through real tool iterations — no source inspection, only observable compaction counts. @@ -29,7 +23,7 @@ from unittest.mock import MagicMock, patch import pytest -from agent.conversation_loop import _should_refund_compression_budget +from agent.conversation_loop import _should_rearm_compression_budget from run_agent import AIAgent @@ -38,23 +32,35 @@ from run_agent import AIAgent # --------------------------------------------------------------------------- -class TestRefundDecision: - def test_no_attempts_no_refund(self): - assert not _should_refund_compression_budget(0, 100, 10_000) +class TestRearmDecision: + def test_provider_confirmed_recovery_rearms(self): + assert _should_rearm_compression_budget( + 2, + completed_compaction_pending=True, + prompt_tokens=7_999, + threshold_tokens=10_000, + ) - def test_zero_threshold_no_refund(self): - assert not _should_refund_compression_budget(2, 100, 0) - - def test_barely_under_threshold_no_refund(self): - # 9,999 of 10,000 is inside the anti-thrash belt (margin 0.8). - assert not _should_refund_compression_budget(2, 9_999, 10_000) - - def test_at_margin_no_refund(self): - assert not _should_refund_compression_budget(2, 8_000, 10_000) - - def test_comfortably_under_margin_refunds(self): - assert _should_refund_compression_budget(2, 7_999, 10_000) - assert _should_refund_compression_budget(1, 100, 10_000) + @pytest.mark.parametrize( + ("attempts", "pending", "prompt_tokens", "threshold_tokens"), + [ + (0, True, 7_999, 10_000), + (2, False, 7_999, 10_000), + (2, True, 0, 10_000), + (2, True, 10_000, 10_000), + (2, True, 10_001, 10_000), + (2, True, 7_999, 0), + ], + ) + def test_unverified_or_pressured_response_keeps_budget_burned( + self, attempts, pending, prompt_tokens, threshold_tokens + ): + assert not _should_rearm_compression_budget( + attempts, + completed_compaction_pending=pending, + prompt_tokens=prompt_tokens, + threshold_tokens=threshold_tokens, + ) # --------------------------------------------------------------------------- @@ -70,7 +76,17 @@ def _tool_call(i: int): ) -def _tool_response(i: int): +def _usage(prompt_tokens: int | None): + if prompt_tokens is None: + return None + return SimpleNamespace( + prompt_tokens=prompt_tokens, + completion_tokens=1, + total_tokens=prompt_tokens + 1, + ) + + +def _tool_response(i: int, prompt_tokens: int | None): msg = SimpleNamespace( content=None, reasoning_content=None, @@ -78,10 +94,12 @@ def _tool_response(i: int): tool_calls=[_tool_call(i)], ) choice = SimpleNamespace(message=msg, finish_reason="tool_calls") - return SimpleNamespace(choices=[choice], model="test/model", usage=None) + return SimpleNamespace( + choices=[choice], model="test/model", usage=_usage(prompt_tokens) + ) -def _stop_response(): +def _stop_response(prompt_tokens: int | None): msg = SimpleNamespace( content="done", reasoning_content=None, @@ -89,7 +107,9 @@ def _stop_response(): tool_calls=None, ) choice = SimpleNamespace(message=msg, finish_reason="stop") - return SimpleNamespace(choices=[choice], model="test/model", usage=None) + return SimpleNamespace( + choices=[choice], model="test/model", usage=_usage(prompt_tokens) + ) def _make_tool_defs(*names: str) -> list: @@ -129,9 +149,18 @@ def _coherent_compressor() -> MagicMock: compressor.threshold_tokens = THRESHOLD compressor.context_length = 200_000 compressor.last_prompt_tokens = 0 + compressor._verify_compaction_cleared_threshold = False + compressor.awaiting_real_usage_after_compression = False compressor.should_compress.side_effect = lambda t=None: (t or 0) >= THRESHOLD compressor.should_defer_preflight_to_real_usage.return_value = False compressor.get_active_compression_failure_cooldown.return_value = None + + def _update_from_response(usage): + compressor.last_prompt_tokens = int(usage.get("prompt_tokens", 0) or 0) + compressor._verify_compaction_cleared_threshold = False + compressor.awaiting_real_usage_after_compression = False + + compressor.update_from_response.side_effect = _update_from_response return compressor @@ -161,10 +190,14 @@ def agent(): return a -def _run_marathon_turn(agent, n_tool_iterations: int): +def _run_marathon_turn( + agent, n_tool_iterations: int, *, provider_prompt_tokens: int | None +): """Drive one turn of ``n_tool_iterations`` oversized tool results.""" - responses = [_tool_response(i) for i in range(n_tool_iterations)] - responses.append(_stop_response()) + responses = [ + _tool_response(i, provider_prompt_tokens) for i in range(n_tool_iterations) + ] + responses.append(_stop_response(provider_prompt_tokens)) agent.client.chat.completions.create.side_effect = responses compress_calls = [] @@ -172,8 +205,11 @@ def _run_marathon_turn(agent, n_tool_iterations: int): def _fake_compress(messages, system_message, **_kwargs): # Model a compaction that works: blank out every oversized payload, # keeping roles and tool-call pairing intact so sanitization is - # unaffected. Pressure drops far below the refund margin. + # unaffected. Arm the same provider-verification boundary as the real + # compression path. compress_calls.append(len(messages)) + agent.context_compressor._verify_compaction_cleared_threshold = True + agent.context_compressor.awaiting_real_usage_after_compression = True compacted = [ dict(m, content="[summarized]") if isinstance(m, dict) and len(str(m.get("content") or "")) > 5_000 @@ -209,7 +245,11 @@ class TestCompressionBudgetRefund: completes. """ assert agent.max_compression_attempts == 3 # config default - result, compress_calls = _run_marathon_turn(agent, n_tool_iterations=8) + result, compress_calls = _run_marathon_turn( + agent, + n_tool_iterations=8, + provider_prompt_tokens=THRESHOLD - 1, + ) assert result["completed"] is True assert len(compress_calls) > 3, ( @@ -217,19 +257,19 @@ class TestCompressionBudgetRefund: f"got only {len(compress_calls)} compactions for 8 pressure spikes" ) - def test_no_refund_when_compressor_still_reports_pressure(self, agent): - """Divergent signals: compressor demands compression regardless of - the local estimate → budget stays burnt at the hard cap. - - Mirrors the always-True stub of the attempt-cap regression tests — - the refund must not reopen that runaway.""" - agent.context_compressor.should_compress.side_effect = None - agent.context_compressor.should_compress.return_value = True - - result, compress_calls = _run_marathon_turn(agent, n_tool_iterations=8) + @pytest.mark.parametrize("provider_prompt_tokens", [None, THRESHOLD]) + def test_unverified_or_pressured_compaction_stays_capped( + self, agent, provider_prompt_tokens + ): + """Missing usage or real usage at threshold cannot recycle the cap.""" + result, compress_calls = _run_marathon_turn( + agent, + n_tool_iterations=8, + provider_prompt_tokens=provider_prompt_tokens, + ) assert result["completed"] is True assert len(compress_calls) <= agent.max_compression_attempts, ( - "with should_compress pinned True the per-turn cap must hold; " + "without provider-confirmed headroom the per-turn cap must hold; " f"got {len(compress_calls)} compactions" )