Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tests/agent/test_replay_budget_accounting.py

@@ -0,0 +1,158 @@
+"""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).
+"""
+
+from agent.context_compressor import (
+    _ALWAYS_REPLAYED_BUDGET_KEYS,
+    _NEWEST_TURN_ONLY_BUDGET_KEYS,
+    _REPLAY_BUDGET_KEYS,
+    _estimate_msg_budget_tokens,
+    _last_assistant_index,
+)
+
+
+BIG_THINKING = "deliberation " * 400  # ~1.3K tokens of stale thinking text
+BIG_BLOB = [{"type": "reasoning", "encrypted_content": "x" * 4000}]
+
+
+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_always_charged(self):
+        """Wire-replayed Codex blobs (incl. native compaction checkpoints)
+        must stay in the budget even for stale turns — #55572's invariant."""
+        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 + 500  # blob still charged on the stale path
+
+    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 TestLastAssistantIndex:
+    def test_finds_newest_assistant(self):
+        msgs = [
+            {"role": "user", "content": "u"},
+            {"role": "assistant", "content": "a1"},
+            {"role": "user", "content": "u2"},
+            {"role": "assistant", "content": "a2"},
+            {"role": "tool", "content": "t"},
+        ]
+        assert _last_assistant_index(msgs) == 3
+
+    def test_no_assistant_returns_minus_one(self):
+        assert _last_assistant_index([{"role": "user", "content": "u"}]) == -1
+        assert _last_assistant_index([]) == -1
+
+
+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