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/anthropic): run tool-search history repair after turn hooks

Browse Source 
## Description

`strip_unsupported_tool_search_blocks` (#2807) validates every replayed
`tool_reference` in the transcript against the request's `tools` array.
It ran *before* the turn-hooks block in `handlers/anthropic.py`, and a
registered turn hook may rewrite that array — the hook surface is
documented as "a registered hook may inspect or rewrite the outbound
tools/messages before we send upstream".

So a hook that drops a tool named by a replayed reference leaves the
repair having validated against a stale view, and upstream rejects the
request:

```text
400 Tool reference 'X' not found in available tools
```

The repair's correctness argument is that it validates against exactly
the `tools` array upstream will see. That was true at the old call site
and stopped being true one block later.

### Fix

Move the repair to after the turn-hooks block, so it is the last stage
that can invalidate a reference:

- It still runs **after** the deferral injection, so the tool just
injected counts as present — the main loop strips nothing and the frozen
prefix stays byte-identical.
- Nothing past the new call site mutates `body["tools"]` on the outbound
path. (The two later `continuation_body["tools"]` assignments build a
*derived* body from the already-repaired `body`, so they inherit the
repair.)
- It still runs before the consistency token re-count, so `tok_after`
continues to reflect the repaired messages.
- It remains unconditional (not gated on `HEADROOM_TOOL_SEARCH`, not
gated on `_bypass`), so transcripts poisoned before the flag was turned
off still recover.

`strip_unsupported_tool_search_blocks` is copy-on-write and returns the
original `messages` object by identity when nothing is removed, so
relocating the call does not change the no-op path.

### Severity

Latent. No turn hook ships in-tree, so this cannot fire on a default
install — it is reachable only through a third-party registered hook
that shrinks the tools array. Filing the fix now so the ordering
constraint is enforced by a test rather than rediscovered.

Closes #2888

## 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)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `headroom/proxy/handlers/anthropic.py`: the tool-search history repair
block moves from just after the deferral injection to just after the
turn-hooks block. The comment now states the ordering constraint in both
directions (after injection, after hooks) so the next person to add a
stage knows where the boundary is. No logic change.
- `tests/test_proxy/test_tool_search_repair_after_turn_hooks.py` (new):
two handler-level regressions. Ordering is the whole property under
test, so a unit test of the helper cannot see it — these drive the real
handler through `TestClient` and assert on the forwarded body.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed — no in-tree turn hook exists to exercise
this against a live API key; the handler-level test below is the
substitute, see Not tested.

### Test Output

```text
$ uv run --extra dev pytest tests/test_proxy/ -q
======================= 241 passed, 1 warning in 35.82s ========================

$ uvx ruff check headroom tests
All checks passed!

$ uv run --extra dev mypy headroom
Success: no issues found in 515 source files
```

## Real Behavior Proof

- Environment: macOS 15 (darwin 24.6.0), Python 3.10.18, pytest 9.0.3,
in a worktree off `upstream/main` at 2f2950a6. No live Anthropic key:
the proxy handler is driven end to end through
`fastapi.testclient.TestClient` with `_retry_request` stubbed, so the
assertion is on the exact body that would have been sent upstream.
- Exact command / steps: (1) on the branch as submitted, `uv run --extra
dev pytest tests/test_proxy/test_tool_search_repair_after_turn_hooks.py
-q` -> 2 passed; (2) revert ONLY the handler ordering change while
keeping the new tests, `git checkout HEAD~1 --
headroom/proxy/handlers/anthropic.py`, and re-run the same command. The
request under test carries a `tool_search_tool_result` referencing
`Grep`, a `tools` array containing `Grep`, and a registered turn hook
that removes `Grep`.
- Observed result: with the fix reverted the primary test fails on
exactly the shape upstream 400s on, because the forwarded body still
carries a `tool_reference` naming a tool the turn hook had already
removed from `tools`. Restoring the handler change turns it green. The
second test passes in both states by design: it pins the converse (a
hook that leaves `tools` alone must not cause over-stripping), so the
fix cannot regress into "strip always". Verbatim output of the reverted
run:

```text
$ git checkout HEAD~1 -- headroom/proxy/handlers/anthropic.py   # revert ONLY the ordering fix
$ uv run --extra dev pytest tests/test_proxy/test_tool_search_repair_after_turn_hooks.py -q
collected 2 items
tests/test_proxy/test_tool_search_repair_after_turn_hooks.py F.          [100%]
=================================== FAILURES ===================================
____________ test_repair_sees_the_tools_array_the_hook_left_behind _____________
tests/test_proxy/test_tool_search_repair_after_turn_hooks.py:166: in test_repair_sees_the_tools_array_the_hook_left_behind
    assert _referenced_tool_names(forwarded) == []
E   AssertionError: assert ['Grep'] == []
E
E     Left contains one more item: 'Grep'
FAILED tests/test_proxy/test_tool_search_repair_after_turn_hooks.py::test_repair_sees_the_tools_array_the_hook_left_behind
```

- Not tested: no live-API reproduction of the 400 itself, since
triggering it needs a third-party turn hook that removes a tool and none
ships in-tree (the assertion above is on the forwarded body, which is
the input that produces the 400); no streaming-path variant, since the
repair mutates `body` upstream of the stream/buffered split so both
inherit it but only the buffered path is asserted; no performance
measurement, since the change moves an existing call ~60 lines later in
the same function and adds no work.

## 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
- [ ] I have made corresponding changes to the documentation — N/A, no
user-facing or configuration surface changes
- [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`

## Additional Notes

- **CI:** `test (4)` failed on
`tests/test_tokenizer_count_offload.py::test_count_tokens_offloaded_keeps_loop_responsive`
with `assert 0 >= 5`. That is an event-loop-responsiveness timing
assertion under a shared runner, and it is unrelated to this diff —
nothing here touches the tokenizer or the offload path. It passes
locally (`10 passed in 1.43s`). I do not have rerun permission on this
fork PR (`gh run rerun` → `cannot be rerun`), so a maintainer rerun is
needed to clear it.
- **Codecov:** reports "All modified and coverable lines are covered by
tests". The accompanying warning is the repo-level "install the Codecov
app" notice, not a finding against this PR.
- Surfaced while confirming that #2807 and #2848 supersede #2507, which
is now closed as such.

🤖 Generated with [Claude Code](https://claude.com/claude-code)
This commit is contained in:
gglucass
2026-08-11 18:10:27 +02:00
committed by GitHub
parent 1a04c957f5
commit c6f99482e1
2 changed files with 210 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
headroom/proxy/handlers/anthropic.py
+34 -29
View File
@@ -2504,35 +2504,6 @@ 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
@@ -2579,6 +2550,40 @@ class AnthropicHandlerMixin:
int(tags.get("turn_hook_tools_saved_tokens", 0) or 0) + _th_saved
)
# 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. Unconditional (not gated on the flag) so transcripts poisoned
# before the flag was turned off still recover.
#
# ORDERING (#2888): this must be the LAST stage that can invalidate a
# tool_reference, so it runs after BOTH the deferral injection above (the
# tool we just added counts as present, so the main loop strips nothing
# and the prefix is untouched) AND the turn hooks (a hook may rewrite the
# tools array, and repairing before it validated against a stale view).
# Nothing past this point mutates `body["tools"]` on the outbound path.
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,
)
# Consistency: report tok_before/tok_after with ONE tokenizer. The pipeline
# and the handler use different token estimators, and cache-mode branches
# can leave original_tokens (handler, line ~1049) and optimized_tokens
tests/test_proxy/test_tool_search_repair_after_turn_hooks.py
+176
View File
@@ -0,0 +1,176 @@
"""Tool-search history repair must run AFTER the turn hooks (#2888).
``strip_unsupported_tool_search_blocks`` (#2807) validates every replayed
``tool_reference`` against the request's ``tools`` array. A registered turn hook
may rewrite that array, so repairing before the hook validates against a stale
view: the reference looks resolvable, the hook then drops the tool it named, and
upstream 400s with ``Tool reference 'X' not found in available tools``.
These drive the real handler and assert on the forwarded body, because ordering
is the whole property under test -- a unit test of the helper cannot see it.
"""
from __future__ import annotations
from types import SimpleNamespace
import httpx
import pytest
pytest.importorskip("fastapi")
from fastapi.testclient import TestClient
from headroom.proxy.server import ProxyConfig, create_app
from headroom.proxy.turn_hooks import clear_turn_hooks, register_turn_hook
_SEARCH_TOOL = {"type": "tool_search_tool_20250917", "name": "tool_search"}
_GREP = {"name": "Grep", "description": "search files", "input_schema": {"type": "object"}}
_READ = {"name": "Read", "description": "read a file", "input_schema": {"type": "object"}}
# A transcript that already carries a resolved tool-search round trip for `Grep`.
_POISONED_MESSAGES = [
{"role": "user", "content": "find the thing"},
{
"role": "assistant",
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_1",
"name": "tool_search_tool_20250917",
"input": {"query": "grep"},
},
{
"type": "tool_search_tool_result",
"tool_use_id": "srvtoolu_1",
"content": {
"type": "tool_search_tool_result_content",
"tool_references": [{"type": "tool_reference", "tool_name": "Grep"}],
},
},
],
},
{"role": "user", "content": "now use it"},
]
@pytest.fixture(autouse=True)
def _clean_registry():
clear_turn_hooks()
yield
clear_turn_hooks()
class _DropToolHook:
"""Turn hook that removes one tool from the outbound array."""
def __init__(self, name: str):
self._name = name
def on_request(self, ctx) -> None: # noqa: ANN001
if ctx.tools:
ctx.tools = [t for t in ctx.tools if t.get("name") != self._name]
class _InertHook:
def on_request(self, ctx) -> None: # noqa: ANN001, ARG002
return None
def _run(hook) -> dict: # noqa: ANN001
"""POST a poisoned transcript through the handler, return the forwarded body."""
captured: dict[str, object] = {}
register_turn_hook(hook)
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
image_optimize=False,
)
with TestClient(create_app(config)) as client:
proxy = client.app.state.proxy
proxy.pipeline_extensions.emit = lambda *args, **kwargs: SimpleNamespace(
messages=kwargs.get("messages"),
tools=kwargs.get("tools"),
headers=kwargs.get("headers"),
metadata=kwargs.get("metadata"),
)
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
captured["body"] = body
return httpx.Response(
200,
json={
"id": "msg_repair_order",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "ok"}],
"usage": {"input_tokens": 20, "output_tokens": 3},
},
)
proxy._retry_request = _fake_retry
response = client.post(
"/v1/messages",
headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"},
json={
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"messages": _POISONED_MESSAGES,
"tools": [_SEARCH_TOOL, _GREP, _READ],
},
)
assert response.status_code == 200
return captured["body"] # type: ignore[return-value]
def _referenced_tool_names(body: dict) -> list[str]:
names = []
for message in body.get("messages", []):
content = message.get("content")
if not isinstance(content, list):
continue
for block in content:
if not isinstance(block, dict) or block.get("type") != "tool_search_tool_result":
continue
inner = block.get("content")
entries = inner.get("tool_references") if isinstance(inner, dict) else inner
for entry in entries or []:
names.append(str(entry.get("tool_name") or entry.get("name")))
return names
def _block_types(body: dict) -> list[str]:
types = []
for message in body.get("messages", []):
content = message.get("content")
if isinstance(content, list):
types.extend(str(b.get("type")) for b in content if isinstance(b, dict))
return types
def test_repair_sees_the_tools_array_the_hook_left_behind() -> None:
"""A hook that drops `Grep` must leave no dangling reference to it."""
forwarded = _run(_DropToolHook("Grep"))
assert "Grep" not in [t.get("name") for t in forwarded["tools"]]
# The whole pair goes: an orphaned server_tool_use 400s on its own.
assert _referenced_tool_names(forwarded) == []
assert "tool_search_tool_result" not in _block_types(forwarded)
assert "server_tool_use" not in _block_types(forwarded)
def test_repair_leaves_resolvable_history_alone_when_the_hook_keeps_the_tool() -> None:
"""The converse: no over-stripping when the hook does not touch `tools`."""
forwarded = _run(_InertHook())
assert _referenced_tool_names(forwarded) == ["Grep"]
assert "tool_search_tool_result" in _block_types(forwarded)