github-featured/chopratejas--headroom
2
0
Fork 0
Code Issues Pull Requests Actions 16 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

fix(proxy): enable tool search by default and repair poisoned transcripts (#2807)

Browse Source 
## Description

Server-side tool search poisons the Claude Code transcript: once the
proxy injects deferral and the model runs one search, Anthropic's
`server_tool_use` + `tool_search_tool_result` pair lives in the message
history forever. Upstream validates **every `tool_reference` in that
history against the *current* request's `tools` array** — and Claude
Code replays one transcript across requests with wildly different tools
arrays (main loop: hundreds of tools; prompt-type Stop hook evaluator,
`/compact`, other side-requests: a handful). Every one of those
side-requests 400s with `Tool reference 'X' not found in available
tools`.

This PR keeps tool search **on** — it's the whole point of the feature,
and the default `coding` savings profile already turned it on at proxy
startup — and instead repairs the transcript per request, statelessly.

The issue author's preferred fix (never inject for Claude Code clients)
would disable the feature for its main audience. A session-sticky
approach was also considered and rejected: it needs session state, it
can't re-add ~500 tool definitions to a 5-tool side-request without
erasing the savings, and it can't heal transcripts already poisoned
before the upgrade.

Closes #2805

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [x] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- **`headroom/proxy/helpers.py`** — new
`strip_unsupported_tool_search_blocks(messages, tools)`. Builds the set
of names this request can resolve, drops any `tool_search_tool_result`
whose `tool_reference` entries aren't all resolvable (or when no search
tool is present at all), and drops the paired `server_tool_use` by
`tool_use_id`. Other server tools (`web_search`, code execution) are
untouched. Turns left with zero content blocks are removed rather than
forwarded empty. Copy-on-write: returns the **original** `messages`
object by identity when nothing was removed.
- **`headroom/proxy/handlers/anthropic.py`** — runs the repair right
after the injection block, so the tool just injected counts as present
and the main loop is a no-op with a byte-identical prefix. Deliberately
**not** gated on `HEADROOM_TOOL_SEARCH`, so transcripts poisoned before
an upgrade (or before someone sets the flag to `0`) still recover. Logs
and tags `router:tool_search_repair:Nblocks` when it fires.
- **`headroom/proxy/handlers/anthropic.py`** — `HEADROOM_TOOL_SEARCH`
now defaults to `1`. This matches the posture
`seed_proxy_env_defaults()` already established for the default `coding`
profile; the flip only affects entry points that never seeded.
- **`docs/content/docs/proxy.mdx`** — documents on-by-default plus
`HEADROOM_TOOL_SEARCH=0` as the opt-out.
- **`tests/test_issue_746_tool_search.py`** — 6 tests covering the
repair.

### Answering the issue's open question

> we could not determine what enables it — `/proc/<pid>/environ` shows
no `HEADROOM_TOOL_SEARCH`

`seed_proxy_env_defaults()` calls
`os.environ.setdefault("HEADROOM_TOOL_SEARCH", "1")` at proxy startup
because the default savings profile is `coding`, which has
`tool_search=True` (`headroom/agent_savings.py`). In-process mutation of
`os.environ` never appears in the process's environ snapshot, which is
why the flag looked unset.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ python -m pytest tests/test_issue_746_tool_search.py -q
45 passed, 1 warning in 1.56s

$ python -m pytest tests/test_*anthropic*.py tests/test_*tool*.py -q
4 failed, 459 passed, 2 skipped, 7 warnings in 27.40s
# the 4 failures are in tests/test_bedrock_tool_result_cache_and_streaming_stats.py
# and reproduce identically on this branch's merge-base with the changes stashed:
#   4 failed, 9 passed, 5 warnings in 3.02s

$ ruff check headroom/proxy/helpers.py headroom/proxy/handlers/anthropic.py tests/test_issue_746_tool_search.py
All checks passed!

$ ruff format --check <same three files>
3 files already formatted

$ mypy --python-version 3.12 headroom/proxy/helpers.py headroom/proxy/handlers/anthropic.py
Success: no issues found in 2 source files
# --python-version 3.12 only to skip a pre-existing numpy-stub syntax error that
# the repo's python_version = "3.10" setting triggers on this machine.
```

New tests:

| Test | Asserts |
|---|---|
| `test_repair_drops_blocks_the_hook_evaluator_cannot_resolve` | small
tools array → both blocks dropped, surrounding assistant text survives |
| `test_repair_is_noop_on_the_main_loop` | search tool + referenced tool
present → `removed == 0` and `messages is transcript` (prefix cache
untouched) |
| `test_repair_drops_a_turn_left_with_no_blocks` | a turn that was
*only* the search round-trip is removed, not forwarded empty |
| `test_repair_leaves_other_server_tools_alone` | `web_search`
`server_tool_use` blocks survive |
| `test_repair_is_idempotent` | second pass over a repaired transcript
removes nothing |
| `test_repair_strips_search_history_when_only_the_tool_is_missing` |
references resolvable but no search tool in the array → still stripped |

## Real Behavior Proof

- **Environment:** macOS 25.4.0, Python 3.12 venv, live
`api.anthropic.com`, `claude-sonnet-4-6`, local proxy on
`127.0.0.1:8799` built from this branch.
- **Exact command / steps:** one request body — a poisoned transcript
(`server_tool_use` + `tool_search_tool_result` referencing
`AskUserQuestion`) with a **1-tool** `tools` array (`Read`), exactly the
shape a Claude Code side-request replays — sent twice: once straight to
`https://api.anthropic.com`, once to the proxy.

```text
$ python /tmp/hr-2805-repro.py https://api.anthropic.com
HTTP 400
{"type": "invalid_request_error", "message": "Tool reference 'AskUserQuestion' not found in available tools"}

$ python /tmp/hr-2805-repro.py http://127.0.0.1:8799
HTTP 200
content: [{"type": "text", "text": "OK"}]
```

- **Observed result:** the exact 400 from the issue reproduces against
upstream; the identical body through the proxy returns 200. The proxy's
savings event for that request records `before: 133, after: 32, saved:
101` tokens — the two dropped blocks. The one-tool array is below
`_TOOL_SEARCH_MIN_TOOLS = 12`, so no injection ran; the repair alone is
what made the request valid.
- **Not tested:** a full end-to-end Claude Code session with a real Stop
hook (the synthetic replay above is the same request shape the hook
evaluator produces); non-Anthropic providers, which don't have
server-side tool search.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md`

## Screenshots (if applicable)

N/A — proxy-side behavior, covered by the command output above.

## Additional Notes

- **Cache cost is zero on the hot path.** The repair only rewrites
requests whose transcripts reference tools they don't carry — request
families that were 400ing anyway. The main loop takes the identity path
and its prefix stays byte-identical.
- **Out of scope, spotted while here:** `run-all-plugins.sh` exports
`HEADROOM_TOOL_SEARCH_MIN_TOOLS=5`, but nothing in Python reads it —
`_TOOL_SEARCH_MIN_TOOLS` is a hardcoded `12`. Worth a follow-up.
This commit is contained in:
Tejas Chopra
2026-08-05 14:36:34 -07:00
committed by GitHub
parent 303e0522c4
commit 0237cbffbb
4 changed files with 269 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/content/docs/proxy.mdx
+1 -1
View File
@@ -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
headroom/proxy/handlers/anthropic.py
+34 -2
View File
@@ -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
headroom/proxy/helpers.py
+111
View File
@@ -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+).
#
tests/test_issue_746_tool_search.py
+123
View File
@@ -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