diff --git a/agent/anthropic_thinking_replay.py b/agent/anthropic_thinking_replay.py index da925699de..b4257ce77a 100644 --- a/agent/anthropic_thinking_replay.py +++ b/agent/anthropic_thinking_replay.py @@ -55,14 +55,29 @@ def rejected_thinking_fingerprints(session_db: Any, session_id: Any) -> set[str] return {value for value in raw if isinstance(value, str) and value} if isinstance(raw, list) else set() +def session_rejected_thinking(holder: Any, session_db: Any, session_id: Any) -> set[str]: + """``holder``'s in-memory fingerprints for ``session_id``, else the persisted ones. The in-memory + set is authoritative: it outlives a failed or disabled persist.""" + cached = getattr(holder, "_anthropic_rejected_thinking", None) + if isinstance(cached, tuple) and cached[0] == session_id: + return cached[1] + return rejected_thinking_fingerprints(session_db, session_id) + + +def _bind(agent: Any, session_id: Any, rejected: set[str]) -> None: + # The compressor's tail walk must price the same replay as the agent's preflight, so it shares + # the agent's set object (later in-place updates reach both). + agent._anthropic_rejected_thinking = (session_id, rejected) + compressor = getattr(agent, "context_compressor", None) + if compressor is not None: + compressor._anthropic_rejected_thinking = agent._anthropic_rejected_thinking + + def _rejected(agent: Any) -> set[str]: # Keyed by session: /new, /resume, /branch and compression rotate ``session_id`` on a live agent. session_id = getattr(agent, "session_id", None) - cached = getattr(agent, "_anthropic_rejected_thinking", None) - if isinstance(cached, tuple) and cached[0] == session_id: - return cached[1] - rejected = rejected_thinking_fingerprints(getattr(agent, "_session_db", None), session_id) - agent._anthropic_rejected_thinking = (session_id, rejected) + rejected = session_rejected_thinking(agent, getattr(agent, "_session_db", None), session_id) + _bind(agent, session_id, rejected) return rejected @@ -177,5 +192,5 @@ def carry_rejected_thinking_to_session(agent: Any, old_session_id: str) -> None: if not rejected: return rejected |= _rejected(agent) - agent._anthropic_rejected_thinking = (getattr(agent, "session_id", None), rejected) + _bind(agent, getattr(agent, "session_id", None), rejected) _persist(agent, rejected) diff --git a/agent/context_compressor.py b/agent/context_compressor.py index ddb2e2fab6..85965470a5 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -4983,10 +4983,10 @@ Write only the summary body. Do not include any preamble or prefix.""" self.base_url, self.model ): return None - from agent.anthropic_thinking_replay import rejected_thinking_fingerprints, strip_rejected_thinking + from agent.anthropic_thinking_replay import session_rejected_thinking, strip_rejected_thinking from agent.message_sanitization import native_anthropic_accounting_projection - rejected = rejected_thinking_fingerprints(getattr(self, "_session_db", None), getattr(self, "_session_id", "")) + rejected = session_rejected_thinking(self, getattr(self, "_session_db", None), getattr(self, "_session_id", "")) def budget(msg: Dict[str, Any]) -> int: request_copy = dict(msg) diff --git a/tests/agent/test_anthropic_preserved_thinking_replay.py b/tests/agent/test_anthropic_preserved_thinking_replay.py index 1055efd5fa..8873bc37a2 100644 --- a/tests/agent/test_anthropic_preserved_thinking_replay.py +++ b/tests/agent/test_anthropic_preserved_thinking_replay.py @@ -690,7 +690,9 @@ def _readable_tokens(blocks): return sum(estimate_tokens_rough(block.get("thinking", "")) for block in blocks) -@pytest.mark.parametrize("shape", ["rejected", "reasoning_only", "invalid_ordered_then_details"]) +@pytest.mark.parametrize( + "shape", ["rejected", "rejected_unpersisted", "reasoning_only", "invalid_ordered_then_details"] +) def test_native_preflight_and_tail_walk_charge_exactly_the_replayed_thinking(tmp_path, shape): """Invariant: on a preserved-thinking native route, growing thinking text moves the preflight estimate and the compressor tail walk by exactly what reaches the wire: nothing for @@ -706,6 +708,8 @@ def test_native_preflight_and_tail_walk_charge_exactly_the_replayed_thinking(tmp api_mode=agent.api_mode, quiet_mode=True, config_context_length=200_000, ) compressor.bind_session_state(db, "s1") + agent.context_compressor = compressor + agent._persist_disabled = shape == "rejected_unpersisted" def history(size): thinking = "x" * size @@ -725,7 +729,7 @@ def test_native_preflight_and_tail_walk_charge_exactly_the_replayed_thinking(tmp {"role": "user", "content": "continue"}, ] - if shape == "rejected": + if shape.startswith("rejected"): _reject_signatures(agent, history(1)) def measure(size):