Merge pull request #26 from nangsontay/fix/ccr-transcript-tool-reinjection

fix(proxy/ccr): self-heal dangling headroom_retrieve reference after …

Author: nangsontay

docs/content/docs/troubleshooting.mdx

@@ -213,6 +213,33 @@ See [issue #746](https://github.com/headroomlabs-ai/headroom/issues/746) for the
   deployment. See [issue #2028](https://github.com/headroomlabs-ai/headroom/issues/2028).
 </Callout>
 
+## `Tool reference 'headroom_retrieve' not found` after a `/model` switch
+
+**Symptom**: Through the proxy, switching models mid-session (`/model` in Claude Code) —
+or restarting the proxy — makes every subsequent turn fail with:
+
+```
+400 Tool reference 'headroom_retrieve' not found in available tools
+```
+
+**Cause**: Once a session has done CCR compression, the proxy registers the
+`headroom_retrieve` retrieval tool and the client transcript starts carrying a
+`tool_reference` block naming it. Anthropic validates every `tool_reference` in the
+messages against the request's `tools` array. The proxy's sticky guarantee that keeps
+re-injecting the tool was keyed by a **model-scoped, in-memory** session id, so a
+`/model` switch (or a proxy restart) rotated/lost that key — the tool stopped being
+injected while the transcript kept re-sending the reference, leaving it dangling.
+
+**Fix**: Upgrade to a build that self-heals this. The proxy now scans the
+about-to-forward transcript for a dangling `headroom_retrieve` `tool_reference`/`tool_use`
+and re-injects the tool independent of tracker state, so `/model` switches and restarts
+recover automatically (logged as `decision=inject_transcript_recovery`).
+
+**Workarounds on older versions**:
+
+- `/model` back to the model you started the session on, or
+- `/clear` to start a fresh session (drops the dangling reference).
+
 ## Remote Control unavailable through custom ANTHROPIC_BASE_URL
 
 **Symptom**: When Claude Code runs with `ANTHROPIC_BASE_URL` set to a custom host (for example, Headroom), the Remote Control menu is absent.

headroom/proxy/ccr_marker_policy.py

@@ -35,11 +35,101 @@ def should_inject_ccr_tool(
     configured_inject_tool: bool,
     frozen_message_count: int,
     has_compressed_content: bool,
+    transcript_requires_tool: bool = False,
 ) -> tuple[bool, bool]:
-    """Decide whether the CCR retrieval tool must be injected this turn."""
+    """Decide whether the CCR retrieval tool must be injected this turn.
+
+    ``transcript_requires_tool`` forces injection when the about-to-forward
+    transcript already names ``headroom_retrieve`` but tracker state was lost
+    (a ``/model`` switch or proxy restart) — the dangling reference would
+    otherwise 400. It does NOT set ``is_marker_override``: that flag stays
+    specific to fresh #1006 markers so the caller can log each cause distinctly.
+    """
 
     inject_tool = configured_inject_tool
     if inject_tool and frozen_message_count > 0:
         inject_tool = False
     is_marker_override = not inject_tool and has_compressed_content
-    return (inject_tool or is_marker_override), is_marker_override
+    should_inject = inject_tool or is_marker_override or transcript_requires_tool
+    return should_inject, is_marker_override
+
+
+def transcript_references_ccr_tool(
+    messages: list[dict[str, Any]] | None,
+    *,
+    tool_name: str | None = None,
+    provider: Literal["anthropic", "openai", "google"] = "anthropic",
+) -> bool:
+    """Whether the about-to-forward transcript already names the CCR retrieve tool.
+
+    Once an Anthropic ``tool_reference``/``tool_use`` block, or an assistant
+    ``tool_calls`` entry (OpenAI chat), names ``headroom_retrieve``, the request's
+    ``tools`` array MUST still carry that tool or Anthropic 400s ("Tool reference
+    'headroom_retrieve' not found in available tools"). The sticky guarantee is
+    model-scoped and in-memory, so a ``/model`` switch or proxy restart loses it
+    while the client transcript keeps re-sending the reference. This scan lets
+    injection self-heal from the transcript, per request, independent of tracker
+    state.
+
+    ``provider="google"`` is accepted for signature parity with the sibling
+    policy fns but currently falls through the Anthropic matcher — Gemini's
+    ``functionCall`` parts are not matched (no google CCR handler calls this yet).
+
+    Only the exact bare ``headroom_retrieve`` name matches — a client-owned
+    ``mcp__headroom__headroom_retrieve`` (registered via MCP, lifecycle owned by
+    the client) must not trigger proxy injection. Tolerates string content and
+    malformed blocks.
+    """
+    if not messages:
+        return False
+    if tool_name is None:
+        from headroom.ccr.tool_injection import CCR_TOOL_NAME
+
+        tool_name = CCR_TOOL_NAME
+
+    for msg in messages:
+        if not isinstance(msg, dict):
+            continue
+        if provider == "openai":
+            if _openai_message_references_tool(msg, tool_name):
+                return True
+        elif _anthropic_content_references_tool(msg.get("content"), tool_name):
+            return True
+    return False
+
+
+def _anthropic_content_references_tool(content: Any, tool_name: str) -> bool:
+    """Match a bare ``tool_name`` in tool_reference/tool_use blocks (one level deep)."""
+    if not isinstance(content, list):
+        return False
+    for block in content:
+        if not isinstance(block, dict):
+            continue
+        if block.get("type") in ("tool_reference", "tool_use") and block.get("name") == tool_name:
+            return True
+        # tool_search_tool_result / tool_result nest blocks one level down.
+        nested = block.get("content")
+        if isinstance(nested, list):
+            for inner in nested:
+                if (
+                    isinstance(inner, dict)
+                    and inner.get("type") in ("tool_reference", "tool_use")
+                    and inner.get("name") == tool_name
+                ):
+                    return True
+    return False
+
+
+def _openai_message_references_tool(msg: dict[str, Any], tool_name: str) -> bool:
+    """Match a bare ``tool_name`` in an assistant message's ``tool_calls`` (chat shape)."""
+    tool_calls = msg.get("tool_calls")
+    if not isinstance(tool_calls, list):
+        return False
+    for call in tool_calls:
+        if not isinstance(call, dict):
+            continue
+        fn = call.get("function")
+        name = fn.get("name") if isinstance(fn, dict) else call.get("name")
+        if name == tool_name:
+            return True
+    return False

headroom/proxy/handlers/anthropic.py

@@ -1947,6 +1947,7 @@ class _DeferredCompressionResult:
                 from headroom.proxy.helpers import (
                     has_new_ccr_markers,
                     should_inject_ccr_tool,
+                    transcript_references_ccr_tool,
                 )
 
                 # #1850: only markers NEW this turn justify overriding the
@@ -1961,10 +1962,22 @@ class _DeferredCompressionResult:
                     provider="anthropic",
                 )
 
+                # Self-heal a dangling headroom_retrieve reference: a /model
+                # switch or proxy restart rotates the model-scoped sticky key,
+                # so the tracker stops injecting while the client transcript
+                # still carries a tool_reference/tool_use naming the tool.
+                # Anthropic 400s every such turn. Scanning the about-to-forward
+                # transcript recovers injection independent of tracker state.
+                transcript_requires_tool = transcript_references_ccr_tool(
+                    optimized_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=has_new_compressed_content,
+                    transcript_requires_tool=transcript_requires_tool,
                 )
                 if should_inject:
                     if is_marker_override:
@@ -1974,6 +1987,13 @@ class _DeferredCompressionResult:
                             f"(frozen_message_count={frozen_message_count}); injecting to "
                             "prevent unredeemable markers (#1006)"
                         )
+                    elif transcript_requires_tool:
+                        logger.info(
+                            f"[{request_id}] CCR: recovering headroom_retrieve from "
+                            "transcript — a dangling tool_reference exists but sticky "
+                            "tracker state was lost (/model switch or restart); "
+                            "re-injecting to avoid a 400"
+                        )
                     from headroom.proxy.helpers import apply_session_sticky_ccr_tool
 
                     tools, ccr_tool_injected = apply_session_sticky_ccr_tool(
@@ -1982,6 +2002,7 @@ class _DeferredCompressionResult:
                         request_id=request_id,
                         existing_tools=tools,
                         has_compressed_content_this_turn=has_new_compressed_content,
+                        transcript_requires_tool=transcript_requires_tool,
                     )
                     if ccr_tool_injected:
                         logger.debug(

headroom/proxy/handlers/openai.py

@@ -3230,6 +3230,7 @@ async def handle_openai_chat(
                 from headroom.proxy.helpers import (
                     apply_session_sticky_ccr_tool,
                     has_new_ccr_markers,
+                    transcript_references_ccr_tool,
                 )
 
                 # #1850: markers replayed from overlay_cached_prefix are
@@ -3242,12 +3243,23 @@ async def handle_openai_chat(
                     previous_forwarded_messages=openai_prefix_tracker.get_last_forwarded_messages(),
                     provider="openai",
                 )
+
+                # Self-heal a dangling headroom_retrieve reference after tracker
+                # state loss (proxy restart): the chat history still carries a
+                # prior headroom_retrieve tool call in assistant.tool_calls.
+                # OpenAI does not hard-400 on this the way Anthropic does, but
+                # re-injecting restores marker redeemability regardless.
+                transcript_requires_tool = transcript_references_ccr_tool(
+                    optimized_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=has_new_compressed_content,
+                    transcript_requires_tool=transcript_requires_tool,
                 )
                 if ccr_tool_injected:
                     logger.debug(

headroom/proxy/helpers.py

@@ -77,6 +77,9 @@
 from headroom.proxy.ccr_marker_policy import (
     should_inject_ccr_tool as _should_inject_ccr_tool,
 )
+from headroom.proxy.ccr_marker_policy import (
+    transcript_references_ccr_tool as transcript_references_ccr_tool,  # noqa: F401 - compatibility export
+)
 from headroom.proxy.ccr_session_tracker import SessionCcrTracker as _SessionCcrTracker
 from headroom.proxy.ccr_session_tracker import (
     SessionExpansionDedupTracker as _SessionExpansionDedupTracker,
@@ -2324,6 +2327,7 @@ def should_inject_ccr_tool(
     configured_inject_tool: bool,
     frozen_message_count: int,
     has_compressed_content: bool,
+    transcript_requires_tool: bool = False,
 ) -> tuple[bool, bool]:
     """Decide whether the ``headroom_retrieve`` tool must be injected this turn.
 
@@ -2337,6 +2341,11 @@ def should_inject_ccr_tool(
     that case we override the deferral and inject anyway (one cache miss is
     cheaper than dropped content).
 
+    ``transcript_requires_tool`` similarly overrides the deferral when the
+    transcript still carries a dangling ``headroom_retrieve`` reference but the
+    sticky tracker state was lost (a ``/model`` switch or proxy restart) — the
+    reference would otherwise 400.
+
     Returns ``(should_inject, is_marker_override)``. ``is_marker_override`` is
     True only when injection happens *because* of new markers despite a deferral,
     so the caller can log the override distinctly.
@@ -2345,6 +2354,7 @@ def should_inject_ccr_tool(
         configured_inject_tool=configured_inject_tool,
         frozen_message_count=frozen_message_count,
         has_compressed_content=has_compressed_content,
+        transcript_requires_tool=transcript_requires_tool,
     )
 
 
@@ -2355,6 +2365,7 @@ def apply_session_sticky_ccr_tool(
     request_id: str | None,
     existing_tools: list[dict[str, Any]] | None,
     has_compressed_content_this_turn: bool,
+    transcript_requires_tool: bool = False,
 ) -> tuple[list[dict[str, Any]], bool]:
     """Apply sticky-on CCR retrieval-tool injection per :class:`SessionCcrTracker`.
 
@@ -2365,13 +2376,16 @@ def apply_session_sticky_ccr_tool(
     Logic:
 
       * If ``session_id`` is None: tracker is bypassed and the per-turn
-        ``has_compressed_content_this_turn`` flag drives the decision
-        verbatim (matching legacy behaviour for WS / pre-session paths).
+        ``has_compressed_content_this_turn`` flag (or ``transcript_requires_tool``)
+        drives the decision verbatim (matching legacy behaviour for WS paths).
       * If the session has previously done CCR (``has_done_ccr``):
         ALWAYS inject the recorded golden bytes — even if this turn has
         no fresh compression. That is the load-bearing PR-B7 fix.
-      * Otherwise, inject only when this turn produced compressed content.
-        The first injection records the golden bytes for future turns.
+      * Otherwise, inject when this turn produced compressed content OR
+        ``transcript_requires_tool`` is set (the transcript still names the tool
+        but tracker state was lost — a ``/model`` switch or restart). The first
+        injection records the golden bytes so subsequent turns resume the normal
+        sticky-replay path.
 
     Tools whose name already equals ``CCR_TOOL_NAME`` (e.g. the client
     pre-registered it via MCP) are not re-appended; the client's bytes
@@ -2405,7 +2419,7 @@ def apply_session_sticky_ccr_tool(
 
     # No session_id (e.g. WS path): per-turn decision drives directly.
     if not session_id:
-        if not has_compressed_content_this_turn:
+        if not has_compressed_content_this_turn and not transcript_requires_tool:
             log_tool_injection_decision(
                 provider=provider,
                 session_id=None,
@@ -2419,7 +2433,11 @@ def apply_session_sticky_ccr_tool(
         log_tool_injection_decision(
             provider=provider,
             session_id=None,
-            decision="inject_first_time",
+            decision=(
+                "inject_transcript_recovery"
+                if transcript_requires_tool and not has_compressed_content_this_turn
+                else "inject_first_time"
+            ),
             tool_definition_bytes_count=len(replay.canonical_bytes),
             request_id=request_id,
         )
@@ -2469,8 +2487,10 @@ def apply_session_sticky_ccr_tool(
         )
         return tools_out, True
 
-    # Fresh session — only inject when this turn produced compressed content.
-    if not has_compressed_content_this_turn:
+    # Fresh session — inject when this turn produced compressed content, or when
+    # the transcript still references the tool but tracker state was lost
+    # (transcript recovery: a /model switch or proxy restart).
+    if not has_compressed_content_this_turn and not transcript_requires_tool:
         log_tool_injection_decision(
             provider=provider,
             session_id=session_id,
@@ -2486,7 +2506,11 @@ def apply_session_sticky_ccr_tool(
     log_tool_injection_decision(
         provider=provider,
         session_id=session_id,
-        decision="inject_first_time",
+        decision=(
+            "inject_transcript_recovery"
+            if transcript_requires_tool and not has_compressed_content_this_turn
+            else "inject_first_time"
+        ),
         tool_definition_bytes_count=len(replay.canonical_bytes),
         request_id=request_id,
     )

headroom/proxy/tool_injection_logging.py

@@ -8,6 +8,7 @@
 ToolInjectionDecision = Literal[
     "inject_first_time",
     "inject_sticky_replay",
+    "inject_transcript_recovery",
     "skip",
     "skip_disabled_via_env",
 ]