diff --git a/headroom/ccr/tool_injection.py b/headroom/ccr/tool_injection.py index 20337285d..86f6cac3f 100644 --- a/headroom/ccr/tool_injection.py +++ b/headroom/ccr/tool_injection.py @@ -16,12 +16,24 @@ from __future__ import annotations import json import re from dataclasses import dataclass, field -from typing import Any +from typing import Any, Protocol, runtime_checkable # Tool name constant - used for matching tool calls CCR_TOOL_NAME = "headroom_retrieve" +@runtime_checkable +class _HashOwnershipStore(Protocol): + """Structural type for verify_ownership()'s store dependency. + + Only needs the existence check — matches CompressionStore.exists() + without coupling this module to the concrete cache implementation + (or requiring test doubles to subclass it). + """ + + def exists(self, hash_key: str, clean_expired: bool = False) -> bool: ... + + def create_ccr_tool_definition( provider: str = "anthropic", ) -> dict[str, Any]: @@ -170,6 +182,11 @@ class CCRToolInjector: inject_tool: bool = True inject_system_instructions: bool = True retrieval_endpoint: str = "/v1/retrieve" + # Store used to verify a scanned marker's hash is actually ours before + # advertising it (issue #2836). None resolves lazily to + # get_compression_store() — request-scoped store if one is set, else the + # global singleton — matching how every other CCR call site resolves it. + compression_store: _HashOwnershipStore | None = None # Detected compression markers _detected_hashes: list[str] = field(default_factory=list) @@ -281,7 +298,16 @@ class CCRToolInjector: return self._detected_hashes def _scan_text(self, text: str) -> None: - """Scan text for compression markers from any compressor.""" + """Scan text for compression markers from any compressor. + + Shape-only: this matches the bracket format any compressor (or, + as it turns out, any *other* context tool) can produce. Callers + that need to know the hash is actually ours — i.e. before + advertising it to the model via the retrieve tool — must call + :meth:`verify_ownership` afterward. Kept separate so this method + stays a pure, store-independent text scan (that's what the + existing marker-format test suite exercises). + """ for pattern in self._marker_patterns: matches = pattern.findall(text) for match in matches: @@ -293,6 +319,59 @@ class CCRToolInjector: if hash_key and hash_key not in self._detected_hashes: self._detected_hashes.append(hash_key) + def verify_ownership(self, store: _HashOwnershipStore | None = None) -> list[str]: + """Drop any detected hash the compression store doesn't recognize. + + The bracket-marker shape (``[... hash=...]``) is not unique to + Headroom — other context tools emit visually identical markers. + Matching shape alone (what :meth:`scan_for_markers` does) adopts + their hashes too: ``has_compressed_content`` goes true and the + retrieve tool + "Available hashes" instruction get injected for a + hash this proxy never stored. The model then calls + ``headroom_retrieve``, gets a guaranteed miss, and re-does the work + it already had (issue #2836). + + Call this after :meth:`scan_for_markers` and before checking + ``has_compressed_content`` / injecting the tool. Uses the same + ``store.exists()`` check the retrieve endpoint itself performs, so + a hash that survives this filter is provably redeemable right now + (or, if it expires between this check and the model's next call, + fails the same way a genuinely-ours stale hash already would — + this only removes hashes that were never ours to begin with). + + Args: + store: Compression store to verify against. Defaults to + ``get_compression_store()`` (request-scoped if set, else + the global singleton) — the same resolution every other + CCR call site uses. + + Returns: + The filtered ``detected_hashes`` list (also updates + ``self.detected_hashes`` in place). + """ + if not self._detected_hashes: + return self._detected_hashes + if store is None: + store = self.compression_store + if store is None: + from headroom.cache.compression_store import get_compression_store + + store = get_compression_store() + + def _safe_exists(hash_key: str) -> bool: + try: + return store.exists(hash_key) + except Exception: + # A store lookup failure must not make CCR verification + # blow up the request; treat as "not ours" (drop the + # marker) — the safe direction, since a dropped real + # marker just means the model can't use the retrieve tool + # for it this turn, the same failure mode as CCR being off. + return False + + self._detected_hashes = [h for h in self._detected_hashes if _safe_exists(h)] + return self._detected_hashes + def inject_tool_definition( self, tools: list[dict[str, Any]] | None, @@ -437,6 +516,10 @@ class CCRToolInjector: tool_was_injected is False if tool was already present (e.g., from MCP). """ self.scan_for_markers(messages) + # Shape-only scanning also matches markers from other context tools; + # drop hashes this proxy never actually stored before they can + # drive tool injection (issue #2836). + self.verify_ownership() if not (self.has_compressed_content or session_has_done_ccr): return messages, tools, False diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 52db0a846..5426d0cf5 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2006,6 +2006,10 @@ class AnthropicHandlerMixin: inject_system_instructions=inject_system_instructions, ) injector.scan_for_markers(optimized_messages) + # Shape-only scanning also matches markers from other context + # tools; drop hashes this proxy never actually stored before + # they can drive tool injection (issue #2836). + injector.verify_ownership() if inject_system_instructions and injector.has_compressed_content: optimized_messages = injector.inject_into_system_message(optimized_messages) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 8e38dce35..636d3a69c 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -3603,6 +3603,10 @@ class OpenAIHandlerMixin: ), ) injector.scan_for_markers(optimized_messages) + # Shape-only scanning also matches markers from other context + # tools; drop hashes this proxy never actually stored before + # they can drive tool injection (issue #2836). + injector.verify_ownership() if ( self.config.ccr_inject_system_instructions and not stream diff --git a/tests/test_ccr_tool_injection.py b/tests/test_ccr_tool_injection.py index 6bf2ccc13..e5e5b5a76 100644 --- a/tests/test_ccr_tool_injection.py +++ b/tests/test_ccr_tool_injection.py @@ -11,6 +11,16 @@ from headroom.ccr import ( ) +class _AlwaysOwnStore: + """Stub compression store for verify_ownership() (issue #2836) in tests + that only exercise injection plumbing with hand-typed marker hashes, + not real CompressionStore-backed storage. + """ + + def exists(self, hash_key: str, clean_expired: bool = False) -> bool: + return True + + class TestCCRToolDefinition: """Test tool definition creation for different providers.""" @@ -271,6 +281,10 @@ class TestCCRToolInjector: provider="anthropic", inject_tool=True, inject_system_instructions=True, + # verify_ownership() (issue #2836) requires the store to + # recognize the hash; this test only exercises injection + # plumbing, not real storage, so stub ownership as always-true. + compression_store=_AlwaysOwnStore(), ) updated_messages, updated_tools, was_injected = injector.process_request(messages, None) @@ -620,3 +634,164 @@ class TestAlternativeMarkerFormats: assert len(hashes) == 1 assert "fedcba9876543210fedcba98" in hashes + + +class TestVerifyOwnership: + """Regression tests for issue #2836. + + Shape-only marker scanning (``scan_for_markers``) matches markers from + ANY context tool that happens to use the same bracket format, not just + Headroom's own. ``verify_ownership`` closes that gap by checking each + detected hash against the actual compression store before it can drive + retrieve-tool injection. + """ + + def test_foreign_marker_is_dropped(self): + """The exact repro from issue #2836: a marker Headroom never + created must not be adopted, even though its shape matches. + """ + from headroom.cache.compression_store import reset_compression_store + + reset_compression_store() + try: + foreign = ( + "[374 items compressed to 267 (from 65 source lines). " + "Retrieve more: hash=ddc3d69afad7bc53fbee11e2]" + ) + injector = CCRToolInjector(provider="anthropic") + injector.scan_for_markers([{"role": "user", "content": foreign}]) + + # Shape-only scan still finds it — that's the bug surface. + assert injector.detected_hashes == ["ddc3d69afad7bc53fbee11e2"] + + injector.verify_ownership() + + assert injector.detected_hashes == [] + assert injector.has_compressed_content is False + finally: + reset_compression_store() + + def test_real_hash_survives_verification(self): + """A hash Headroom actually stored must still be recognized.""" + from headroom.cache.compression_store import ( + get_compression_store, + reset_compression_store, + ) + + reset_compression_store() + try: + store = get_compression_store() + real_hash = store.store( + original="original content", + compressed="compressed content", + explicit_hash="abc123def456abc123def456", + ) + + marker = f"[100 items compressed to 10. Retrieve more: hash={real_hash}]" + injector = CCRToolInjector(provider="anthropic") + injector.scan_for_markers([{"role": "user", "content": marker}]) + injector.verify_ownership() + + assert injector.detected_hashes == [real_hash] + assert injector.has_compressed_content is True + finally: + reset_compression_store() + + def test_mixed_own_and_foreign_hashes_keeps_only_own(self): + """One own hash and one foreign hash in the same scan — only the + own hash survives verification. + """ + from headroom.cache.compression_store import ( + get_compression_store, + reset_compression_store, + ) + + reset_compression_store() + try: + store = get_compression_store() + store.store( + original="mine", + compressed="mine-compressed", + explicit_hash="111111111111111111111111", + ) + messages = [ + { + "role": "user", + "content": ( + "[10 items compressed to 5. Retrieve more: hash=111111111111111111111111]" + "\n[20 items compressed to 8. Retrieve more: hash=222222222222222222222222]" + ), + } + ] + injector = CCRToolInjector(provider="anthropic") + injector.scan_for_markers(messages) + assert set(injector.detected_hashes) == { + "111111111111111111111111", + "222222222222222222222222", + } + + injector.verify_ownership() + + assert injector.detected_hashes == ["111111111111111111111111"] + finally: + reset_compression_store() + + def test_explicit_store_takes_precedence_over_global(self): + """A store passed to verify_ownership() overrides the default + (global/request-scoped) resolution — matches the constructor's + compression_store field too. + """ + + class _NeverOwnStore: + def exists(self, hash_key, clean_expired=False): # noqa: ANN001 + return False + + injector = CCRToolInjector(provider="anthropic", compression_store=_NeverOwnStore()) + injector.scan_for_markers( + [ + { + "role": "user", + "content": "[1 items compressed to 1. Retrieve more: hash=abcabcabcabcabcabcabcabc]", + } + ] + ) + injector.verify_ownership() + + assert injector.detected_hashes == [] + + def test_store_lookup_exception_is_treated_as_not_owned(self): + """A store lookup failure must not crash CCR verification — it + should drop the marker (the safe direction), not raise. + """ + + class _BrokenStore: + def exists(self, hash_key, clean_expired=False): # noqa: ANN001 + raise RuntimeError("store backend unavailable") + + injector = CCRToolInjector(provider="anthropic", compression_store=_BrokenStore()) + injector.scan_for_markers( + [ + { + "role": "user", + "content": "[1 items compressed to 1. Retrieve more: hash=abcabcabcabcabcabcabcabc]", + } + ] + ) + injector.verify_ownership() # must not raise + + assert injector.detected_hashes == [] + + def test_verify_ownership_is_noop_on_empty_hashes(self): + """No detected hashes -> verify_ownership must not touch the store + at all (nothing to verify). + """ + + class _ExplodingStore: + def exists(self, hash_key, clean_expired=False): # noqa: ANN001 + raise AssertionError("should not be called with no detected hashes") + + injector = CCRToolInjector(provider="anthropic", compression_store=_ExplodingStore()) + injector.scan_for_markers([{"role": "user", "content": "no markers here"}]) + result = injector.verify_ownership() + + assert result == [] diff --git a/tests/test_proxy/test_anthropic_ccr_deferred_injection.py b/tests/test_proxy/test_anthropic_ccr_deferred_injection.py index 46a334113..a8821a9ba 100644 --- a/tests/test_proxy/test_anthropic_ccr_deferred_injection.py +++ b/tests/test_proxy/test_anthropic_ccr_deferred_injection.py @@ -9,11 +9,16 @@ pytest.importorskip("fastapi") from fastapi.testclient import TestClient +from headroom.cache.compression_store import get_compression_store, reset_compression_store 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)) +# The hash most fixtures below embed in a "[... Retrieve more: hash=...]" +# marker to drive CCR tool injection. +_MARKER_HASH = "abc123def456abc123def456" + @pytest.fixture(autouse=True) def _reset_ccr_tracker(): @@ -30,6 +35,28 @@ def _reset_ccr_tracker(): _reset_session_ccr_tracker_for_test() +@pytest.fixture(autouse=True) +def _seed_marker_hash_in_store(): + """Make ``_MARKER_HASH`` a real, verifiable compression-store entry. + + CCRToolInjector.verify_ownership() (issue #2836) only advertises the + retrieve tool for hashes the compression store actually recognizes. + These fixtures hand-type marker text rather than compressing real + content through the store, so without this the hash would (correctly) + be treated as foreign and the tool would never get injected — these + tests are about the deferred-injection *policy*, not about exercising + real storage, so seed the one hash they all key off of. + """ + reset_compression_store() + get_compression_store().store( + original="original tool output", + compressed="[100 items compressed to 10]", + explicit_hash=_MARKER_HASH, + ) + yield + reset_compression_store() + + class _FakePrefixTracker: def __init__(self, frozen_count: int): self._frozen_count = frozen_count diff --git a/tests/test_proxy_anthropic_cache_stability.py b/tests/test_proxy_anthropic_cache_stability.py index d04972563..b7a26d68f 100644 --- a/tests/test_proxy_anthropic_cache_stability.py +++ b/tests/test_proxy_anthropic_cache_stability.py @@ -515,6 +515,9 @@ def test_ccr_system_instruction_injection_disabled_when_prefix_frozen(monkeypatc def scan_for_markers(self, messages): # noqa: ANN001 return [] + def verify_ownership(self, store=None): # noqa: ANN001 + return self.detected_hashes + monkeypatch.setattr("headroom.ccr.CCRToolInjector", _FakeInjector) async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001 @@ -582,6 +585,9 @@ def test_ccr_tool_injection_disabled_when_prefix_frozen(monkeypatch) -> None: def scan_for_markers(self, messages): # noqa: ANN001 return [] + def verify_ownership(self, store=None): # noqa: ANN001 + return self.detected_hashes + monkeypatch.setattr("headroom.ccr.CCRToolInjector", _FakeInjector) async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001 @@ -630,12 +636,24 @@ def test_ccr_tool_stays_in_forwarded_tools_across_frozen_transition() -> None: value: unit-testing the old policy in isolation is exactly what let a wrong-but-self-consistent decision pass. """ + from headroom.cache.compression_store import get_compression_store, reset_compression_store from headroom.ccr.tool_injection import CCR_TOOL_NAME from headroom.proxy.helpers import ( _reset_session_ccr_tracker_for_test, serialize_tool_definition_canonical, ) + # verify_ownership() (issue #2836) requires the marker's hash to be a + # real store entry — seed one with the exact hash the marker text below + # references, via explicit_hash (the store's own hash generation from + # `original` content wouldn't match this hand-typed literal). + reset_compression_store() + get_compression_store().store( + original="original tool output", + compressed="[50 items compressed to 5]", + explicit_hash="abc123def456abc123def456", + ) + marker_message = { "role": "user", "content": [ @@ -711,6 +729,7 @@ def test_ccr_tool_stays_in_forwarded_tools_across_frozen_transition() -> None: assert _post().status_code == 200 finally: _reset_session_ccr_tracker_for_test() + reset_compression_store() assert len(forwarded) == 2, "expected exactly two forwarded requests" diff --git a/tests/test_proxy_handlers_batch.py b/tests/test_proxy_handlers_batch.py index a20a18d1b..53e823fcf 100644 --- a/tests/test_proxy_handlers_batch.py +++ b/tests/test_proxy_handlers_batch.py @@ -6,7 +6,11 @@ from types import SimpleNamespace import pytest -from headroom.cache.compression_store import CompressionEntry +from headroom.cache.compression_store import ( + CompressionEntry, + get_compression_store, + reset_compression_store, +) from headroom.ccr import response_handler as response_handler_module from headroom.proxy.handlers import batch as batch_module from headroom.proxy.handlers import gemini as gemini_module @@ -295,6 +299,15 @@ async def test_gemini_native_ccr_continuation(monkeypatch: pytest.MonkeyPatch) - @pytest.mark.asyncio async def test_gemini_native_ccr_tools(monkeypatch: pytest.MonkeyPatch) -> None: install_native_gemini_compression(monkeypatch) + # verify_ownership() (issue #2836) requires the marker's hash to be a + # real store entry; NativeGeminiHandler's mocked pipeline hand-types + # "hash=aaaa...aaaa" rather than compressing through the real store. + reset_compression_store() + get_compression_store().store( + original="original content", + compressed="compressed [100 items compressed to 1]", + explicit_hash="aaaaaaaaaaaaaaaaaaaaaaaa", + ) handler = NativeGeminiHandler( [FakeResponse(json_data={"candidates": [{"content": {"parts": [{"text": "answer"}]}}]})] ) @@ -319,6 +332,7 @@ async def test_gemini_native_ccr_tools(monkeypatch: pytest.MonkeyPatch) -> None: declarations = forwarded_tools[0]["functionDeclarations"] assert {item["name"] for item in declarations} == {"client_tool", "headroom_retrieve"} assert forwarded_tools[1]["functionDeclarations"] == [{"name": "second_tool"}] + reset_compression_store() @pytest.mark.asyncio