mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
fix(proxy): keep cache_control bounded + stable so the freeze overlay stops busting (#1852)
Some checks failed
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime name:]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code name:code]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim name:code-slim]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-nonroot name:nonroot]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim name:slim]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime name:]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code name:code]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim name:code-slim]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-nonroot name:nonroot]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim name:slim]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Waiting to run
Docker / docker-manifest (map[bake_target:runtime name:]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-code name:code]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-code-slim name:code-slim]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-nonroot name:nonroot]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-slim name:slim]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Blocked by required conditions
Docker / promote-latest (push) Blocked by required conditions
CI / changes (push) Failing after 7s
Init E2E / docker-init-e2e (push) Failing after 7s
CI / commitlint (push) Failing after 3s
Init Native E2E / init-native (ubuntu-latest, claude) (push) Failing after 6s
Init Native E2E / init-native (ubuntu-latest, codex) (push) Failing after 5s
Install Native E2E / install-native (ubuntu-latest) (push) Failing after 4s
Init Native E2E / init-native (ubuntu-latest, copilot) (push) Failing after 6s
Merge Conflicts / merge-conflicts (push) Failing after 6s
Release Please / release-please (push) Failing after 5s
rust / test (ubuntu) (push) Failing after 5s
rust / wheels (x86_64-unknown-linux-gnu) (push) Failing after 6s
rust / audit (push) Failing after 5s
rust / parity (nightly, allowed to fail during Phase 0) (push) Failing after 2s
Security / CodeQL (python) (push) Failing after 25s
Security / CodeQL (javascript-typescript) (push) Failing after 30s
Security / Dependency audit (pip-audit) (push) Failing after 40s
Wrap E2E / docker-wrap-e2e (push) Failing after 4s
Wrap Native E2E / wrap-native (ubuntu-latest) (push) Failing after 8s
Security / Secret scan (gitleaks) (push) Failing after 25s
CI / lint (push) Failing after 3s
CI / build-wheel (push) Failing after 3s
CI / prefetch-model (push) Failing after 2s
CI / build (push) Failing after 4s
CI / workflow-validation (push) Failing after 4s
CI / docker-native-e2e (push) Failing after 4s
CI / test-extras (push) Failing after 1s
CI / test-agno (push) Failing after 2s
CI / test-dashboard-ui (push) Failing after 2s
CI / test (1) (push) Failing after 1s
CI / test (2) (push) Failing after 3s
CI / test (3) (push) Failing after 2s
CI / test (4) (push) Failing after 2s
CI / windows-native-wrapper (push) Has been cancelled
CI / macos-native-wrapper (push) Has been cancelled
Init Native E2E / init-native (macos-latest, claude) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, codex) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, copilot) (push) Has been cancelled
Install Native E2E / install-native (macos-latest) (push) Has been cancelled
rust / wheels (aarch64-apple-darwin) (push) Has been cancelled
rust / wheels (x86_64-apple-darwin) (push) Has been cancelled
Wrap Native E2E / wrap-native (macos-latest) (push) Has been cancelled
Some checks failed
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime name:]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code name:code]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim name:code-slim]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-nonroot name:nonroot]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim name:slim]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime name:]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code name:code]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim name:code-slim]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-nonroot name:nonroot]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim name:slim]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Waiting to run
Docker / docker-manifest (map[bake_target:runtime name:]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-code name:code]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-code-slim name:code-slim]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-nonroot name:nonroot]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-slim name:slim]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Blocked by required conditions
Docker / promote-latest (push) Blocked by required conditions
CI / changes (push) Failing after 7s
Init E2E / docker-init-e2e (push) Failing after 7s
CI / commitlint (push) Failing after 3s
Init Native E2E / init-native (ubuntu-latest, claude) (push) Failing after 6s
Init Native E2E / init-native (ubuntu-latest, codex) (push) Failing after 5s
Install Native E2E / install-native (ubuntu-latest) (push) Failing after 4s
Init Native E2E / init-native (ubuntu-latest, copilot) (push) Failing after 6s
Merge Conflicts / merge-conflicts (push) Failing after 6s
Release Please / release-please (push) Failing after 5s
rust / test (ubuntu) (push) Failing after 5s
rust / wheels (x86_64-unknown-linux-gnu) (push) Failing after 6s
rust / audit (push) Failing after 5s
rust / parity (nightly, allowed to fail during Phase 0) (push) Failing after 2s
Security / CodeQL (python) (push) Failing after 25s
Security / CodeQL (javascript-typescript) (push) Failing after 30s
Security / Dependency audit (pip-audit) (push) Failing after 40s
Wrap E2E / docker-wrap-e2e (push) Failing after 4s
Wrap Native E2E / wrap-native (ubuntu-latest) (push) Failing after 8s
Security / Secret scan (gitleaks) (push) Failing after 25s
CI / lint (push) Failing after 3s
CI / build-wheel (push) Failing after 3s
CI / prefetch-model (push) Failing after 2s
CI / build (push) Failing after 4s
CI / workflow-validation (push) Failing after 4s
CI / docker-native-e2e (push) Failing after 4s
CI / test-extras (push) Failing after 1s
CI / test-agno (push) Failing after 2s
CI / test-dashboard-ui (push) Failing after 2s
CI / test (1) (push) Failing after 1s
CI / test (2) (push) Failing after 3s
CI / test (3) (push) Failing after 2s
CI / test (4) (push) Failing after 2s
CI / windows-native-wrapper (push) Has been cancelled
CI / macos-native-wrapper (push) Has been cancelled
Init Native E2E / init-native (macos-latest, claude) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, codex) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, copilot) (push) Has been cancelled
Install Native E2E / install-native (macos-latest) (push) Has been cancelled
rust / wheels (aarch64-apple-darwin) (push) Has been cancelled
rust / wheels (x86_64-apple-darwin) (push) Has been cancelled
Wrap Native E2E / wrap-native (macos-latest) (push) Has been cancelled
Follow-up to #1850. Two residual cache-bust sources, both `cache_control`-related: 1. **Guard too strict.** `overlay_cached_prefix` decided "is this turn an append-only extension?" by comparing whole message dicts — including `cache_control`. Clients (Claude Code, litellm) move the cache breakpoint to the newest message every call, so a marker landing in the frozen prefix made the guard fail, the overlay skip its replay, and the raw freeze forward ORIGINAL bytes over the cached COMPRESSED prefix → partial bust (the ~42% residual on the a10 run, `prefix_change=0`). Fix: run the append-only guard on **content only** (strip `cache_control` before comparing) — content is what the provider's cache keys on. 2. **Marker accumulation.** The overlay replays the markers that rode on each turn's then-newest message, so `cache_control` blocks pile up ~1/turn; Anthropic hard-errors at >4 total. Fix: `normalize_message_cache_control` strips every message-level marker and re-places a single ephemeral breakpoint on the last block (one breakpoint caches the whole prefix; cache is content-keyed so re-placing never busts). Wired into the Anthropic handler after the overlay. **Per-provider (deliberately scoped):** - **Anthropic**: `cache_control` markers → both fixes apply. - **OpenAI**: AUTOMATIC prefix caching, no markers → overlay (byte-identity) only; normalize is NOT applied (Anthropic markers on an OpenAI request would be wrong). - **Bedrock**: serves Claude via the pipeline but has no cachePoint/freeze-replay path → not affected; a cachePoint analog would be needed if caching is expanded. - **Gemini**: explicit Cache API (`cachedContent`), no inline markers/freeze → N/A. > Stacked on #1850 — review that first; the diff against `main` includes its overlay + `has_new_ccr_markers` work. ## Description Keeps the freeze overlay's cache-safety intact against real clients that relocate the `cache_control` breakpoint each turn, and prevents `cache_control` blocks from accumulating past Anthropic's 4-marker limit. See the two fixes above. Closes #<!-- none --> — follow-up to #1850 (no separate issue). ## 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/cache/prefix_tracker.py`: append-only guard in `overlay_cached_prefix` now compares **content only** (ignores `cache_control`); new `normalize_message_cache_control()` collapses message-level markers to a single ephemeral breakpoint on the last block. - `headroom/proxy/handlers/anthropic.py`: apply `normalize_message_cache_control` after the overlay (Anthropic only). - `tests/test_cache_control_move_bust.py`: reproduces the moved-marker bust + proves both fixes. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed (local, see below) ### Test Output ```text $ pytest tests/test_cache_control_move_bust.py -q ....... [100%] 7 passed in 0.19s # broader cache-safety suite (overlay + cross-turn + CCR deferred + openai/anthropic cache-stability + helpers) $ pytest tests/test_cache_control_move_bust.py tests/test_cache_prefix_overlay.py \ tests/test_cross_turn_cache_safety.py tests/test_proxy/test_anthropic_ccr_deferred_injection.py \ tests/test_proxy_handler_helpers.py tests/test_proxy_openai_cache_stability.py \ tests/test_proxy_anthropic_cache_stability.py -q 91 passed, 2 warnings in 29.98s $ ruff check . # ruff 0.15.17 (CI-pinned) All checks passed! $ ruff format --check . # ruff 0.15.17 1057 files already formatted $ mypy headroom --ignore-missing-imports Success: no issues found (changed modules: prefix_tracker, anthropic, openai, helpers) ``` ## Real Behavior Proof - **Environment:** local (`.venv`, Python 3.12), ruff 0.15.17 / mypy pinned to CI versions. - **Exact command / steps:** `tests/test_cache_control_move_bust.py` drives the REAL tracker + freeze + `overlay_cached_prefix` + `normalize_message_cache_control` across multiple append-only turns where the client moves the `cache_control` breakpoint each turn. - **Observed result:** with a moved marker in the frozen prefix, the content-only guard keeps the overlay replaying (forwarded prefix stays byte-identical → no bust); `cache_control` blocks stay ≤4 across many turns and content is never altered. The reproduction test fails without the fix and passes with it. - **Not tested (this PR):** the end-to-end a10 SWE-bench run is the field observation motivating fix #1 (~42% residual, `prefix_change=0`); not re-run here. ## 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 - [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 ## Additional Notes Stacked on #1850; land that first. Docs/CHANGELOG untouched (behavioral cache-safety fix; no user-facing surface change). N/A: no screenshots (no UI).
This commit is contained in:
parent
5d14080c94
commit
48201345be
3 changed files with 278 additions and 4 deletions
77
headroom/cache/prefix_tracker.py
vendored
77
headroom/cache/prefix_tracker.py
vendored
|
|
@ -112,6 +112,22 @@ class CacheMissAttribution:
|
|||
ttl_exceeded: bool = False
|
||||
|
||||
|
||||
def _strip_cache_control(obj: Any) -> Any:
|
||||
"""Recursively drop ``cache_control`` for content-only equality checks.
|
||||
|
||||
Clients (notably Claude Code) move the cache_control breakpoint to the newest
|
||||
message on every call, so the exact same message carries cache_control on one
|
||||
turn and not the next. That per-call annotation must be ignored when deciding
|
||||
whether this turn append-only-extends the previous one — otherwise a moved
|
||||
marker spuriously fails the check and we skip the byte-identical replay,
|
||||
busting the cache."""
|
||||
if isinstance(obj, dict):
|
||||
return {k: _strip_cache_control(v) for k, v in obj.items() if k != "cache_control"}
|
||||
if isinstance(obj, list):
|
||||
return [_strip_cache_control(v) for v in obj]
|
||||
return obj
|
||||
|
||||
|
||||
def overlay_cached_prefix(
|
||||
optimized_messages: list[dict[str, Any]],
|
||||
current_original_messages: list[dict[str, Any]],
|
||||
|
|
@ -151,13 +167,68 @@ def overlay_cached_prefix(
|
|||
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:
|
||||
# Append-only guard on CONTENT ONLY: the frozen region must be the same
|
||||
# messages we cached. Compare with cache_control stripped — clients move that
|
||||
# breakpoint to the newest message each turn, so a raw dict compare would
|
||||
# spuriously fail whenever a marker lands in the frozen prefix, skip the
|
||||
# replay, and bust the cache (the residual busts observed after the first
|
||||
# fix). Content stability is what the provider's prefix cache actually keys on.
|
||||
if _strip_cache_control(current_original_messages[:n]) != _strip_cache_control(prev_orig):
|
||||
return optimized_messages
|
||||
# Replay the cached (compressed) prefix; keep this turn's compressed tail.
|
||||
# Replay the cached (compressed) prefix byte-identical; keep this turn's tail.
|
||||
return list(prev_fwd) + list(optimized_messages[n:])
|
||||
|
||||
|
||||
def normalize_message_cache_control(
|
||||
messages: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Own message-level cache_control placement so breakpoints stay bounded.
|
||||
|
||||
Two forces pile up cache_control markers turn over turn: clients move the
|
||||
breakpoint to the newest message each call, and ``overlay_cached_prefix``
|
||||
replays the markers that rode on each turn's then-newest message. Anthropic
|
||||
hard-errors at **>4 cache_control blocks total** (system + tools + messages),
|
||||
so on a long conversation the accumulation eventually 400s.
|
||||
|
||||
Fix: strip EVERY message-level cache_control and re-place a **single**
|
||||
ephemeral breakpoint on the last block of the last block-style message. One
|
||||
breakpoint caches the whole message prefix up to it, and — because the
|
||||
provider's cache key is message CONTENT, not marker presence (moving the
|
||||
breakpoint forward is the documented client pattern and it hits) — stripping
|
||||
and re-placing markers never busts. system/tools breakpoints live outside
|
||||
``messages`` and are left untouched (they still count toward the 4 limit, so
|
||||
holding messages to one breakpoint leaves room for them).
|
||||
|
||||
Only block-style (list) content can carry cache_control; string content is
|
||||
left as-is. Returns the input unchanged when there is nothing to normalize.
|
||||
"""
|
||||
changed = False
|
||||
out: list[dict[str, Any]] = []
|
||||
last_block_idx = -1
|
||||
for i, msg in enumerate(messages):
|
||||
content = msg.get("content") if isinstance(msg, dict) else None
|
||||
if isinstance(content, list):
|
||||
had = any(isinstance(b, dict) and "cache_control" in b for b in content)
|
||||
stripped = [
|
||||
{k: v for k, v in b.items() if k != "cache_control"} if isinstance(b, dict) else b
|
||||
for b in content
|
||||
]
|
||||
out.append({**msg, "content": stripped} if had else msg)
|
||||
changed = changed or had
|
||||
if stripped and isinstance(stripped[-1], dict):
|
||||
last_block_idx = i
|
||||
else:
|
||||
out.append(msg)
|
||||
# Re-place exactly one breakpoint on the last block-style message.
|
||||
if last_block_idx >= 0:
|
||||
msg = out[last_block_idx]
|
||||
content = list(msg["content"])
|
||||
content[-1] = {**content[-1], "cache_control": {"type": "ephemeral"}}
|
||||
out[last_block_idx] = {**msg, "content": content}
|
||||
changed = True
|
||||
return out if changed else messages
|
||||
|
||||
|
||||
class PrefixCacheTracker:
|
||||
"""Tracks provider prefix cache state across turns in a session.
|
||||
|
||||
|
|
|
|||
|
|
@ -1348,7 +1348,10 @@ class AnthropicHandlerMixin:
|
|||
# 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
|
||||
from headroom.cache.prefix_tracker import (
|
||||
normalize_message_cache_control,
|
||||
overlay_cached_prefix,
|
||||
)
|
||||
|
||||
_ov = overlay_cached_prefix(
|
||||
optimized_messages,
|
||||
|
|
@ -1360,6 +1363,16 @@ class AnthropicHandlerMixin:
|
|||
optimized_messages = _ov
|
||||
optimized_tokens = tokenizer.count_messages(optimized_messages)
|
||||
|
||||
# Own cache_control placement: the client moves the breakpoint each
|
||||
# turn and the overlay replays past markers, so they accumulate ~1/turn
|
||||
# and Anthropic hard-errors at >4. Strip message-level markers and keep
|
||||
# a single breakpoint on the last block (caches the whole prefix;
|
||||
# content-keyed cache so re-placing never busts). Applied last so the
|
||||
# forwarded AND recorded (next_forwarded) messages stay bounded.
|
||||
_norm = normalize_message_cache_control(optimized_messages)
|
||||
if _norm is not optimized_messages:
|
||||
optimized_messages = _norm
|
||||
|
||||
# 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):
|
||||
|
|
|
|||
190
tests/test_cache_control_move_bust.py
Normal file
190
tests/test_cache_control_move_bust.py
Normal file
|
|
@ -0,0 +1,190 @@
|
|||
"""Reproduce the residual cache-bust: client MOVING cache_control defeats the
|
||||
prefix overlay.
|
||||
|
||||
Real clients (Claude Code, litellm) move the cache_control breakpoint to the
|
||||
newest message every turn — so a message that was marked last turn is unmarked
|
||||
this turn (its dict bytes change). The first overlay fix compared *raw* message
|
||||
dicts for its append-only guard, so a moved marker in the frozen prefix made the
|
||||
guard fail → the overlay skipped the replay → the raw freeze forwarded ORIGINAL
|
||||
bytes over the cached COMPRESSED prefix → partial bust (the ~42% residual seen
|
||||
on the a10 run, with prefix_change=0).
|
||||
|
||||
These tests pin the exact scenario, prove the content-only guard fixes it, and
|
||||
document the remaining piece (marker accumulation > 4 → needs stable placement).
|
||||
"""
|
||||
|
||||
from headroom.cache.prefix_tracker import (
|
||||
PrefixCacheTracker,
|
||||
PrefixFreezeConfig,
|
||||
overlay_cached_prefix,
|
||||
)
|
||||
|
||||
|
||||
def M(role, text, cc=False):
|
||||
m = {"role": role, "content": text}
|
||||
if cc:
|
||||
m["cache_control"] = {"type": "ephemeral"}
|
||||
return m
|
||||
|
||||
|
||||
def _toklen(m):
|
||||
return max(1, len(str(m.get("content", ""))))
|
||||
|
||||
|
||||
def _compress(m):
|
||||
c = str(m.get("content", ""))
|
||||
return {**m, "content": c[: max(1, len(c) // 2)]}
|
||||
|
||||
|
||||
def _freeze(original, frozen):
|
||||
# content_router freeze model: frozen prefix = ORIGINAL bytes, rest compressed.
|
||||
return [(original[i] if i < frozen else _compress(original[i])) for i in range(len(original))]
|
||||
|
||||
|
||||
# ── Unit reproduction ────────────────────────────────────────────────────────
|
||||
# Last turn we forwarded the compressed prefix; the client had marked msg1.
|
||||
PREV_ORIG = [M("user", "READ foo:\n<big>"), M("assistant", "ok", cc=True)]
|
||||
PREV_FWD = [M("user", "READ foo:\n<compressed>"), M("assistant", "ok", cc=True)]
|
||||
# This turn the client MOVED the marker off msg1 onto the new last message (msg2).
|
||||
CUR_ORIG = [M("user", "READ foo:\n<big>"), M("assistant", "ok"), M("user", "grep:\n<big>", cc=True)]
|
||||
# Freeze forwarded ORIGINAL bytes for the frozen prefix + compressed tail.
|
||||
OPTIMIZED = [M("user", "READ foo:\n<big>"), M("assistant", "ok"), M("user", "grep:\n<compressed>")]
|
||||
|
||||
|
||||
def test_marker_move_would_fail_a_raw_dict_guard():
|
||||
# This is the exact condition the old (raw) guard tripped on: the frozen
|
||||
# prefix differs ONLY because cache_control moved off msg1.
|
||||
assert CUR_ORIG[:2] != PREV_ORIG
|
||||
# ...but with cache_control stripped, the content is an append-only extension.
|
||||
from headroom.cache.prefix_tracker import _strip_cache_control
|
||||
|
||||
assert _strip_cache_control(CUR_ORIG[:2]) == _strip_cache_control(PREV_ORIG)
|
||||
|
||||
|
||||
def test_overlay_replays_despite_moved_marker():
|
||||
out = overlay_cached_prefix(OPTIMIZED, CUR_ORIG, PREV_ORIG, PREV_FWD)
|
||||
# The content-only guard lets the replay happen: the forwarded prefix is now
|
||||
# byte-identical to what the provider cached (compressed), NOT the freeze's
|
||||
# original bytes → cache hits instead of busting.
|
||||
assert out[:2] == PREV_FWD
|
||||
assert out[:2] != OPTIMIZED[:2]
|
||||
assert out[2] == OPTIMIZED[2] # compressed tail preserved
|
||||
|
||||
|
||||
# ── Cross-turn: client moves the marker every turn, provider keys on full bytes ─
|
||||
def _client_convo(t):
|
||||
msgs = [{"role": "user", "content": f"turn-{k}:" + "X" * 300} for k in range(1, t + 1)]
|
||||
msgs[-1] = {**msgs[-1], "cache_control": {"type": "ephemeral"}} # mark ONLY the newest
|
||||
return msgs
|
||||
|
||||
|
||||
def _cache_read(fwd, prev_fwd):
|
||||
# cache_control-AWARE (worst case): a moved marker changes the block's bytes,
|
||||
# so it breaks the byte-identical prefix.
|
||||
if not prev_fwd:
|
||||
return 0
|
||||
matched = 0
|
||||
for a, b in zip(fwd, prev_fwd):
|
||||
if a == b:
|
||||
matched += _toklen(a)
|
||||
else:
|
||||
break
|
||||
return matched
|
||||
|
||||
|
||||
def _drive(use_overlay, turns=5):
|
||||
tracker = PrefixCacheTracker("anthropic", PrefixFreezeConfig(min_cached_tokens=0))
|
||||
prev_fwd = None
|
||||
results = []
|
||||
last_fwd = None
|
||||
for t in range(1, turns + 1):
|
||||
cur = _client_convo(t)
|
||||
frozen = tracker.get_frozen_message_count()
|
||||
fwd = _freeze(cur, frozen)
|
||||
if use_overlay:
|
||||
fwd = overlay_cached_prefix(
|
||||
fwd,
|
||||
cur,
|
||||
tracker.get_last_original_messages(),
|
||||
tracker.get_last_forwarded_messages(),
|
||||
)
|
||||
exp = sum(_toklen(m) for m in prev_fwd) if prev_fwd else 0
|
||||
act = _cache_read(fwd, prev_fwd)
|
||||
results.append((exp, act))
|
||||
counts = [_toklen(m) for m in fwd]
|
||||
tracker.update_from_response(
|
||||
act, sum(counts) - act, fwd, message_token_counts=counts, original_messages=cur
|
||||
)
|
||||
prev_fwd = fwd
|
||||
last_fwd = fwd
|
||||
return results, last_fwd
|
||||
|
||||
|
||||
def test_moving_marker_busts_without_overlay():
|
||||
results, _ = _drive(use_overlay=False)
|
||||
assert any(exp > act for exp, act in results[1:]), "moving marker should bust the raw freeze"
|
||||
|
||||
|
||||
def test_moving_marker_no_bust_with_overlay():
|
||||
results, _ = _drive(use_overlay=True)
|
||||
for exp, act in results[1:]:
|
||||
assert act >= exp, f"cache bust under moved marker: expected {exp} read {act}"
|
||||
|
||||
|
||||
# ── fix-2: Headroom owns cache_control placement (realistic block content) ────
|
||||
from headroom.cache.prefix_tracker import ( # noqa: E402
|
||||
_strip_cache_control,
|
||||
normalize_message_cache_control,
|
||||
)
|
||||
|
||||
|
||||
def B(role, text, cc=False):
|
||||
"""Anthropic block-style message (cache_control lives on a content block)."""
|
||||
blk = {"type": "text", "text": text}
|
||||
if cc:
|
||||
blk["cache_control"] = {"type": "ephemeral"}
|
||||
return {"role": role, "content": [blk]}
|
||||
|
||||
|
||||
def _markers(messages):
|
||||
return sum(
|
||||
1
|
||||
for m in messages
|
||||
if isinstance(m.get("content"), list)
|
||||
for b in m["content"]
|
||||
if isinstance(b, dict) and "cache_control" in b
|
||||
)
|
||||
|
||||
|
||||
def test_normalize_strips_all_and_keeps_one_on_last():
|
||||
# 5 accumulated markers (the pile-up the overlay would produce).
|
||||
msgs = [
|
||||
B("user", "a", cc=True),
|
||||
B("assistant", "b", cc=True),
|
||||
B("user", "c", cc=True),
|
||||
B("user", "d", cc=True),
|
||||
B("user", "e", cc=True),
|
||||
]
|
||||
out = normalize_message_cache_control(msgs)
|
||||
assert _markers(out) == 1 # bounded — no >4 error
|
||||
assert "cache_control" in out[-1]["content"][-1] # on the last block
|
||||
assert _strip_cache_control(out) == _strip_cache_control(msgs) # content untouched
|
||||
|
||||
|
||||
def test_normalize_stays_bounded_across_many_turns():
|
||||
"""The accumulation that would 400 Anthropic is now capped at 1 every turn."""
|
||||
conv = []
|
||||
forwarded = []
|
||||
for t in range(1, 12):
|
||||
conv = conv + [B("user", f"turn-{t}", cc=True)] # client marks the newest
|
||||
forwarded = normalize_message_cache_control(conv)
|
||||
assert _markers(forwarded) <= 4 # never exceeds Anthropic's limit
|
||||
assert _markers(forwarded) == 1 # exactly one, on the last message
|
||||
|
||||
|
||||
def test_normalize_is_noop_when_no_block_markers():
|
||||
plain = [B("user", "a"), B("assistant", "b")] # no cache_control
|
||||
out = normalize_message_cache_control(plain)
|
||||
# places exactly one breakpoint (so the prefix gets cached), content stable
|
||||
assert _markers(out) == 1
|
||||
assert _strip_cache_control(out) == _strip_cache_control(plain)
|
||||
Loading…
Add table
Add a link
Reference in a new issue