Merge pull request #3 from nangsontay/feat/token-mode-prefix-cache

fix(proxy): restore token-mode provider prefix-cache hit rate

Author: nangsontay

CHANGELOG.md

@@ -18,6 +18,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **cache/prefix-freeze:** resolve `PrefixCacheTracker`s per conversation lineage within a session id, so concurrent conversations sharing a fallback id no longer thrash one tracker's frozen-prefix state ([#2085](https://github.com/headroomlabs-ai/headroom/issues/2085)). Without an `x-headroom-session-id` header the fallback id hashes `model + system prompt` — identical across a Claude Code session and every one of its parallel subagents (and any set of sessions reusing one system prompt). On the shared tracker their interleaved histories cross-contaminate the freeze state: the forwarded prefix is byte-unstable on nearly every turn and the provider prompt cache is re-written instead of read — reported as ~4.4x cache-creation inflation and a 2.5–3x net cost increase under Claude Code. `SessionTrackerStore.resolve_tracker` now reuses the tracker whose previous request messages are a prefix of the incoming history (client histories are append-only, so a conversation's next request always extends its previous one), starts a fresh lineage when the history diverges or was rewritten (client-side compaction — that provider cache line is gone anyway), and caps lineages per session id (`PrefixFreezeConfig.max_lineages_per_session`, default 32; over-cap conversations share one overflow tracker instead of evicting established lineages, so a fan-out storm past the cap degrades only its own tail — and `0` disables lineage splitting). Matching compares the original client bytes under the same canonical cross-turn equivalence as the cache-stable delta path (`_canonicalize_for_prefix_compare`), so a moved cache breakpoint, string<->block content sugar, or per-turn transport annotations do not read as a rewrite. Separately, the fallback id now hashes only the LEADING run of `role:"system"` messages: agentic clients interleave `<system-reminder>` turns into the history as actual system-role messages (hook output, skills lists, truncation notices), and hashing those rotated the session id mid-conversation — orphaning the prefix tracker and every other session-sticky subsystem (beta headers, CCR/memory registries, the compression cache) each time a reminder landed. Both handler paths now derive the session id and the lineage from the same original client bytes, so a turn-dependent hook rewrite cannot rotate one without the other. The session id itself never changes: session-sticky state keyed on it (beta-header stickiness, CCR and memory-tool registries, the compression cache) is untouched, and a single-conversation session keeps its exact previous behavior (the first lineage lives under the bare id).
 - **proxy/bedrock:** wire `PrefixCacheTracker` updates into both Bedrock backend paths (`handle_anthropic_messages`'s non-streaming branch in `anthropic.py`, and `_stream_response_bedrock` in `streaming.py`). `update_from_response()` was previously only called from the direct-Anthropic-API branch; both Bedrock branches returned before ever reaching it, so the tracker's state stayed permanently empty for the life of a session on any `--backend bedrock` deployment: `extract_cache_stable_delta()` always saw no previous turn, and `--mode cache` fell back to full unmodified passthrough on every turn instead of freezing the already-cached prefix and compressing only the new suffix.
 - **install:** `install_supervisor`'s macOS branch did an unconditional `launchctl bootout` followed by a bare `bootstrap` with no retry, unlike `start_supervisor`, which already rides out the ~15s EIO (error 5) window launchd exhibits for several seconds after a bootout. This left `install apply`'s own reinstall path (and anything that re-applies a deployment, e.g. a future `headroom doctor --fix`) exposed to a race that previously required manual recovery (bootout + remove the plist + reapply). Extracted the retry loop already used by `start_supervisor` into a shared `_bootstrap_with_retry()` helper, now used by both call sites.
+- **proxy/anthropic:** CCR system-instruction injection (`ccr_inject_system_instructions`, default off) now appends to the top-level `body["system"]` field instead of prepending a `role:"system"` message to `messages` — the Anthropic Messages API rejects that message shape outright, so every marker-bearing turn with the flag enabled previously got a 400 from upstream. The injected text is now a static, hash-free block (`create_stable_system_instructions`) appended idempotently (existing `cache_control` breakpoints on prior system blocks are left untouched), so the system segment stays byte-identical turn over turn instead of busting the cache via the old per-turn hash list.
+- **proxy/ccr:** the messages-based system-instruction injection path (`CCRToolInjector.inject_into_system_message`, used by OpenAI-format clients when `ccr_inject_system_instructions` is enabled) now injects the same static, hash-free block via `create_stable_system_instructions` instead of `create_system_instructions`'s per-turn `Available hashes` list, so an OpenAI-format system prompt stays byte-identical turn over turn instead of busting the provider's prompt cache each time the enumerated hashes change.
+- **proxy/ccr:** CCR proactive-expansion and memory-context injection now skip appending to the live-zone tail message when that exact position was already forwarded (and byte-replayed by `overlay_cached_prefix`) last turn. Previously both injections ran unconditionally in token mode, so a same-messages re-request appended a second copy on top of the one `overlay_cached_prefix` had just replayed — busting the final segment and growing it unboundedly turn over turn. Proactive-expansion recommendations are also deduped per session (a `SessionExpansionDedupTracker`, mirroring the existing sticky-CCR-tool session store) so the same compressed-content hash is never injected twice into one conversation, and `ContextTracker.analyze_query`/`track_compression` now accept an optional `session_id` so a fresh session started after Claude Code's `/compact` no longer has pre-compact compressed content surfaced back into it ([#2186](https://github.com/headroomlabs-ai/headroom/issues/2186)).
+- **proxy/openai:** the OpenAI-format chat-completions path (`handle_openai_chat`, both the LiteLLM/any-llm backend branch and the direct-backend branch) and its streaming finalizer now pass the raw client messages as `original_messages` to `PrefixCacheTracker.update_from_response`. All three call sites previously omitted it, so the tracker fell back to recording the FORWARDED (compressed) bytes as "originals" — the next turn's append-only cache-safety check then compared raw client bytes against compressed ones, diverged immediately, and never replayed the cached prefix, silently busting the provider's prompt cache on every turn for OpenAI-format clients in token mode.
 - **proxy/savings:** `SavingsTracker.record_request()` only appended a history point when `tokens_saved > 0` (headroom's own lossy compression). In `--mode cache`, `tokens_saved` is near-always 0 by design, since the frozen prefix is byte-replayed rather than compressed to keep the provider's prompt cache warm. That silently dropped every history point on a cache-mode deployment even when `cache_read_tokens`/`cache_savings_usd` were large, making `headroom-monthly`-style tooling read as a total savings collapse. The guard now fires on `tokens_saved` OR `cache_read_tokens`, and the appended entry carries `cache_read_tokens`/`cache_savings_usd` so downstream consumers can show them; `_normalize_history_entry` defaults both fields to 0/0.0 for legacy entries that predate this change.
 - **litellm:** vendor-specific top-level fields on `/v1/chat/completions`, including vLLM's `chat_template_kwargs` for per-request Qwen3 thinking-mode toggles, now reach OpenAI-compatible backends through LiteLLM `extra_body` instead of being dropped by the standard-parameter allowlist ([#2128](https://github.com/headroomlabs-ai/headroom/issues/2128)).
 - **cache aligner:** hash the actual frozen Claude Code prefix instead of only system-message text, so `stable_prefix_hash` / `prefix_changed` now surface prompt-cache churn when a cached tool-result block changes without any system-prompt edit ([#2085](https://github.com/headroomlabs-ai/headroom/issues/2085)).

headroom/cache/prefix_tracker.py

@@ -359,6 +359,33 @@ def overlay_cached_prefix(
     return list(prev_fwd[:k]) + list(optimized_messages[k:])
 
 
+def is_append_only_extension(
+    current_original_messages: list[dict[str, Any]],
+    previous_original_messages: list[dict[str, Any]] | None,
+) -> bool:
+    """Return True iff ``previous_original_messages`` is a full canonical
+    prefix of ``current_original_messages``.
+
+    This is the same append-only condition ``overlay_cached_prefix`` requires
+    before it will replay ANY previously-forwarded content — a position is
+    only guaranteed to hold last turn's cached bytes when this holds for the
+    ENTIRE previous message list, not just a leading run of it. Used by
+    handler-level injection guards to detect "this exact tail position was
+    already forwarded last turn" without duplicating the canonicalization
+    logic.
+    """
+    if not previous_original_messages:
+        return False
+    n = len(previous_original_messages)
+    if len(current_original_messages) < n:
+        return False
+    return all(
+        _canonicalize_for_prefix_compare(current_original_messages[i])
+        == _canonicalize_for_prefix_compare(previous_original_messages[i])
+        for i in range(n)
+    )
+
+
 def normalize_message_cache_control(
     messages: list[dict[str, Any]],
 ) -> list[dict[str, Any]]:
@@ -490,6 +517,14 @@ def update_from_response(
         """
         self._last_activity = time.time()
         self._turn_number += 1
+        if original_messages is None:
+            logger.warning(
+                "PrefixCacheTracker[%s]: update_from_response called without "
+                "original_messages — falling back to forwarded messages as "
+                "originals, which busts the overlay's append-only cache-safety "
+                "check on the next turn.",
+                self.provider,
+            )
         self._last_original_messages = copy.deepcopy(original_messages or messages)
         self._last_forwarded_messages = copy.deepcopy(messages)
 

headroom/ccr/context_tracker.py

@@ -88,6 +88,7 @@ class CompressedContext:
     query_context: str  # The query/context when compression happened
     sample_content: str  # Preview of what was compressed (for relevance matching)
     workspace_key: str  # Stable per-project identity (see ProjectResolver in storage_router)
+    session_id: str | None = None  # Conversation identity (see analyze_query)
 
 
 @dataclass
@@ -167,6 +168,7 @@ def track_compression(
         compressed_count: int,
         *,
         workspace_key: str,
+        session_id: str | None = None,
         query_context: str = "",
         sample_content: str = "",
     ) -> None:
@@ -206,6 +208,7 @@ def track_compression(
             query_context=query_context,
             sample_content=sample_content[:2000],  # Limit sample size
             workspace_key=workspace_key,
+            session_id=session_id,
         )
 
         # Add or update context
@@ -232,6 +235,7 @@ def analyze_query(
         current_turn: int | None = None,
         *,
         workspace_key: str,
+        session_id: str | None = None,
     ) -> list[ExpansionRecommendation]:
         """Analyze a query to find relevant compressed contexts.
 
@@ -279,6 +283,18 @@ def analyze_query(
             if context.workspace_key != workspace_key:
                 continue
 
+            # Session filter (#2186): a NEW session (e.g. started after
+            # Claude Code's /compact) must not have old compressed content
+            # surfaced back into it — the whole point of compacting was to
+            # drop it. Only enforced when both sides have a concrete id;
+            # legacy/no-session callers keep workspace-only scoping.
+            if (
+                session_id is not None
+                and context.session_id is not None
+                and context.session_id != session_id
+            ):
+                continue
+
             # Check age
             age = now - context.timestamp
             if age > self.config.max_context_age_seconds:

headroom/ccr/tool_injection.py

@@ -144,6 +144,42 @@ def create_system_instructions(
 """
 
 
+def create_stable_system_instructions(
+    retrieval_endpoint: str = "/v1/retrieve",
+) -> str:
+    """Create session-stable CCR retrieval instructions (no hash list).
+
+    Same purpose as ``create_system_instructions`` but WITHOUT the per-turn
+    ``**Available hashes:**`` list. Once a session injects this text, it must
+    stay byte-identical every subsequent turn or the segment it lives in
+    busts the provider's prompt cache each time the enumerated hashes change.
+    The model can always find a compressed output's hash from the marker
+    already present in that tool result (``[N items compressed to M.
+    Retrieve more: hash=abc123]``), so the separate list was redundant, not
+    load-bearing.
+
+    Args:
+        retrieval_endpoint: The endpoint path for retrieval (unused in the
+            text today, kept for signature parity with
+            ``create_system_instructions``).
+
+    Returns:
+        Static instruction text to append to the system prompt.
+    """
+    return f"""
+## Compressed Context Available
+
+Some tool outputs have been compressed to reduce context size. If you need
+the full uncompressed data, you can retrieve it using the `{CCR_TOOL_NAME}` tool.
+
+**How to retrieve:**
+- Call `{CCR_TOOL_NAME}(hash="<hash>")` to get the full original content back
+
+Look for markers like `[N items compressed to M. Retrieve more: hash=abc123]`
+in tool results to find the hash for each compressed output.
+"""
+
+
 @dataclass
 class CCRToolInjector:
     """Manages CCR tool injection into LLM requests.
@@ -362,8 +398,12 @@ def inject_into_system_message(
         if not self.inject_system_instructions or not self.has_compressed_content:
             return messages
 
-        instructions = create_system_instructions(
-            self._detected_hashes,
+        # Use the hash-free stable variant. create_system_instructions embeds
+        # an "Available hashes" list that changes every turn as new content is
+        # compressed, which busts the provider's prompt cache on the system
+        # segment. The list is redundant — each compressed tool result carries
+        # its own retrievable hash in its marker — so drop it here.
+        instructions = create_stable_system_instructions(
             self.retrieval_endpoint,
         )
 

headroom/proxy/ccr_session_tracker.py

@@ -85,3 +85,56 @@ def reset(self) -> None:
 
         with self._lock:
             self._sessions.clear()
+
+
+class SessionExpansionDedupTracker:
+    """Bounded LRU tracker for per-session CCR proactive-expansion dedup.
+
+    Prevents the same compressed-content hash from being proactively
+    re-injected more than once per session (#2186): a same-messages
+    re-request or a continued conversation must not receive duplicate
+    expansion blocks for content the agent already saw this session.
+    Dedup is per-session, not global — a different conversation may
+    legitimately need the same expansion.
+    """
+
+    def __init__(self, max_sessions: int) -> None:
+        if max_sessions <= 0:
+            raise ValueError("max_sessions must be > 0")
+        self._max_sessions = max_sessions
+        self._lock = threading.RLock()
+        self._sessions: OrderedDict[str, set[str]] = OrderedDict()
+
+    def filter_new(self, session_id: str, hash_keys: list[str]) -> list[str]:
+        """Return the subset of hash_keys not yet injected for this session."""
+
+        if not session_id:
+            raise ValueError("session_id must be non-empty")
+        with self._lock:
+            seen = self._sessions.get(session_id)
+            if seen is None:
+                return list(hash_keys)
+            return [h for h in hash_keys if h not in seen]
+
+    def record_injected(self, session_id: str, hash_keys: list[str]) -> None:
+        """Mark hash_keys as injected for this session."""
+
+        if not session_id:
+            raise ValueError("session_id must be non-empty")
+        if not hash_keys:
+            return
+        with self._lock:
+            seen = self._sessions.get(session_id)
+            if seen is None:
+                seen = set()
+                self._sessions[session_id] = seen
+            seen.update(hash_keys)
+            self._sessions.move_to_end(session_id)
+            while len(self._sessions) > self._max_sessions:
+                self._sessions.popitem(last=False)
+
+    def reset(self) -> None:
+        """Clear all session state."""
+
+        with self._lock:
+            self._sessions.clear()