Merge branch 'feat/token-mode-prefix-cache' into main

Author: nangsontay

CHANGELOG.md

@@ -14,6 +14,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 - **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/session:** `compute_session_id`'s fallback (used when no `x-headroom-session-id` header is sent) now folds in a conversation discriminator — the caller-supplied hint (Anthropic `metadata.user_id` / OpenAI `user`) or, failing that, the first user message's text — instead of hashing only `model + system prompt`. Two concurrent or successive conversations sharing a model and system prompt (two Claude Code tabs on one repo, a new session started after ending one) previously collapsed onto the same `PrefixCacheTracker`, so interleaved turns overwrote each other's recorded original/forwarded messages and busted the prompt cache for both ([#1808](https://github.com/headroomlabs-ai/headroom/issues/1808)).
+- **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)).

docs/content/docs/configuration.mdx

@@ -317,7 +317,7 @@ For provider-only proxying, prefer `HEADROOM_HTTP_PROXY` over process-wide varia
 
 ### Session Beta Header Tracking
 
-When running as a proxy, Headroom maintains a per-session union of `anthropic-beta` (and `OpenAI-Beta`) tokens via `SessionBetaTracker`. The session key is derived from the `x-headroom-session-id` header if present, otherwise from `md5(model + system_prompt[:500])[:16]` — stable across turns of the same conversation.
+When running as a proxy, Headroom maintains a per-session union of `anthropic-beta` (and `OpenAI-Beta`) tokens via `SessionBetaTracker`. The session key is derived from the `x-headroom-session-id` header if present, otherwise from `md5(model + system_prompt + conversation_discriminator)[:16]`, where the discriminator is a client-supplied hint (Anthropic `metadata.user_id` / OpenAI `user`) or, failing that, the first user message's text. This keeps the key stable across turns of one conversation while staying distinct across concurrent conversations that share a model and system prompt.
 
 **Why:** clients such as Claude Code and Codex CLI may drop a beta token between consecutive turns. Because `anthropic-beta` is part of the request bytes that determine the upstream prefix-cache key, a dropped token would bust the cache mid-conversation. The tracker re-injects any token seen earlier in the session so the cache key stays stable.
 

headroom/cache/prefix_tracker.py

@@ -351,6 +351,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]]:
@@ -482,6 +509,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)
 
@@ -764,6 +799,30 @@ def _estimate_message_tokens(messages: list[dict[str, Any]]) -> list[int]:
         return counts
 
 
+def _first_user_message_text(messages: list[dict[str, Any]]) -> str:
+    """Extract the first user message's text, capped for use as a hash input.
+
+    Stable within one append-only conversation (the client always resends the
+    first message unchanged) but differs across conversations in practice, so
+    it serves as a conversation discriminator when no explicit hint is given.
+    """
+    for msg in messages:
+        if msg.get("role") != "user":
+            continue
+        content = msg.get("content", "")
+        if isinstance(content, str):
+            return content[:500]
+        if isinstance(content, list):
+            parts = [
+                block.get("text", "")
+                for block in content
+                if isinstance(block, dict) and block.get("type") == "text"
+            ]
+            return "".join(parts)[:500]
+        return ""
+    return ""
+
+
 class SessionTrackerStore:
     """Manages PrefixCacheTracker instances across sessions.
 
@@ -800,12 +859,19 @@ def compute_session_id(
         request: Any,
         model: str,
         messages: list[dict[str, Any]],
+        *,
+        conversation_hint: str | None = None,
     ) -> str:
         """Compute a session ID from the request.
 
         Priority:
         1. x-headroom-session-id header (explicit)
-        2. Hash of (model + system prompt) — stable per conversation
+        2. Hash of (model + system prompt + conversation discriminator)
+
+        The discriminator is ``conversation_hint`` when the caller supplies one
+        (Anthropic ``metadata.user_id`` / OpenAI ``user``), else the first user
+        message's text. It keeps concurrent or successive conversations that
+        share a model+system prompt from colliding onto one tracker.
 
         The system prompt is harvested from ``role:"system"`` entries in
         ``messages``. Anthropic carries the system prompt as a top-level
@@ -834,7 +900,17 @@ def compute_session_id(
                             system_parts.append(block.get("text", ""))
 
         system_content = json.dumps(system_parts, ensure_ascii=False, separators=(",", ":"))
-        key = f"{model}:{system_content}"
+        # Conversation discriminator: two conversations sharing model+system
+        # (concurrent tabs, a new session started after ending one) must NOT
+        # collapse onto the same tracker/CCR-sticky state. Prefer a
+        # client-supplied hint (Anthropic metadata.user_id / OpenAI user);
+        # fall back to the first user message's text, which is stable across
+        # append-only turns of ONE conversation but differs across
+        # conversations in practice. Hash only — never log raw content.
+        discriminator = str(conversation_hint)[:200] if conversation_hint else ""
+        if not discriminator:
+            discriminator = _first_user_message_text(messages)
+        key = f"{model}:{system_content}:{discriminator}"
         return hashlib.md5(key.encode()).hexdigest()[:16]  # nosec B324
 
     def _maybe_cleanup(self) -> None:

headroom/ccr/context_tracker.py

@@ -56,6 +56,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
@@ -135,6 +136,7 @@ def track_compression(
         compressed_count: int,
         *,
         workspace_key: str,
+        session_id: str | None = None,
         query_context: str = "",
         sample_content: str = "",
     ) -> None:
@@ -167,6 +169,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
@@ -193,6 +196,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.
 
@@ -240,6 +244,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()