github-featured/chopratejas--headroom
2
0
Fork 0
Code Issues Pull Requests Actions 16 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

fix(ccr): verify a scanned marker's hash before advertising it (#2908)

Browse Source 
## 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.
This commit is contained in:
Ashish Patel
2026-08-13 22:16:21 +05:30
committed by GitHub
parent d76fce04a3
commit 41dab2d099
7 changed files with 329 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
headroom/ccr/tool_injection.py
+85 -2
View File
@@ -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
headroom/proxy/handlers/anthropic.py
+4
View File
@@ -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)
headroom/proxy/handlers/openai.py
+4
View File
@@ -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
tests/test_ccr_tool_injection.py
+175
View File
@@ -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 == []
tests/test_proxy/test_anthropic_ccr_deferred_injection.py
+27
View File
@@ -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
tests/test_proxy_anthropic_cache_stability.py
+19
View File
@@ -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"
tests/test_proxy_handlers_batch.py
+15 -1
View File
@@ -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