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

## 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
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 269 additions and 3 deletions

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