"""Tests for provider-aware replay-field accounting in tail-budget walks (#73624). Generic thinking fields (``reasoning`` / ``reasoning_content`` + the ``reasoning_details`` text charge) are replayed for at most the NEWEST assistant turn on every transport — Anthropic strips all-but-newest at convert time, Bedrock never replays thinking, strict chat-completions providers reject or one-space-pad the field. Charging them on every message spent 19-24% of the tail budget on bytes that never reach the wire. Codex sidecar fields (``codex_reasoning_items`` / ``codex_message_items``) ARE wire-replayed on every retained turn and stay charged unconditionally (#55572) — including native compaction checkpoints (#81747). """ import pytest from agent.context_compressor import ( _ALWAYS_REPLAYED_BUDGET_KEYS, _NEWEST_TURN_ONLY_BUDGET_KEYS, _REPLAY_BUDGET_KEYS, _estimate_msg_budget_tokens, ) from agent.model_metadata import estimate_tokens_rough from agent.turn_context import substitute_api_content BIG_THINKING = "deliberation " * 400 # ~1.3K tokens of stale thinking text BIG_BLOB = [{"type": "reasoning", "encrypted_content": "x" * 4000}] @pytest.mark.parametrize( "message", [ {"role": "user", "content": "clean", "api_content": "wire user"}, {"role": "assistant", "content": "clean", "api_content": "wire assistant"}, {"role": "system", "content": "clean", "api_content": "ignored"}, {"role": "tool", "content": "clean", "api_content": "ignored"}, {"role": "user", "content": "clean", "api_content": ""}, {"role": "user", "content": "clean", "api_content": ["ignored"]}, ], ids=[ "user-sidecar", "assistant-sidecar", "system-role", "tool-role", "empty-sidecar", "non-string-sidecar", ], ) def test_api_content_matches_wire_substitution_without_mutation(message): """Tail budgeting must mirror ``substitute_api_content`` exactly.""" original = dict(message) wire_message = dict(message) substitute_api_content(wire_message) tokens = _estimate_msg_budget_tokens(message) expected_content = wire_message.get("content") or "" assert tokens == estimate_tokens_rough(expected_content) + 10 assert message == original def _assistant(thinking=False, codex=False): msg = {"role": "assistant", "content": "done"} if thinking: msg["reasoning"] = BIG_THINKING msg["reasoning_content"] = BIG_THINKING if codex: msg["codex_reasoning_items"] = BIG_BLOB return msg class TestChargeStaleThinking: def test_stale_turn_thinking_not_charged(self): msg = _assistant(thinking=True) full = _estimate_msg_budget_tokens(msg, charge_stale_thinking=True) stale = _estimate_msg_budget_tokens(msg, charge_stale_thinking=False) assert stale < full # The delta is the thinking text — a substantial chunk, not noise. assert full - stale > 300 def test_codex_sidecar_is_not_thinking_text(self): """The Codex sidecar is not stale thinking: the stale path drops nothing from it. Its ciphertext is priced only by real usage (#104462), so the row costs the same as a bare one.""" msg = _assistant(codex=True) full = _estimate_msg_budget_tokens(msg, charge_stale_thinking=True) stale = _estimate_msg_budget_tokens(msg, charge_stale_thinking=False) assert stale == full # nothing thinking-only to drop bare = _estimate_msg_budget_tokens( {"role": "assistant", "content": "done"}, charge_stale_thinking=False ) assert stale < bare + 100 def test_default_is_conservative_full_charge(self): msg = _assistant(thinking=True) assert _estimate_msg_budget_tokens(msg) == _estimate_msg_budget_tokens( msg, charge_stale_thinking=True ) def test_reasoning_details_text_skipped_on_stale_path(self): msg = { "role": "assistant", "content": "done", "reasoning_details": [ {"type": "reasoning.text", "text": "long plan " * 300} ], } full = _estimate_msg_budget_tokens(msg, charge_stale_thinking=True) stale = _estimate_msg_budget_tokens(msg, charge_stale_thinking=False) assert stale < full class TestKeyPartition: def test_partition_covers_replay_budget_keys_exactly(self): """Invariant: the two accounting classes partition _REPLAY_BUDGET_KEYS. A future key added to the replay budget must be classified.""" assert set(_ALWAYS_REPLAYED_BUDGET_KEYS) | set( _NEWEST_TURN_ONLY_BUDGET_KEYS ) == set(_REPLAY_BUDGET_KEYS) assert not set(_ALWAYS_REPLAYED_BUDGET_KEYS) & set( _NEWEST_TURN_ONLY_BUDGET_KEYS ) def test_codex_fields_are_always_replayed_class(self): assert "codex_reasoning_items" in _ALWAYS_REPLAYED_BUDGET_KEYS assert "codex_message_items" in _ALWAYS_REPLAYED_BUDGET_KEYS class TestTailCutBehavior: """The tail cut must protect MORE real transcript when stale turns carry heavy thinking — the #73624 symptom was the cut landing early.""" def _compressor(self): from agent.context_compressor import ContextCompressor cc = ContextCompressor( model="claude-opus-5", quiet_mode=True, config_context_length=200_000, ) return cc def test_stale_thinking_does_not_shrink_tail(self): cc = self._compressor() # Build a transcript where every assistant turn drags huge stale # thinking. Under the old accounting these bloat the walk and the # cut lands early; with newest-turn-only accounting the same budget # protects more messages. msgs = [{"role": "system", "content": "sys"}] for i in range(30): msgs.append({"role": "user", "content": f"question {i}"}) msgs.append( { "role": "assistant", "content": f"answer {i}", "reasoning": BIG_THINKING, "reasoning_content": BIG_THINKING, } ) budget = 3_000 cut = cc._find_tail_cut_by_tokens(msgs, 1, token_budget=budget) # Compute what the OLD accounting (charge everything) would protect. old_accumulated = 0 old_cut = len(msgs) soft = int(budget * 1.5) for i in range(len(msgs) - 1, 0, -1): t = _estimate_msg_budget_tokens(msgs[i], charge_stale_thinking=True) if old_accumulated + t > soft and (len(msgs) - i) >= 3: break old_accumulated += t old_cut = i # New accounting must protect at least as much transcript (lower cut # index = more messages in the tail), and strictly more here because # the stale thinking dominates each message's old-cost. assert cut < old_cut