From 41dab2d09925658b96fed492d534346ce1930f4c Mon Sep 17 00:00:00 2001 From: Ashish Patel Date: Thu, 13 Aug 2026 22:16:21 +0530 Subject: [PATCH] fix(ccr): verify a scanned marker's hash before advertising it (#2908) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `CCRToolInjector.scan_for_markers()` decides whether a compression marker is Headroom's own by *shape* alone — any bracket marker carrying a 24-hex hash counts, per the generic fallback pattern (`\[.*?compressed.*?hash=([a-f0-9]{24})\]`). Other context tools emit exactly that shape. Once a foreign hash is scanned, `has_compressed_content` flips true and the retrieve tool + "Available hashes" system instruction get injected for a hash this proxy never stored — the model calls `headroom_retrieve`, gets a guaranteed miss, and re-does work it already had. Two wasted turns per adopted foreign hash. Closes #2836 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/ccr/tool_injection.py`: added `CCRToolInjector.verify_ownership()` — filters `detected_hashes` down to hashes the compression store actually recognizes, via the same `store.exists()` check the retrieve endpoint itself performs. Added a small `_HashOwnershipStore` Protocol (structural typing, not a hard dependency on the concrete `CompressionStore` class) and a `compression_store` constructor field for dependency injection/testing. `scan_for_markers()` itself is untouched — kept store-independent (pure regex) rather than baking the check into the scan loop, since that approach broke 24 existing tests that correctly test "does this shape match" in isolation. - `headroom/proxy/handlers/anthropic.py`, `headroom/proxy/handlers/openai.py`: call `injector.verify_ownership()` right after `scan_for_markers()` — the two real per-request call sites. - `verify_ownership()` is also called inside `process_request()` (the convenience wrapper `batch.py`'s Google path uses), so that path is covered without a separate call site edit. - `tests/test_ccr_tool_injection.py`: new `TestVerifyOwnership` class (6 tests) — the exact issue repro, a real-hash-survives case, mixed own/foreign hashes, explicit store override, store-exception safety (must not raise), and no-op-on-empty-hashes. - `tests/test_proxy_anthropic_cache_stability.py`: 3 pre-existing tests needed updating for the new (correct) behavior — two `_FakeInjector` test doubles needed a `verify_ownership()` stub added, and one real end-to-end test needed a genuine store entry seeded (via `explicit_hash`) for the hash its hand-typed marker references, instead of asserting on an unverified shape-only match. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) — not run locally, will confirm via CI - [x] New tests added for new functionality - [x] Manual testing performed (see Real Behavior Proof) ### Test Output ```text $ .venv/Scripts/python -m pytest tests/test_ccr_marker_policy.py tests/test_ccr_tool_always_on.py tests/test_ccr_tool_injection.py tests/test_proxy/test_ccr_frozen_prefix_coupling.py tests/test_proxy_anthropic_cache_stability.py tests/test_proxy_handlers_batch.py tests/test_proxy_handler_helpers.py -q 151 passed in 15.87s $ .venv/Scripts/python -m pytest tests/test_proxy_ccr.py tests/test_proxy_openai_responses_stream_ccr.py tests/test_anthropic_ccr_workspace_unbound.py tests/test_compression_store.py tests/test_no_ccr_lossy.py tests/test_proxy_handlers_batch.py -q 121 passed in 20.73s $ .venv/Scripts/ruff check . && .venv/Scripts/ruff format --check . # touched files only All checks passed / already formatted ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.14.5, local venv - Exact command / steps: ran the issue's exact 3-line repro (`CCRToolInjector.scan_for_markers()` on the foreign marker text, then `verify_ownership()`) before and after the fix; separately verified a genuinely-Headroom-stored hash (via `store.store(..., explicit_hash=...)`) still survives verification and still drives injection - Observed result: before the fix (scan only, no verify step exists yet) `has_compressed_content` is `True` for the foreign marker — matches the bug report exactly. After adding `verify_ownership()`: foreign marker → `detected_hashes == []`, `has_compressed_content is False`; real stored hash → `detected_hashes == [real_hash]`, `has_compressed_content is True`. - Not tested: have not driven this through a live two-context-tool proxy session (e.g. Headroom alongside another CCR-shaped tool in the same conversation) — verified at the unit/integration level (the exact repro plus the real proxy handler call sites via `test_proxy_anthropic_cache_stability.py`'s end-to-end `TestClient` tests), not via a live multi-tool session. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation (N/A — internal CCR safety behavior, no user-facing docs reference the marker-adoption mechanism) - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable (release-please generates this automatically from commit messages) ## Additional Notes Design note on why `verify_ownership()` is a separate step rather than baked into `scan_for_markers()`: my first attempt did exactly that and broke 24 tests across `test_ccr_tool_injection.py`, `test_ccr_marker_policy.py`, `test_proxy_anthropic_cache_stability.py`, and `test_proxy_handler_helpers.py` — all of them legitimately testing "does the regex detect this marker shape" independent of any store state. Keeping the scan pure and adding an explicit, separately-testable verification step kept that test surface intact while still closing the real gap at the three places that actually decide whether to advertise the retrieve tool. --- headroom/ccr/tool_injection.py | 87 ++++++++- headroom/proxy/handlers/anthropic.py | 4 + headroom/proxy/handlers/openai.py | 4 + tests/test_ccr_tool_injection.py | 175 ++++++++++++++++++ .../test_anthropic_ccr_deferred_injection.py | 27 +++ tests/test_proxy_anthropic_cache_stability.py | 19 ++ tests/test_proxy_handlers_batch.py | 16 +- 7 files changed, 329 insertions(+), 3 deletions(-) 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