mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): freeze must forward cached (compressed) prefix byte-identical — stop token-mode cache busting (#1850)
The freeze path (both providers) emits the agent's ORIGINAL bytes for a frozen message, but the provider cached whatever we FORWARDED last turn (the compressed form). Forwarding original then mismatches the cached prefix and busts it from that point — re-creating the whole suffix. Measured on a real SWE-bench run: 100% of attributed misses were prefix_change, ~56% of ALL cache-writes were bust-induced (2.8M tokens), driving cache_create +150% and cost +41% vs baseline. Cache mode already avoided this via _extract_cache_stable_delta (replay the previously-forwarded prefix, compress only the delta). Token mode called apply(frozen_count) directly, which forwards original for the frozen region. Fix: add a shared, provider-agnostic overlay_cached_prefix() that replays the previously-forwarded (cached, compressed) prefix byte-identical, append-only guarded and idempotent, and apply it in BOTH the Anthropic and OpenAI handlers right before forwarding. This makes freezing byte-identical in every mode, so the only remaining difference between "token" and "cache" mode is how large a mutable (still-compressible) tail each leaves — not whether the frozen prefix busts the cache. Tests: - test_cache_prefix_overlay.py: the helper (replay, append-only guard, idempotence). - test_cross_turn_cache_safety.py: the invariant that was missing — drive the REAL tracker + freeze + overlay over multiple append-only turns against a simulated provider prefix cache and assert the forwarded prefix stays byte-identical turn-over-turn. Load-bearing: it fails (detects the bust) without the overlay. ## Description <!-- Briefly explain the change and why it is needed. --> Closes # ## Type of Change - [ ] 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 - ## Testing <!-- Check what you actually ran, then paste the real command output below. --> - [ ] Unit tests pass (`pytest`) - [ ] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) - [ ] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text # Paste relevant command output or artifact links here ``` ## Real Behavior Proof - Environment: - Exact command / steps: - Observed result: - Not tested: ## Review Readiness - [ ] I have performed a self-review - [ ] This PR is ready for human review ## Checklist - [ ] My code follows the project's style guidelines - [ ] I have performed a self-review of my code - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) Add screenshots to help explain your changes. ## Additional Notes <!-- Mention any N/A checklist items, tradeoffs, follow-ups, or maintainer context. -->
This commit is contained in:
parent
7208792ee8
commit
248ae0f3e0
9 changed files with 471 additions and 9 deletions
46
headroom/cache/prefix_tracker.py
vendored
46
headroom/cache/prefix_tracker.py
vendored
|
|
@ -112,6 +112,52 @@ class CacheMissAttribution:
|
|||
ttl_exceeded: bool = False
|
||||
|
||||
|
||||
def overlay_cached_prefix(
|
||||
optimized_messages: list[dict[str, Any]],
|
||||
current_original_messages: list[dict[str, Any]],
|
||||
previous_original_messages: list[dict[str, Any]] | None,
|
||||
previous_forwarded_messages: list[dict[str, Any]] | None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Replay the previously-forwarded (cached, compressed) prefix byte-identical.
|
||||
|
||||
Provider-agnostic cache-safety guard for the freeze path. When a message is
|
||||
"frozen", the compression pipeline may emit the agent's ORIGINAL bytes for
|
||||
it — but the provider cached whatever we FORWARDED last turn (the compressed
|
||||
form). Forwarding the original then mismatches the cached prefix and busts
|
||||
the prompt cache from that point (100% of observed misses were this
|
||||
``prefix_change``). This overlays the exact previously-forwarded prefix onto
|
||||
the corresponding leading messages so the forwarded prefix stays byte-for-byte
|
||||
what the provider hashed for its cache key.
|
||||
|
||||
Safe only when this turn append-only-extends the previous turn (the standard
|
||||
growing-conversation shape): the previous ORIGINAL messages must be an exact
|
||||
prefix of the current ORIGINAL messages, and there is exactly one forwarded
|
||||
message per original. Otherwise the previous forwarded bytes may not
|
||||
correspond to the same positions, so we return ``optimized_messages``
|
||||
unchanged (accept a possible bust rather than forward wrong content).
|
||||
|
||||
This makes freezing byte-identical in BOTH proxy modes, so the only remaining
|
||||
difference between them is how large a mutable (still-compressible) tail each
|
||||
leaves — not whether the frozen prefix busts the cache.
|
||||
"""
|
||||
prev_orig = previous_original_messages
|
||||
prev_fwd = previous_forwarded_messages
|
||||
if not prev_orig or not prev_fwd:
|
||||
return optimized_messages
|
||||
n = len(prev_orig)
|
||||
# One forwarded message per original, and the frozen prefix must fit within
|
||||
# both the current originals and this turn's optimized output.
|
||||
if len(prev_fwd) != n:
|
||||
return optimized_messages
|
||||
if len(current_original_messages) < n or len(optimized_messages) < n:
|
||||
return optimized_messages
|
||||
# Append-only guard: the frozen region must be the same messages we cached.
|
||||
if current_original_messages[:n] != prev_orig:
|
||||
return optimized_messages
|
||||
# Replay the cached (compressed) prefix; keep this turn's compressed tail.
|
||||
return list(prev_fwd) + list(optimized_messages[n:])
|
||||
|
||||
|
||||
class PrefixCacheTracker:
|
||||
"""Tracks provider prefix cache state across turns in a session.
|
||||
|
||||
|
|
|
|||
|
|
@ -1339,6 +1339,27 @@ class AnthropicHandlerMixin:
|
|||
# Flag compression failure for observability
|
||||
_compression_failed = True
|
||||
|
||||
# Cache-safety (ALL modes): forward the previously-cached (compressed)
|
||||
# prefix byte-identical. The freeze path can emit the agent's ORIGINAL
|
||||
# bytes for a frozen message, but the provider cached whatever we
|
||||
# FORWARDED last turn (the compressed form); forwarding original then
|
||||
# mismatches the cached prefix and busts it (prefix_change was 100% of
|
||||
# observed misses, ~56% of all cache-writes). Replaying the exact
|
||||
# previously-forwarded prefix keeps it byte-identical → cache hits.
|
||||
# Append-only-guarded and idempotent (cache mode already replays), so
|
||||
# it is safe to run unconditionally here.
|
||||
from headroom.cache.prefix_tracker import overlay_cached_prefix
|
||||
|
||||
_ov = overlay_cached_prefix(
|
||||
optimized_messages,
|
||||
original_client_messages,
|
||||
prefix_tracker.get_last_original_messages(),
|
||||
prefix_tracker.get_last_forwarded_messages(),
|
||||
)
|
||||
if _ov != optimized_messages:
|
||||
optimized_messages = _ov
|
||||
optimized_tokens = tokenizer.count_messages(optimized_messages)
|
||||
|
||||
# Guard: if "optimization" inflated tokens, revert to originals.
|
||||
# Skip in cache mode where prefix-stability may legitimately shift counts.
|
||||
if optimized_tokens > original_tokens and not is_cache_mode(self.config.mode):
|
||||
|
|
@ -1553,12 +1574,27 @@ class AnthropicHandlerMixin:
|
|||
# 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 should_inject_ccr_tool
|
||||
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=injector.has_compressed_content,
|
||||
has_compressed_content=has_new_compressed_content,
|
||||
)
|
||||
if should_inject:
|
||||
if is_marker_override:
|
||||
|
|
@ -1575,7 +1611,7 @@ class AnthropicHandlerMixin:
|
|||
session_id=session_id,
|
||||
request_id=request_id,
|
||||
existing_tools=tools,
|
||||
has_compressed_content_this_turn=injector.has_compressed_content,
|
||||
has_compressed_content_this_turn=has_new_compressed_content,
|
||||
)
|
||||
if ccr_tool_injected:
|
||||
logger.debug(
|
||||
|
|
|
|||
|
|
@ -2256,6 +2256,22 @@ class OpenAIHandlerMixin:
|
|||
# Flag compression failure for observability
|
||||
_compression_failed = True
|
||||
|
||||
# Cache-safety (ALL modes): forward the previously-cached (compressed)
|
||||
# prefix byte-identical, so freezing can't bust the prompt cache. See the
|
||||
# matching guard in the Anthropic handler for the full rationale. Append-
|
||||
# only-guarded and idempotent (cache mode already replays).
|
||||
from headroom.cache.prefix_tracker import overlay_cached_prefix
|
||||
|
||||
_ov = overlay_cached_prefix(
|
||||
optimized_messages,
|
||||
original_client_messages,
|
||||
openai_prefix_tracker.get_last_original_messages(),
|
||||
openai_prefix_tracker.get_last_forwarded_messages(),
|
||||
)
|
||||
if _ov != optimized_messages:
|
||||
optimized_messages = _ov
|
||||
optimized_tokens = tokenizer.count_messages(optimized_messages)
|
||||
|
||||
# Guard: if "optimization" inflated tokens, revert to originals
|
||||
if optimized_tokens > original_tokens:
|
||||
logger.warning(
|
||||
|
|
@ -2352,14 +2368,27 @@ class OpenAIHandlerMixin:
|
|||
optimized_messages = injector.inject_into_system_message(optimized_messages)
|
||||
|
||||
if self.config.ccr_inject_tool:
|
||||
from headroom.proxy.helpers import apply_session_sticky_ccr_tool
|
||||
from headroom.proxy.helpers import (
|
||||
apply_session_sticky_ccr_tool,
|
||||
has_new_ccr_markers,
|
||||
)
|
||||
|
||||
# #1850: markers replayed from overlay_cached_prefix are
|
||||
# historical; only markers NEW this turn should drive injection,
|
||||
# else we re-inject the tool 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=openai_prefix_tracker.get_last_forwarded_messages(),
|
||||
provider="openai",
|
||||
)
|
||||
tools, ccr_tool_injected = apply_session_sticky_ccr_tool(
|
||||
provider="openai",
|
||||
session_id=openai_session_id,
|
||||
request_id=request_id,
|
||||
existing_tools=tools,
|
||||
has_compressed_content_this_turn=injector.has_compressed_content,
|
||||
has_compressed_content_this_turn=has_new_compressed_content,
|
||||
)
|
||||
if ccr_tool_injected:
|
||||
logger.debug(
|
||||
|
|
|
|||
|
|
@ -2566,6 +2566,43 @@ def _reset_session_ccr_tracker_for_test() -> None:
|
|||
_session_ccr_tracker = None
|
||||
|
||||
|
||||
def has_new_ccr_markers(
|
||||
*,
|
||||
current_detected_hashes: list[str],
|
||||
previous_forwarded_messages: list[dict[str, Any]] | None,
|
||||
provider: Literal["anthropic", "openai", "google"],
|
||||
) -> bool:
|
||||
"""Whether the about-to-forward content carries CCR markers NOT already forwarded.
|
||||
|
||||
``overlay_cached_prefix`` (#1850) replays the previously-forwarded (compressed)
|
||||
prefix byte-identical to keep the prompt cache warm — which reintroduces the
|
||||
``hash=…`` markers that prefix already carried. Those markers are *historical*:
|
||||
the agent saw them last turn and the retrieve-tool state was already settled
|
||||
for them. Only markers that are genuinely NEW this turn justify overriding the
|
||||
tool-injection deferral (#1006); counting the replayed ones would re-inject the
|
||||
tool on every frozen turn and bust the *tools* cache segment (undoing the very
|
||||
cache-safety the overlay provides).
|
||||
|
||||
Returns True iff ``current_detected_hashes`` contains a hash that is not present
|
||||
in ``previous_forwarded_messages``.
|
||||
"""
|
||||
current = set(current_detected_hashes)
|
||||
if not current:
|
||||
return False
|
||||
if not previous_forwarded_messages:
|
||||
# No prior forward → every marker is new (genuine first CCR turn).
|
||||
return True
|
||||
from headroom.ccr.tool_injection import CCRToolInjector
|
||||
|
||||
prev = CCRToolInjector(
|
||||
provider=provider,
|
||||
inject_tool=False,
|
||||
inject_system_instructions=False,
|
||||
)
|
||||
prev.scan_for_markers(previous_forwarded_messages)
|
||||
return bool(current - set(prev.detected_hashes))
|
||||
|
||||
|
||||
def should_inject_ccr_tool(
|
||||
*,
|
||||
configured_inject_tool: bool,
|
||||
|
|
|
|||
79
tests/test_cache_prefix_overlay.py
Normal file
79
tests/test_cache_prefix_overlay.py
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
"""overlay_cached_prefix: freeze must forward the CACHED (compressed) bytes.
|
||||
|
||||
The freeze path can emit the agent's ORIGINAL bytes for a frozen message, but
|
||||
the provider cached whatever we FORWARDED last turn (the compressed form).
|
||||
Forwarding original then mismatches the cached prefix and busts the prompt cache
|
||||
(observed: 100% of misses were this ``prefix_change``, ~56% of all cache-writes).
|
||||
``overlay_cached_prefix`` replays the previously-forwarded prefix byte-identical
|
||||
so the cache still hits — in BOTH proxy modes.
|
||||
"""
|
||||
|
||||
from headroom.cache.prefix_tracker import overlay_cached_prefix
|
||||
|
||||
|
||||
def M(role, text):
|
||||
return {"role": role, "content": text}
|
||||
|
||||
|
||||
# Previous turn: 2 messages. Original was big; we FORWARDED the compressed form,
|
||||
# so that compressed form is what the provider cached.
|
||||
PREV_ORIG = [M("user", "READ foo.py:\n<2000 original lines>"), M("assistant", "ok")]
|
||||
PREV_FWD = [M("user", "READ foo.py:\n<compressed>"), M("assistant", "ok")]
|
||||
# This turn: agent appended one new message (append-only growth).
|
||||
CUR_ORIG = PREV_ORIG + [M("user", "grep result:\n<800 original lines>")]
|
||||
# What apply() produced in the buggy freeze path: ORIGINAL bytes for the frozen
|
||||
# prefix (== PREV_ORIG) + compressed new tail.
|
||||
OPTIMIZED_BUGGY = [PREV_ORIG[0], PREV_ORIG[1], M("user", "grep result:\n<compressed>")]
|
||||
|
||||
|
||||
def test_replays_cached_compressed_prefix_byte_identical():
|
||||
out = overlay_cached_prefix(OPTIMIZED_BUGGY, CUR_ORIG, PREV_ORIG, PREV_FWD)
|
||||
# The frozen prefix now equals what the provider cached (compressed), NOT the
|
||||
# agent's original bytes → cache hits instead of busting.
|
||||
assert out[:2] == PREV_FWD
|
||||
assert out[:2] != PREV_ORIG
|
||||
# This turn's compressed tail is preserved.
|
||||
assert out[2] == OPTIMIZED_BUGGY[2]
|
||||
assert len(out) == len(CUR_ORIG)
|
||||
|
||||
|
||||
def test_is_a_noop_relative_to_cache_when_already_correct():
|
||||
# If the freeze path already forwarded the compressed (cached) prefix, the
|
||||
# overlay reproduces exactly that — idempotent.
|
||||
already_correct = [PREV_FWD[0], PREV_FWD[1], M("user", "grep result:\n<compressed>")]
|
||||
out = overlay_cached_prefix(already_correct, CUR_ORIG, PREV_ORIG, PREV_FWD)
|
||||
assert out == already_correct
|
||||
|
||||
|
||||
def test_not_append_only_returns_unchanged():
|
||||
# An early message changed → previous forwarded bytes may not correspond to
|
||||
# the same positions; do NOT overlay (accept a possible bust over corruption).
|
||||
changed = [M("user", "TOTALLY DIFFERENT"), PREV_ORIG[1], M("user", "x")]
|
||||
out = overlay_cached_prefix(OPTIMIZED_BUGGY, changed, PREV_ORIG, PREV_FWD)
|
||||
assert out == OPTIMIZED_BUGGY
|
||||
|
||||
|
||||
def test_no_previous_state_returns_unchanged():
|
||||
assert overlay_cached_prefix(OPTIMIZED_BUGGY, CUR_ORIG, None, None) == OPTIMIZED_BUGGY
|
||||
assert overlay_cached_prefix(OPTIMIZED_BUGGY, CUR_ORIG, [], []) == OPTIMIZED_BUGGY
|
||||
|
||||
|
||||
def test_forwarded_count_mismatch_returns_unchanged():
|
||||
# Defensive: not exactly one forwarded message per original → bail.
|
||||
assert (
|
||||
overlay_cached_prefix(OPTIMIZED_BUGGY, CUR_ORIG, PREV_ORIG, PREV_FWD[:1]) == OPTIMIZED_BUGGY
|
||||
)
|
||||
|
||||
|
||||
def test_shorter_current_or_optimized_returns_unchanged():
|
||||
assert overlay_cached_prefix([M("user", "x")], [M("user", "x")], PREV_ORIG, PREV_FWD) == [
|
||||
M("user", "x")
|
||||
]
|
||||
|
||||
|
||||
def test_cache_hit_property_prefix_matches_last_forward():
|
||||
# The invariant that guarantees a cache hit: forwarded[:n] this turn ==
|
||||
# forwarded[:n] last turn (== what the provider cached).
|
||||
out = overlay_cached_prefix(OPTIMIZED_BUGGY, CUR_ORIG, PREV_ORIG, PREV_FWD)
|
||||
n = len(PREV_FWD)
|
||||
assert out[:n] == PREV_FWD # exact byte-identical prefix → provider cache hit
|
||||
140
tests/test_cross_turn_cache_safety.py
Normal file
140
tests/test_cross_turn_cache_safety.py
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
"""Cross-turn cache-safety invariant — the test class that catches cache busts.
|
||||
|
||||
Why the +150%-cache_create / +41%-cost bug slipped through: every prior cache
|
||||
test was SINGLE-turn and used a fake tracker, so nobody exercised the real
|
||||
multi-turn invariant that actually governs prompt-cache cost:
|
||||
|
||||
Across append-only turns, the forwarded prefix must stay BYTE-IDENTICAL to
|
||||
what was forwarded (and cached) last turn — otherwise the provider re-creates
|
||||
the whole suffix (a cache bust) instead of reading it.
|
||||
|
||||
This simulates the provider's prefix cache (longest byte-identical leading run of
|
||||
messages = cache_read; the rest = cache_create) and drives the REAL
|
||||
``PrefixCacheTracker`` + the freeze model + ``overlay_cached_prefix`` over several
|
||||
turns. It asserts the invariant directly, and proves the guard is load-bearing:
|
||||
WITHOUT the overlay the freeze forwards the agent's original bytes and busts every
|
||||
turn; WITH it the prefix stays stable.
|
||||
"""
|
||||
|
||||
from headroom.cache.prefix_tracker import (
|
||||
PrefixCacheTracker,
|
||||
PrefixFreezeConfig,
|
||||
overlay_cached_prefix,
|
||||
)
|
||||
|
||||
|
||||
def _toklen(m) -> int:
|
||||
return max(1, len(str(m.get("content", ""))))
|
||||
|
||||
|
||||
def _compress(m):
|
||||
"""Deterministic stand-in for a real compressor (kompress is deterministic
|
||||
per content via the result cache): shrink the content by half."""
|
||||
c = str(m.get("content", ""))
|
||||
return {**m, "content": c[: max(1, len(c) // 2)]}
|
||||
|
||||
|
||||
def _apply_freeze(original, frozen_count):
|
||||
"""Faithful model of pipeline.apply()'s freeze: the frozen prefix is
|
||||
forwarded as the agent's ORIGINAL bytes; everything else is compressed.
|
||||
(Mirrors content_router.py: `result_slots[i] = message` for i < frozen.)"""
|
||||
return [
|
||||
(original[i] if i < frozen_count else _compress(original[i])) for i in range(len(original))
|
||||
]
|
||||
|
||||
|
||||
def _provider_cache_read(forwarded, prev_forwarded):
|
||||
"""Longest byte-identical leading run of messages the provider can serve from
|
||||
cache, in tokens. A single differing message breaks the prefix (bust)."""
|
||||
if not prev_forwarded:
|
||||
return 0
|
||||
matched = 0
|
||||
for a, b in zip(forwarded, prev_forwarded):
|
||||
if a == b:
|
||||
matched += _toklen(a)
|
||||
else:
|
||||
break
|
||||
return matched
|
||||
|
||||
|
||||
def _drive_turns(*, use_overlay: bool, turns: int = 5):
|
||||
"""Return per-turn (expected_cache_read, actual_cache_read). A bust is any
|
||||
turn where actual < expected (the previously-cached prefix wasn't reused)."""
|
||||
# min_cached_tokens=0 so freeze activates from turn 2 regardless of size.
|
||||
tracker = PrefixCacheTracker("anthropic", PrefixFreezeConfig(min_cached_tokens=0))
|
||||
convo: list[dict] = []
|
||||
prev_forwarded: list[dict] | None = None
|
||||
out = []
|
||||
for t in range(1, turns + 1):
|
||||
# Append-only growth: one new large tool output per turn.
|
||||
convo = convo + [{"role": "user", "content": f"tool-output-turn-{t}:" + "X" * 400}]
|
||||
|
||||
frozen = tracker.get_frozen_message_count()
|
||||
forwarded = _apply_freeze(convo, frozen)
|
||||
if use_overlay:
|
||||
forwarded = overlay_cached_prefix(
|
||||
forwarded,
|
||||
convo,
|
||||
tracker.get_last_original_messages(),
|
||||
tracker.get_last_forwarded_messages(),
|
||||
)
|
||||
|
||||
expected_read = sum(_toklen(m) for m in prev_forwarded) if prev_forwarded else 0
|
||||
actual_read = _provider_cache_read(forwarded, prev_forwarded)
|
||||
out.append((expected_read, actual_read))
|
||||
|
||||
counts = [_toklen(m) for m in forwarded]
|
||||
write = sum(counts) - actual_read
|
||||
tracker.update_from_response(
|
||||
actual_read, write, forwarded, message_token_counts=counts, original_messages=convo
|
||||
)
|
||||
prev_forwarded = forwarded
|
||||
return out
|
||||
|
||||
|
||||
def test_freeze_busts_cache_every_turn_without_overlay():
|
||||
"""Proves the test is load-bearing: the raw freeze path busts the cache."""
|
||||
results = _drive_turns(use_overlay=False)
|
||||
# From turn 2 on, a hit was expected but the prefix broke (actual < expected).
|
||||
busts = [exp > act for (exp, act) in results[1:]]
|
||||
assert any(busts), "expected the un-fixed freeze path to bust the prefix cache"
|
||||
|
||||
|
||||
def test_overlay_keeps_prefix_byte_identical_no_bust():
|
||||
"""The fix: every turn reuses the full previously-cached prefix — no bust."""
|
||||
results = _drive_turns(use_overlay=True)
|
||||
for exp, act in results[1:]:
|
||||
assert act >= exp, (
|
||||
f"cache bust: expected to read {exp} cached tokens but only read {act} "
|
||||
"— forwarded prefix diverged from last turn"
|
||||
)
|
||||
|
||||
|
||||
def test_cache_create_stays_bounded_to_the_delta_with_overlay():
|
||||
"""Cost proxy: with the fix, per-turn cache_create ≈ the new delta only, not
|
||||
the whole re-created prefix (which is what drove +150% cache_create)."""
|
||||
tracker = PrefixCacheTracker("anthropic", PrefixFreezeConfig(min_cached_tokens=0))
|
||||
convo: list[dict] = []
|
||||
prev_forwarded: list[dict] | None = None
|
||||
creates = []
|
||||
for t in range(1, 6):
|
||||
convo = convo + [{"role": "user", "content": f"turn-{t}:" + "X" * 400}]
|
||||
frozen = tracker.get_frozen_message_count()
|
||||
forwarded = overlay_cached_prefix(
|
||||
_apply_freeze(convo, frozen),
|
||||
convo,
|
||||
tracker.get_last_original_messages(),
|
||||
tracker.get_last_forwarded_messages(),
|
||||
)
|
||||
read = _provider_cache_read(forwarded, prev_forwarded)
|
||||
counts = [_toklen(m) for m in forwarded]
|
||||
create = sum(counts) - read
|
||||
creates.append(create)
|
||||
tracker.update_from_response(
|
||||
read, create, forwarded, message_token_counts=counts, original_messages=convo
|
||||
)
|
||||
prev_forwarded = forwarded
|
||||
# Steady-state cache_create per turn should be ~one delta message, NOT growing
|
||||
# with conversation length. Assert the last turn creates no more than the
|
||||
# first (which had no cache to reuse).
|
||||
assert creates[-1] <= creates[0] + 1
|
||||
|
|
@ -490,7 +490,7 @@ def test_existing_retrieve_tool_keeps_reversible_ccr_path_when_prefix_is_frozen(
|
|||
assert [tool["name"] for tool in forwarded["tools"]] == ["headroom_retrieve"]
|
||||
|
||||
|
||||
def test_cache_mode_skip_forwards_original_prefix_when_tool_injection_is_deferred(
|
||||
def test_cache_mode_skip_replays_cached_compressed_prefix_when_tool_injection_is_deferred(
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
|
|
@ -576,11 +576,15 @@ def test_cache_mode_skip_forwards_original_prefix_when_tool_injection_is_deferre
|
|||
assert response.status_code == 200
|
||||
assert captured.get("compression_calls", []) == []
|
||||
forwarded = captured["body"]
|
||||
assert forwarded["messages"] == original_messages
|
||||
# Tool injection is deferred (no CCR tool this turn), but the frozen
|
||||
# prefix was cached COMPRESSED last turn. Replay it byte-identical so the
|
||||
# prompt cache still hits instead of busting on original bytes (#1850);
|
||||
# the mutable tail stays original. Tool absent AND cache intact.
|
||||
assert forwarded["messages"] == previous_forwarded_messages + original_messages[1:]
|
||||
assert "tools" not in forwarded
|
||||
|
||||
|
||||
def test_cache_mode_exact_prefix_replay_forwards_original_messages_when_tool_injection_is_deferred(
|
||||
def test_cache_mode_exact_prefix_replay_forwards_cached_compressed_prefix_when_tool_injection_is_deferred(
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
|
|
@ -663,7 +667,10 @@ def test_cache_mode_exact_prefix_replay_forwards_original_messages_when_tool_inj
|
|||
assert response.status_code == 200
|
||||
assert captured.get("compression_calls", []) == []
|
||||
forwarded = captured["body"]
|
||||
assert forwarded["messages"] == original_messages
|
||||
# Deferred injection (no CCR tool), single frozen message cached
|
||||
# COMPRESSED last turn: replay it so the cache holds instead of busting
|
||||
# on original bytes (#1850). Tool absent AND cache intact.
|
||||
assert forwarded["messages"] == previous_forwarded_messages
|
||||
assert "tools" not in forwarded
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -877,3 +877,81 @@ def test_resolve_ccr_workspace_malformed_request_returns_empty() -> None:
|
|||
key, label = AnthropicHandlerMixin._resolve_ccr_workspace(request, body)
|
||||
assert key == ""
|
||||
assert label is None
|
||||
|
||||
|
||||
class TestHasNewCcrMarkers:
|
||||
"""#1850: replayed (overlay) markers must not count as new-this-turn.
|
||||
|
||||
``overlay_cached_prefix`` replays the previously-forwarded compressed prefix
|
||||
byte-identical to keep the messages cache warm — which reintroduces its old
|
||||
``hash=…`` markers. If those replayed markers counted as "new", the handler
|
||||
would re-inject the retrieve tool every frozen turn and bust the *tools*
|
||||
cache. ``has_new_ccr_markers`` filters them out.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _hashes(*contents: str) -> list[str]:
|
||||
from headroom.ccr.tool_injection import CCRToolInjector
|
||||
|
||||
inj = CCRToolInjector(
|
||||
provider="anthropic", inject_tool=False, inject_system_instructions=False
|
||||
)
|
||||
inj.scan_for_markers([{"role": "user", "content": c} for c in contents])
|
||||
return inj.detected_hashes
|
||||
|
||||
def test_replayed_markers_are_not_new(self):
|
||||
from headroom.proxy.helpers import has_new_ccr_markers
|
||||
|
||||
marker = "[100 items compressed to 10. Retrieve more: hash=abc123def456abc123def456]"
|
||||
current = self._hashes(marker)
|
||||
assert current, "sanity: the marker must be detected"
|
||||
# Every marker was already in what we forwarded last turn → nothing new.
|
||||
assert (
|
||||
has_new_ccr_markers(
|
||||
current_detected_hashes=current,
|
||||
previous_forwarded_messages=[{"role": "user", "content": marker}],
|
||||
provider="anthropic",
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
def test_genuinely_new_marker_is_detected(self):
|
||||
from headroom.proxy.helpers import has_new_ccr_markers
|
||||
|
||||
old = "[100 items compressed to 10. Retrieve more: hash=abc123def456abc123def456]"
|
||||
new = "[50 items compressed to 5. Retrieve more: hash=deadbeefdeadbeefdeadbeef]"
|
||||
current = self._hashes(old, new)
|
||||
# Only `old` was forwarded before; `new` is fresh → override must fire.
|
||||
assert (
|
||||
has_new_ccr_markers(
|
||||
current_detected_hashes=current,
|
||||
previous_forwarded_messages=[{"role": "user", "content": old}],
|
||||
provider="anthropic",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_no_previous_forward_means_all_new(self):
|
||||
from headroom.proxy.helpers import has_new_ccr_markers
|
||||
|
||||
marker = "[100 items compressed to 10. Retrieve more: hash=abc123def456abc123def456]"
|
||||
assert (
|
||||
has_new_ccr_markers(
|
||||
current_detected_hashes=self._hashes(marker),
|
||||
previous_forwarded_messages=None,
|
||||
provider="anthropic",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_no_markers_means_nothing_new(self):
|
||||
from headroom.proxy.helpers import has_new_ccr_markers
|
||||
|
||||
assert (
|
||||
has_new_ccr_markers(
|
||||
current_detected_hashes=[],
|
||||
previous_forwarded_messages=None,
|
||||
provider="anthropic",
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
|
|
|||
|
|
@ -21,6 +21,16 @@ class _FakePrefixTracker:
|
|||
def get_frozen_message_count(self) -> int:
|
||||
return self._frozen_count
|
||||
|
||||
# Empty history → overlay_cached_prefix() is a no-op here, so these tests
|
||||
# keep asserting the cache-freeze behavior they always have. The cross-turn
|
||||
# overlay itself is exercised in test_cross_turn_cache_safety.py against the
|
||||
# real tracker; these stubs just satisfy the handler's overlay call.
|
||||
def get_last_original_messages(self): # noqa: ANN201
|
||||
return []
|
||||
|
||||
def get_last_forwarded_messages(self): # noqa: ANN201
|
||||
return []
|
||||
|
||||
def update_from_response(self, **kwargs): # noqa: ANN003
|
||||
return None
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue