mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): enable tool search by default and repair poisoned transcripts
Anthropic validates every tool_reference in the message history against the CURRENT request's tools array. Claude Code replays one transcript across requests carrying very different tools arrays -- the main loop sends hundreds of tools, the Stop hook evaluator and other side-requests send a handful -- so once server-side tool search puts a server_tool_use / tool_search_tool_result pair in history, every later small-tools request 400s with "Tool reference 'X' not found in available tools". Add strip_unsupported_tool_search_blocks(): drop tool-search result blocks whose references this request's tools array cannot resolve, plus their paired server_tool_use. Other server tools (web search, code execution) are left alone. Copy-on-write, so the main loop gets the original list back by identity and its prefix cache is untouched. The repair runs unconditionally after injection, so transcripts poisoned before the upgrade self-heal. With it in place, HEADROOM_TOOL_SEARCH now defaults to on -- matching what the default 'coding' savings profile already seeded at proxy startup. Closes #2805
This commit is contained in:
Tejas Chopra
parent
303e0522c4
commit
a6ebeb6db8
4 changed files with 269 additions and 3 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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+).
|
||||
#
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue