peterlodri-sec
diff --git a/‎headroom/cache/prefix_tracker.py‎
Lines changed: 74 additions & 3 deletions b/‎headroom/cache/prefix_tracker.py‎
Lines changed: 74 additions & 3 deletions
diff --git a/‎headroom/proxy/handlers/anthropic.py‎
Lines changed: 14 additions & 1 deletion b/‎headroom/proxy/handlers/anthropic.py‎
Lines changed: 14 additions & 1 deletion
diff --git a/‎tests/test_cache_control_move_bust.py‎
Lines changed: 190 additions & 0 deletions b/‎tests/test_cache_control_move_bust.py‎
Lines changed: 190 additions & 0 deletions
@@ -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 _DeferredCompressionResult:
             # 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 _DeferredCompressionResult:
                 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):
 
@@ -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)