diff --git a/headroom/proxy/ccr_marker_policy.py b/headroom/proxy/ccr_marker_policy.py index ca1be0c8a..1331359c8 100644 --- a/headroom/proxy/ccr_marker_policy.py +++ b/headroom/proxy/ccr_marker_policy.py @@ -1,4 +1,8 @@ -"""CCR marker freshness and retrieval-tool injection policy.""" +"""CCR marker freshness policy. + +Retrieval-tool injection is decided by ``apply_session_sticky_ccr_tool`` in +``headroom.proxy.helpers``, from what the session has actually forwarded. +""" from __future__ import annotations @@ -28,18 +32,3 @@ def has_new_ccr_markers( ) previous.scan_for_markers(previous_forwarded_messages) return bool(current - set(previous.detected_hashes)) - - -def should_inject_ccr_tool( - *, - configured_inject_tool: bool, - frozen_message_count: int, - has_compressed_content: bool, -) -> tuple[bool, bool]: - """Decide whether the CCR retrieval tool must be injected this turn.""" - - inject_tool = configured_inject_tool - if inject_tool and frozen_message_count > 0: - inject_tool = False - is_marker_override = not inject_tool and has_compressed_content - return (inject_tool or is_marker_override), is_marker_override diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index ee5a9092b..a62fc1d80 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -1845,11 +1845,6 @@ class AnthropicHandlerMixin: ) inject_system_instructions = False configured_inject_tool = self.config.ccr_inject_tool - if configured_inject_tool and frozen_message_count > 0: - logger.info( - f"[{request_id}] CCR: deferring tool injection " - f"(frozen_message_count={frozen_message_count}) to preserve cache" - ) # Scan for compression markers + maybe inject system instructions. # Tool-list injection is handled separately via the sticky helper. injector = CCRToolInjector( @@ -1865,47 +1860,34 @@ class AnthropicHandlerMixin: # retrieval tool once a session has done CCR, regardless # of whether THIS turn produced compressed content. # - # #1006: if tool injection was deferred (frozen prefix) but - # compression just emitted NEW markers this turn, override the - # deferral — the agent has no other way to redeem those markers. - # The cache miss on this one request is preferable to silent - # data loss. If the session has already done CCR the tool is - # already in the client's tool list, so sticky replay is a - # no-op and the cache is unaffected. - # ponytail: ceiling is one extra cache miss on the first CCR - # turn in a frozen-prefix session. - from headroom.proxy.helpers import ( - has_new_ccr_markers, - should_inject_ccr_tool, - ) - - # #1850: only markers NEW this turn justify overriding the - # injection deferral (#1006). Markers replayed from the - # previously-forwarded prefix (overlay_cached_prefix) are - # historical — counting them would re-inject the tool on every - # frozen turn and bust the *tools* cache segment, undoing the - # overlay's messages-prefix cache-safety. - has_new_compressed_content = has_new_ccr_markers( - current_detected_hashes=injector.detected_hashes, - previous_forwarded_messages=prefix_tracker.get_last_forwarded_messages(), - provider="anthropic", - ) - - should_inject, is_marker_override = should_inject_ccr_tool( - configured_inject_tool=configured_inject_tool, - frozen_message_count=frozen_message_count, - has_compressed_content=has_new_compressed_content, - ) - if should_inject: - if is_marker_override: - logger.info( - f"[{request_id}] CCR: overriding injection deferral — " - f"new markers emitted but headroom_retrieve unavailable " - f"(frozen_message_count={frozen_message_count}); injecting to " - "prevent unredeemable markers (#1006)" - ) - from headroom.proxy.helpers import apply_session_sticky_ccr_tool + # Injection is deliberately NOT gated on + # ``frozen_message_count``. That counter answers "is the + # prefix warm?", but the decision needs "does the established + # prefix already contain the tool?" — which only + # ``SessionCcrTracker`` knows. Gating on the counter dropped a + # tool that was already inside the provider-cached prefix, and + # ``tools`` is the head of Anthropic's cache key, so every + # toggle invalidated the entire prefix in both directions. + # ``apply_session_sticky_ccr_tool`` carries the correct rule: + # a session that has never compressed still gets no tool, so + # dropping the gate cannot start injecting into non-CCR + # conversations. + if configured_inject_tool: + from headroom.proxy.helpers import ( + apply_session_sticky_ccr_tool, + has_new_ccr_markers, + ) + # #1850: markers replayed from the previously-forwarded + # prefix (overlay_cached_prefix) are historical; only + # markers NEW this turn may drive a first-time injection, + # else a replayed marker injects the tool into a session + # that never actually compressed. + has_new_compressed_content = has_new_ccr_markers( + current_detected_hashes=injector.detected_hashes, + previous_forwarded_messages=prefix_tracker.get_last_forwarded_messages(), + provider="anthropic", + ) tools, ccr_tool_injected = apply_session_sticky_ccr_tool( provider="anthropic", session_id=session_id, diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 39697dc0e..efa1ccc7b 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -73,9 +73,6 @@ from headroom.proxy.ccr_golden_policy import ( from headroom.proxy.ccr_marker_policy import ( has_new_ccr_markers as _has_new_ccr_markers, ) -from headroom.proxy.ccr_marker_policy import ( - should_inject_ccr_tool as _should_inject_ccr_tool, -) from headroom.proxy.ccr_session_tracker import SessionCcrTracker as _SessionCcrTracker from headroom.proxy.internal_header_policy import ( INTERNAL_HEADER_PREFIX, @@ -1715,35 +1712,6 @@ def has_new_ccr_markers( ) -def should_inject_ccr_tool( - *, - configured_inject_tool: bool, - frozen_message_count: int, - has_compressed_content: bool, -) -> tuple[bool, bool]: - """Decide whether the ``headroom_retrieve`` tool must be injected this turn. - - This is the decision the Anthropic handler used to inline. It is extracted - so the #1006 regression can be pinned at the decision point itself. - - Tool injection is normally deferred when there is a frozen message prefix - (``frozen_message_count > 0``) to preserve the prompt cache. But if - compression emitted fresh markers this turn, deferring would hand the agent - a ``<>`` marker with no tool to redeem it — silent data loss. In - that case we override the deferral and inject anyway (one cache miss is - cheaper than dropped content). - - Returns ``(should_inject, is_marker_override)``. ``is_marker_override`` is - True only when injection happens *because* of new markers despite a deferral, - so the caller can log the override distinctly. - """ - return _should_inject_ccr_tool( - configured_inject_tool=configured_inject_tool, - frozen_message_count=frozen_message_count, - has_compressed_content=has_compressed_content, - ) - - def apply_session_sticky_ccr_tool( *, provider: Literal["anthropic", "openai", "google"], diff --git a/tests/test_ccr_marker_policy.py b/tests/test_ccr_marker_policy.py index 09d89d592..722a59dc3 100644 --- a/tests/test_ccr_marker_policy.py +++ b/tests/test_ccr_marker_policy.py @@ -1,7 +1,7 @@ from __future__ import annotations from headroom.ccr.tool_injection import CCRToolInjector -from headroom.proxy.ccr_marker_policy import has_new_ccr_markers, should_inject_ccr_tool +from headroom.proxy.ccr_marker_policy import has_new_ccr_markers def _hashes(*contents: str) -> list[str]: @@ -70,27 +70,3 @@ def test_has_new_ccr_markers_returns_false_without_current_hashes() -> None: ) is False ) - - -def test_should_inject_ccr_tool_overrides_frozen_prefix_deferral_for_markers() -> None: - assert should_inject_ccr_tool( - configured_inject_tool=True, - frozen_message_count=3, - has_compressed_content=True, - ) == (True, True) - - -def test_should_inject_ccr_tool_defers_frozen_prefix_without_markers() -> None: - assert should_inject_ccr_tool( - configured_inject_tool=True, - frozen_message_count=3, - has_compressed_content=False, - ) == (False, False) - - -def test_should_inject_ccr_tool_injects_configured_tool_without_frozen_prefix() -> None: - assert should_inject_ccr_tool( - configured_inject_tool=True, - frozen_message_count=0, - has_compressed_content=False, - ) == (True, False) diff --git a/tests/test_proxy/test_anthropic_ccr_deferred_injection.py b/tests/test_proxy/test_anthropic_ccr_deferred_injection.py index 71b8f1abf..edccdbee2 100644 --- a/tests/test_proxy/test_anthropic_ccr_deferred_injection.py +++ b/tests/test_proxy/test_anthropic_ccr_deferred_injection.py @@ -9,11 +9,27 @@ pytest.importorskip("fastapi") from fastapi.testclient import TestClient +from headroom.proxy.helpers import _reset_session_ccr_tracker_for_test from headroom.proxy.server import ProxyConfig, create_app _RAW_TRANSCRIPT = "\n".join(f"row {idx}: payload payload payload" for idx in range(80)) +@pytest.fixture(autouse=True) +def _reset_ccr_tracker(): + """Isolate the process-global ``SessionCcrTracker`` between tests. + + Several tests here share ``session_id="stable-session"``, and the tracker's + ``has_done_ccr`` flag is monotonic per session. Without this, a test that + injects the tool leaves the flag set and the next test sees a sticky replay + it never set up — order-dependent, and only visible in file order, not when + run alone. Mirrors the fixture in ``tests/test_ccr_tool_always_on.py``. + """ + _reset_session_ccr_tracker_for_test() + yield + _reset_session_ccr_tracker_for_test() + + class _FakePrefixTracker: def __init__(self, frozen_count: int): self._frozen_count = frozen_count diff --git a/tests/test_proxy/test_ccr_frozen_prefix_coupling.py b/tests/test_proxy/test_ccr_frozen_prefix_coupling.py index b46c5950a..50299fa39 100644 --- a/tests/test_proxy/test_ccr_frozen_prefix_coupling.py +++ b/tests/test_proxy/test_ccr_frozen_prefix_coupling.py @@ -1,13 +1,21 @@ """Regression test for #1006: the proxy must not emit unredeemable CCR markers. -When frozen_message_count > 0, the old code deferred headroom_retrieve tool -injection unconditionally — even if compression just emitted NEW <> -markers the agent has no tool to redeem. +If compression emits a fresh ``<>`` marker, the forwarded request must +also carry ``headroom_retrieve`` — a marker the agent has no tool to redeem is +silent data loss. -The fix: if new markers were emitted this turn, override the deferral and inject -the tool (one cache miss is acceptable; silent data loss is not). That decision -lives in ``should_inject_ccr_tool``, which the Anthropic handler calls; this test -pins the decision at that function so removing the override would fail here. +This used to be enforced by an override *inside* a ``frozen_message_count`` +deferral gate. That gate is gone: deferring on the freeze counter dropped a tool +that was already inside the provider-cached prefix, and ``tools`` is the head of +Anthropic's cache key, so every toggle invalidated the whole prefix. +``apply_session_sticky_ccr_tool`` now decides alone, from what the session has +actually forwarded. #1006 is therefore pinned here, at that helper, and the +turn-over-turn cache property is pinned in +``tests/test_proxy_anthropic_cache_stability.py``. + +These cases deliberately drive a real ``CCRToolInjector`` marker scan rather than +passing a hand-set boolean, so the marker -> flag -> tool chain stays covered end +to end. """ from __future__ import annotations @@ -15,50 +23,13 @@ from __future__ import annotations from unittest.mock import MagicMock, patch from headroom.ccr.tool_injection import CCR_TOOL_NAME, CCRToolInjector -from headroom.proxy.helpers import ( - apply_session_sticky_ccr_tool, - should_inject_ccr_tool, -) +from headroom.proxy.helpers import apply_session_sticky_ccr_tool -class TestShouldInjectCCRTool: - """The decision the handler used to inline. This is where #1006 lived.""" +class TestMarkersImplyRedeemableTool: + """A marker emitted this turn must arrive with the tool that redeems it.""" - def test_overrides_deferral_when_markers_emitted(self): - """Frozen prefix would normally defer, but fresh markers force injection.""" - should_inject, is_override = should_inject_ccr_tool( - configured_inject_tool=True, - frozen_message_count=3, - has_compressed_content=True, - ) - assert should_inject, "must inject to keep markers redeemable (#1006)" - assert is_override, "this is the deferral override path" - - def test_defers_when_no_markers(self): - """Frozen prefix with no new markers stays deferred — no spurious tool.""" - should_inject, is_override = should_inject_ccr_tool( - configured_inject_tool=True, - frozen_message_count=3, - has_compressed_content=False, - ) - assert not should_inject - assert not is_override - - def test_injects_normally_without_frozen_prefix(self): - """No frozen prefix → inject as configured, not via the override path.""" - should_inject, is_override = should_inject_ccr_tool( - configured_inject_tool=True, - frozen_message_count=0, - has_compressed_content=False, - ) - assert should_inject - assert not is_override - - -class TestCCRInjectionEndToEnd: - """The decision feeds apply_session_sticky_ccr_tool; assert the tool lands.""" - - def test_marker_in_frozen_prefix_yields_injected_tool(self): + def test_fresh_marker_yields_injected_tool(self): # Injector detects a fresh marker, i.e. compression ran this turn. injector = CCRToolInjector(provider="anthropic") injector.scan_for_markers( @@ -77,14 +48,6 @@ class TestCCRInjectionEndToEnd: ) assert injector.has_compressed_content, "test setup: injector should detect marker" - # Drive the real decision the handler makes under a frozen prefix. - should_inject, _ = should_inject_ccr_tool( - configured_inject_tool=True, - frozen_message_count=3, - has_compressed_content=injector.has_compressed_content, - ) - assert should_inject - with patch("headroom.proxy.helpers.get_session_ccr_tracker") as mock_tracker_fn: mock_tracker = MagicMock() mock_tracker.has_done_ccr.return_value = False # first CCR ever @@ -101,22 +64,19 @@ class TestCCRInjectionEndToEnd: tool_names = [t.get("name") for t in tools_out] assert CCR_TOOL_NAME in tool_names, ( - f"headroom_retrieve not injected when markers emitted and prefix frozen (#1006). " - f"tools={tool_names}" + f"headroom_retrieve not injected when markers were emitted (#1006). tools={tool_names}" ) - def test_no_marker_in_frozen_prefix_skips_tool(self): + def test_no_marker_on_session_that_never_compressed_skips_tool(self): + """The property that makes dropping the freeze gate safe. + + A session with no markers and no CCR history still gets no tool, so + removing the gate cannot start injecting into non-CCR conversations. + """ injector = CCRToolInjector(provider="anthropic") injector.scan_for_markers([{"role": "user", "content": "hello"}]) assert not injector.has_compressed_content, "test setup: no markers expected" - should_inject, _ = should_inject_ccr_tool( - configured_inject_tool=True, - frozen_message_count=3, - has_compressed_content=injector.has_compressed_content, - ) - assert not should_inject, "no markers → no forced injection" - with patch("headroom.proxy.helpers.get_session_ccr_tracker") as mock_tracker_fn: mock_tracker = MagicMock() mock_tracker.has_done_ccr.return_value = False @@ -128,10 +88,10 @@ class TestCCRInjectionEndToEnd: session_id="session-frozen-no-markers", request_id="req-test-2", existing_tools=[], - has_compressed_content_this_turn=False, + has_compressed_content_this_turn=injector.has_compressed_content, ) tool_names = [t.get("name") for t in tools_out] assert CCR_TOOL_NAME not in tool_names, ( - "headroom_retrieve should NOT be injected when no markers and frozen prefix" + "headroom_retrieve should NOT be injected for a session that never compressed" ) diff --git a/tests/test_proxy_anthropic_cache_stability.py b/tests/test_proxy_anthropic_cache_stability.py index 2e00d3f96..326b9a193 100644 --- a/tests/test_proxy_anthropic_cache_stability.py +++ b/tests/test_proxy_anthropic_cache_stability.py @@ -607,6 +607,122 @@ def test_ccr_tool_injection_disabled_when_prefix_frozen(monkeypatch) -> None: assert captured["inject_tool"] is False +def test_ccr_tool_stays_in_forwarded_tools_across_frozen_transition() -> None: + """``tools`` identity must survive the ``frozen 0 -> >0`` transition. + + ``tools`` is the head of Anthropic's cache key, so adding or removing + ``headroom_retrieve`` between turns invalidates 100% of the provider-cached + prefix — in both directions. Turn 1 (cold prefix, fresh markers) injects the + tool; turn 2 (warm prefix, no *new* markers) must forward the same bytes + rather than dropping it. + + Asserts on the forwarded request body, not on a policy function's return + value: unit-testing the old policy in isolation is exactly what let a + wrong-but-self-consistent decision pass. + """ + from headroom.ccr.tool_injection import CCR_TOOL_NAME + from headroom.proxy.helpers import ( + _reset_session_ccr_tracker_for_test, + serialize_tool_definition_canonical, + ) + + marker_message = { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_bash_x", + "content": ( + "[50 items compressed to 5. Retrieve more: hash=abc123def456abc123def456]" + ), + } + ], + } + forwarded: list[dict] = [] + + _reset_session_ccr_tracker_for_test() + try: + with _make_proxy_client() as client: + proxy = client.app.state.proxy + proxy.config.optimize = False + proxy.config.image_optimize = False + proxy.config.ccr_inject_tool = True + proxy.config.ccr_inject_system_instructions = False + + fake_tracker = _FakePrefixTracker(frozen_count=0) + proxy.session_tracker_store.compute_session_id = lambda request, model, messages: ( + "frozen-transition-session" + ) + proxy.session_tracker_store.get_or_create = lambda session_id, provider: fake_tracker + + async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001 + forwarded.append(body) + return httpx.Response( + 200, + json={ + "id": "msg_frozen_transition", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "ok"}], + "usage": { + "input_tokens": 20, + "output_tokens": 3, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + }, + }, + ) + + proxy._retry_request = _fake_retry + + def _post(): + return client.post( + "/v1/messages", + headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"}, + json={ + "model": "claude-sonnet-4-6", + "max_tokens": 64, + "messages": [marker_message], + }, + ) + + # Turn 1 — cold prefix, marker is new: first-time injection. + assert _post().status_code == 200 + + # Turn 2 — the provider cached turn 1's prefix (tool included), and + # the marker is now historical rather than new. Seed both facts + # explicitly instead of relying on ``update_from_response`` plumbing: + # if the marker still counted as new, the old code would have + # injected via its override path and this test would pass against + # the defect. + fake_tracker._frozen_count = 3 + fake_tracker._last_forwarded_messages = [marker_message] + + assert _post().status_code == 200 + finally: + _reset_session_ccr_tracker_for_test() + + assert len(forwarded) == 2, "expected exactly two forwarded requests" + + def _ccr_tools(body: dict) -> list[dict]: + return [t for t in (body.get("tools") or []) if t.get("name") == CCR_TOOL_NAME] + + turn1 = _ccr_tools(forwarded[0]) + turn2 = _ccr_tools(forwarded[1]) + + assert turn1, "test setup: turn 1 should inject headroom_retrieve on fresh markers" + assert turn2, ( + "headroom_retrieve was dropped from the forwarded tools array once the " + "prefix went warm — that removes a tool already inside the cached prefix " + "and busts 100% of it" + ) + # Byte-identity, not ``==``: a re-serialized definition with a different key + # order compares equal as a dict but busts the cache just as hard. + assert serialize_tool_definition_canonical(turn1[0]) == serialize_tool_definition_canonical( + turn2[0] + ), "headroom_retrieve was re-serialized rather than replayed byte-for-byte" + + def test_previous_turns_always_frozen_only_final_turn_mutable() -> None: captured = {} with _make_proxy_client() as client: