diff --git a/docs/content/docs/proxy.mdx b/docs/content/docs/proxy.mdx index 984285146..0351e3181 100644 --- a/docs/content/docs/proxy.mdx +++ b/docs/content/docs/proxy.mdx @@ -270,7 +270,7 @@ Defers large tool schemas so they don't sit in every request. See [MCP](/docs/mc | Env | Scope | Effect | |---|---|---| -| `HEADROOM_TOOL_SEARCH` | proxy (server-side) | Defer MCP/system tool schemas behind a `search_tools` tool. The `coding` profile enables it. | +| `HEADROOM_TOOL_SEARCH` | proxy (server-side) | Defer MCP/system tool schemas behind a `search_tools` tool. **On by default** for Anthropic requests carrying enough tools to be worth it; set `HEADROOM_TOOL_SEARCH=0` to opt out. | | `ENABLE_TOOL_SEARCH` | client (Claude Code) | Keep Claude Code's own deferred tool-loading active behind a custom base URL ([issue #746](https://github.com/headroomlabs-ai/headroom/issues/746)). Set automatically by `headroom wrap`. | ### Cost-aware model routing diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index a62fc1d80..eaf399937 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2401,7 +2401,10 @@ class AnthropicHandlerMixin: optimized_tokens = tokenizer.count_messages(body["messages"]) tokens_saved = max(0, original_tokens - optimized_tokens) - # Server-side Tool Search (opt-in HEADROOM_TOOL_SEARCH): defer the + # Server-side Tool Search (on by default; HEADROOM_TOOL_SEARCH=0 opts + # out — the `coding` savings profile already seeded it on via + # seed_proxy_env_defaults, so default-on here just makes the same + # posture hold for entry points that never seeded): defer the # non-core tool schemas behind a tool_search tool so Anthropic excludes # them from the context window — they stop counting as input tokens until # the model searches for one — while every tool stays callable. @@ -2420,7 +2423,7 @@ class AnthropicHandlerMixin: if ( provider_name == "anthropic" and getattr(self, "anthropic_backend", None) is None - and os.environ.get("HEADROOM_TOOL_SEARCH", "").strip().lower() + and os.environ.get("HEADROOM_TOOL_SEARCH", "1").strip().lower() in ("1", "true", "yes", "on", "auto") ): from headroom.proxy.helpers import inject_tool_search_deferral @@ -2446,6 +2449,35 @@ class AnthropicHandlerMixin: f"{_ts_saved_tokens}tok" ) + # Tool-search history repair (#2805). Once deferral is on, the client + # stores Anthropic's server_tool_use / tool_search_tool_result blocks in + # its transcript forever, and upstream validates every tool_reference in + # that history against THIS request's tools array. Claude Code replays + # the same transcript on side-requests carrying a different, smaller + # tools array (the prompt-type Stop hook evaluator, /compact), which the + # proxy cannot predict — so upstream 400s with "Tool reference 'X' not + # found in available tools". Drop the blocks such a request cannot + # support. Runs AFTER the injection above so the tool we just added + # counts as present: on the main loop nothing is stripped and the prefix + # is untouched. Unconditional (not gated on the flag) so transcripts + # poisoned before the flag was turned off still recover. + from headroom.proxy.helpers import strip_unsupported_tool_search_blocks + + _ts_repaired, _ts_stripped = strip_unsupported_tool_search_blocks( + body.get("messages"), body.get("tools") + ) + if _ts_stripped: + body["messages"] = _ts_repaired + optimized_messages = _ts_repaired + body_mutation_tracker.mark_mutated("tool_search_history_repair") + transforms_applied.append(f"router:tool_search_repair:{_ts_stripped}blocks") + logger.info( + "[%s] Tool search: dropped %d unsupportable history block(s) " + "(tools array cannot resolve their tool_reference entries)", + request_id, + _ts_stripped, + ) + # Turn hooks (opt-in extensions): a registered hook may inspect or # rewrite the outbound tools/messages before we send upstream — the # extensible counterpart to the built-in deferral above. A single diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index efa1ccc7b..adcb846e9 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -2364,6 +2364,117 @@ def inject_tool_search_deferral( return out +# --------------------------------------------------------------------------- +# Tool-search history repair (issue #2805). +# +# Once the deferral above is active, Anthropic answers with ``server_tool_use`` +# (the search) + ``tool_search_tool_result`` (a list of ``tool_reference`` +# entries) blocks, and the client writes them into its transcript permanently. +# Anthropic validates every ``tool_reference`` in the history against the +# request's ``tools`` array and 400s with +# ``Tool reference 'X' not found in available tools`` when one is missing. +# +# That is fine for a client's main loop — the proxy re-injects the same tools +# array every turn — but Claude Code also replays the SAME transcript on +# side-requests that carry a different, smaller tools array (the prompt-type +# Stop hook evaluator, /compact, …). The proxy cannot predict those tool sets, +# so instead we repair the history: when the outbound request cannot support +# the tool-search blocks, drop them. Deterministic (same request → same output, +# so the prefix still caches), stateless (no session bookkeeping), and +# self-healing for transcripts already poisoned before the fix. +# --------------------------------------------------------------------------- + +_TOOL_SEARCH_RESULT_TYPE = "tool_search_tool_result" + + +def _tool_search_reference_names(content: Any) -> list[str]: + """Return the ``tool_reference`` names carried by a tool-search result block. + + Server-side results nest them (``content.tool_references``); a client-side + tool-search implementation returns the bare list. Accept both. + """ + entries = content.get("tool_references") if isinstance(content, dict) else content + if not isinstance(entries, list): + return [] + names = [] + for entry in entries: + if isinstance(entry, dict) and entry.get("type") == "tool_reference": + # Server-side blocks use ``tool_name``; be liberal about ``name``. + name = entry.get("tool_name") or entry.get("name") + if name: + names.append(str(name)) + return names + + +def strip_unsupported_tool_search_blocks(messages: Any, tools: Any) -> tuple[Any, int]: + """Drop tool-search blocks this request's ``tools`` array cannot support. + + A block pair is unsupportable when the request carries no ``tool_search_tool_*`` + tool, or when a ``tool_reference`` names a tool absent from ``tools`` — the two + shapes Anthropic rejects. Both the ``tool_search_tool_result`` and its paired + ``server_tool_use`` are removed (an orphan of either 400s on its own), and a + message left with no content blocks is dropped rather than sent empty. + + Returns ``(messages, blocks_removed)``, and the ORIGINAL ``messages`` object + when nothing was removed — callers rely on identity to skip the write-back. + """ + if not isinstance(messages, list): + return messages, 0 + + tool_list = tools if isinstance(tools, list) else [] + available = {str(t["name"]) for t in tool_list if isinstance(t, dict) and t.get("name")} + has_search_tool = any( + isinstance(t, dict) and str(t.get("type", "")).startswith(_TOOL_SEARCH_TOOL_TYPE_PREFIX) + for t in tool_list + ) + + out: list[Any] = [] + removed = 0 + changed = False + for message in messages: + content = message.get("content") if isinstance(message, dict) else None + if not isinstance(content, list): + out.append(message) + continue + + drop_indexes: set[int] = set() + orphaned_ids: set[str] = set() + for index, block in enumerate(content): + if not isinstance(block, dict) or block.get("type") != _TOOL_SEARCH_RESULT_TYPE: + continue + names = _tool_search_reference_names(block.get("content")) + if has_search_tool and all(name in available for name in names): + continue + drop_indexes.add(index) + use_id = block.get("tool_use_id") + if use_id: + orphaned_ids.add(str(use_id)) + # The search call itself precedes its result, so pair it up in a second + # pass. Only tool-search server calls are eligible — web_search and code + # execution use the same block type and must survive untouched. + for index, block in enumerate(content): + if not isinstance(block, dict) or block.get("type") != "server_tool_use": + continue + is_search_call = str(block.get("name", "")).startswith(_TOOL_SEARCH_TOOL_TYPE_PREFIX) + if str(block.get("id", "")) in orphaned_ids or (is_search_call and not has_search_tool): + drop_indexes.add(index) + + if not drop_indexes: + out.append(message) + continue + + changed = True + removed += len(drop_indexes) + kept = [block for index, block in enumerate(content) if index not in drop_indexes] + if not kept: + continue # the whole turn was tool-search bookkeeping + repaired = dict(message) + repaired["content"] = kept + out.append(repaired) + + return (out, removed) if changed else (messages, 0) + + # --------------------------------------------------------------------------- # Server-side Tool Search injection — OpenAI Responses API (gpt-5.4+). # diff --git a/tests/test_issue_746_tool_search.py b/tests/test_issue_746_tool_search.py index 337447dfa..74c0d4b9a 100644 --- a/tests/test_issue_746_tool_search.py +++ b/tests/test_issue_746_tool_search.py @@ -299,3 +299,126 @@ def test_resident_real_tool_survives_pascal_case_surface() -> None: # own; Anthropic 400s when every real tool is deferred. out = inject_tool_search_deferral(_claude_code_tools()) assert any(not t.get("type") and not t.get("defer_loading") for t in out) + + +# --------------------------------------------------------------------------- +# Tool-search history repair (#2805) +# +# Anthropic validates every tool_reference in the transcript against the +# request's tools array. Claude Code replays one transcript across requests +# with DIFFERENT tools arrays (main loop vs prompt-type Stop hook evaluator), +# so the side-request 400s with "Tool reference 'X' not found in available +# tools". The repair drops blocks a request cannot support. +# --------------------------------------------------------------------------- + +from headroom.proxy.helpers import ( # noqa: E402 + strip_unsupported_tool_search_blocks, +) + +_SEARCH_TOOL = {"type": _TOOL_SEARCH_DEFAULT_TYPE, "name": _TOOL_SEARCH_DEFAULT_NAME} + + +def _poisoned_transcript() -> list[dict]: + """A transcript as Claude Code stores it after one server-side tool search.""" + return [ + {"role": "user", "content": [{"type": "text", "text": "ask the user"}]}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Searching for a tool."}, + { + "type": "server_tool_use", + "id": "srvtoolu_01ABC", + "name": _TOOL_SEARCH_DEFAULT_NAME, + "input": {"pattern": "question|ask"}, + }, + { + "type": "tool_search_tool_result", + "tool_use_id": "srvtoolu_01ABC", + "content": { + "type": "tool_search_tool_search_result", + "tool_references": [ + {"type": "tool_reference", "tool_name": "AskUserQuestion"} + ], + }, + }, + {"type": "text", "text": "Found it."}, + ], + }, + ] + + +def test_repair_drops_blocks_the_hook_evaluator_cannot_resolve() -> None: + # The Stop hook evaluator replays the transcript with a small tools array + # that has neither the search tool nor AskUserQuestion -> upstream 400. + messages, removed = strip_unsupported_tool_search_blocks( + _poisoned_transcript(), [{"name": "Read", "input_schema": {}}] + ) + assert removed == 2 # server_tool_use + tool_search_tool_result + kinds = [b["type"] for b in messages[1]["content"]] + assert kinds == ["text", "text"] # surrounding assistant text survives + assert messages[0]["content"][0]["text"] == "ask the user" + + +def test_repair_is_noop_on_the_main_loop() -> None: + # Same transcript, but the request carries the injected search tool AND the + # referenced tool: nothing to repair, and the object is returned by identity + # so the outbound prefix (and its cache) is untouched. + transcript = _poisoned_transcript() + messages, removed = strip_unsupported_tool_search_blocks( + transcript, + [_SEARCH_TOOL, {"name": "AskUserQuestion", "input_schema": {}, "defer_loading": True}], + ) + assert removed == 0 + assert messages is transcript + + +def test_repair_drops_a_turn_left_with_no_blocks() -> None: + # An assistant turn that was ONLY the search round-trip must be removed, not + # forwarded with an empty content array (which Anthropic also rejects). + transcript = _poisoned_transcript() + transcript[1]["content"] = transcript[1]["content"][1:3] + messages, removed = strip_unsupported_tool_search_blocks(transcript, []) + assert removed == 2 + assert len(messages) == 1 + assert messages[0]["role"] == "user" + + +def test_repair_leaves_other_server_tools_alone() -> None: + # web_search / code execution use the same block type and stay untouched. + transcript = [ + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_web", + "name": "web_search", + "input": {"query": "x"}, + }, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_web", "content": []}, + ], + } + ] + messages, removed = strip_unsupported_tool_search_blocks(transcript, []) + assert removed == 0 + assert messages is transcript + + +def test_repair_is_idempotent() -> None: + # Deterministic: repairing an already-repaired transcript is a no-op, so a + # session's forwarded prefix stays byte-stable turn over turn. + once, _ = strip_unsupported_tool_search_blocks(_poisoned_transcript(), []) + twice, removed = strip_unsupported_tool_search_blocks(once, []) + assert removed == 0 + assert twice is once + + +def test_repair_strips_search_history_when_only_the_tool_is_missing() -> None: + # References all resolve, but the request has no tool_search tool at all + # (e.g. deferral skipped below _TOOL_SEARCH_MIN_TOOLS) -- history still + # cannot be supported, so it goes. + _, removed = strip_unsupported_tool_search_blocks( + _poisoned_transcript(), [{"name": "AskUserQuestion", "input_schema": {}}] + ) + assert removed == 2