diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index c379df778..8e12401b7 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -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 diff --git a/tests/test_proxy/test_tool_search_repair_after_turn_hooks.py b/tests/test_proxy/test_tool_search_repair_after_turn_hooks.py new file mode 100644 index 000000000..30d495760 --- /dev/null +++ b/tests/test_proxy/test_tool_search_repair_after_turn_hooks.py @@ -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)