Tejas/tool search deferral (headroomlabs-ai#1885)

## Description

<!-- Briefly explain the change and why it is needed. -->

Closes #

## Type of Change

- [ ] 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

- 

## Testing

<!-- Check what you actually ran, then paste the real command output
below. -->

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

### Test Output

```text
# Paste relevant command output or artifact links here
```

## Real Behavior Proof

- Environment:
- Exact command / steps:
- Observed result:
- Not tested:

## Review Readiness

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

## Checklist

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

## Screenshots (if applicable)

Add screenshots to help explain your changes.

## Additional Notes

<!-- Mention any N/A checklist items, tradeoffs, follow-ups, or
maintainer context. -->

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: chopratejas

headroom/proxy/handlers/anthropic.py

@@ -38,6 +38,33 @@
 logger = logging.getLogger("headroom.proxy")
 
 
+def _strip_streaming_only_content_fields(messages: Any) -> None:
+    """Remove streaming-only ``index`` keys from request content blocks, in place.
+
+    ``index`` is a field Anthropic emits on streaming RESPONSE content-block deltas
+    (see proxy/handlers/streaming.py). It is not part of the request-message schema, so
+    forwarding it upstream triggers a 400 ("...content.N.text.index: Extra inputs are
+    not permitted") that aborts multi-turn sessions once a client echoes a reconstructed
+    assistant turn back. Strip it (including nested tool_result content) so requests are
+    always schema-valid.
+    """
+    if not isinstance(messages, list):
+        return
+    for message in messages:
+        if isinstance(message, dict):
+            _strip_index_from_content_blocks(message.get("content"))
+
+
+def _strip_index_from_content_blocks(content: Any) -> None:
+    if not isinstance(content, list):
+        return
+    for block in content:
+        if isinstance(block, dict):
+            block.pop("index", None)
+            # tool_result blocks nest their own content list of blocks.
+            _strip_index_from_content_blocks(block.get("content"))
+
+
 class AnthropicHandlerMixin:
     """Mixin providing Anthropic API handler methods for HeadroomProxy."""
 
@@ -680,6 +707,16 @@ async def _finalize_pre_upstream() -> None:
                 body["model"] = model
                 body_mutation_tracker.mark_mutated("sanitize_model_id")
             messages = body.get("messages", [])
+            # Strip streaming-only "index" keys from request content blocks BEFORE any
+            # prefix-cache tracking or compression. The proxy's streaming reconstruction
+            # tags assistant blocks with an "index" for SSE re-emission; clients (e.g.
+            # opencode) persist that assistant message and echo it back next turn, but
+            # "index" is a response-delta field that Anthropic REJECTS in a request
+            # ("messages.N.content.0.text.index: Extra inputs are not permitted", 400),
+            # aborting multi-turn sessions. Canonicalizing here (in place, so body,
+            # original, forwarded, and the recorded/replayed prefix are all identical)
+            # keeps it cache-safe: overlay_cached_prefix replays the same stripped bytes.
+            _strip_streaming_only_content_fields(messages)
             pipeline_provider = provider_name
             pipeline_path = request.url.path if upstream_base_url else "/v1/messages"
             pipeline_stream = bool(body.get("stream", False) or force_stream)
@@ -2044,6 +2081,14 @@ class _DeferredCompressionResult:
                 if remembered_event.headers is not None:
                     headers = remembered_event.headers
 
+            # Final sanitization of the FORWARDED body: strip streaming-only "index"
+            # keys from content blocks. In cache mode the forwarded prefix is replayed
+            # from Headroom's own recorded/reconstructed messages (streaming.py tags
+            # blocks with "index"), so the inbound-side strip above doesn't cover it —
+            # this catches both the client-echoed and the cache-replayed forms. Applied
+            # deterministically to the exact bytes forwarded (and thus recorded as the
+            # next prefix), so Anthropic always caches/matches the same stripped prefix.
+            _strip_streaming_only_content_fields(optimized_messages)
             # Update body
             body["messages"] = optimized_messages
             if tools or _original_tools is not None:
@@ -2085,6 +2130,51 @@ class _DeferredCompressionResult:
                 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
+            # 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.
+            # Deterministic → the tools prefix still prompt-caches. Safe for opencode:
+            # its @ai-sdk/anthropic parses the server_tool_use / tool_search_tool_result
+            # round-trip natively, and its cache breakpoints sit on messages (never
+            # tools), so defer_loading cannot collide with cache_control. We COUNT the
+            # deferred tool-schema tokens as the tool-search input-token saving (those
+            # bytes are excluded from context this turn); the response usage confirms it.
+            #
+            # FIRST-PARTY ANTHROPIC ONLY: the tool_search_tool_* type + defer_loading
+            # here use the first-party Claude API shape (GA, no beta header). Bedrock
+            # (``anthropic_backend``) and Vertex/gateway providers gate tool search
+            # differently, so scope the injection to provider "anthropic" over the
+            # direct API and leave those paths untouched.
+            if (
+                provider_name == "anthropic"
+                and getattr(self, "anthropic_backend", None) is None
+                and os.environ.get("HEADROOM_TOOL_SEARCH", "").strip().lower()
+                in ("1", "true", "yes", "on", "auto")
+            ):
+                from headroom.proxy.helpers import inject_tool_search_deferral
+
+                _ts_before = body.get("tools")
+                _ts_after = inject_tool_search_deferral(_ts_before)
+                if _ts_after is not _ts_before:
+                    _ts_deferred = [
+                        t for t in _ts_after if isinstance(t, dict) and t.get("defer_loading")
+                    ]
+                    try:
+                        _ts_saved_tokens = tokenizer.count_text(
+                            json.dumps(_ts_deferred, default=str)
+                        )
+                    except Exception:
+                        _ts_saved_tokens = 0
+                    body["tools"] = _ts_after
+                    tools = _ts_after
+                    tags["tool_search_deferred_tools"] = len(_ts_deferred)
+                    tags["tool_search_deferred_tokens"] = _ts_saved_tokens
+                    transforms_applied.append(
+                        f"router:tool_search_deferral:{len(_ts_deferred)}tools:"
+                        f"{_ts_saved_tokens}tok"
+                    )
+
             # Output shaping (opt-in via HEADROOM_OUTPUT_SHAPER): verbosity
             # steering appended to the system-prompt tail + effort routing on
             # mechanical tool_result continuations. Runs after every other

headroom/proxy/handlers/openai.py

@@ -1691,6 +1691,24 @@ def _add_timing(name: str, started_at: float) -> None:
                     tools_bytes_saved=tools_before_bytes - tools_after_bytes,
                 )
 
+        # Server-side Tool Search deferral (OpenAI Responses, gpt-5.4+): mark
+        # non-core function/MCP tools defer_loading + inject {"type": "tool_search"}
+        # so OpenAI keeps their heavy parameter schemas out of the model's context
+        # until searched (every tool stays callable, cache preserved). No-op for
+        # older models / small tool sets / clients already using tool search. The
+        # deferred defs still ride in the request body (OpenAI needs them to load
+        # on demand), so this is a provider-side context saving, not a request-byte
+        # one — hence a transform tag but no tokens_saved claim.
+        from headroom.proxy.helpers import inject_tool_search_deferral_openai
+
+        _deferred_tools = inject_tool_search_deferral_openai(working.get("tools"), model)
+        if _deferred_tools is not working.get("tools"):
+            if working is payload:
+                working = copy.deepcopy(payload)
+            working["tools"] = _deferred_tools
+            modified = True
+            transforms.append("openai:responses:tool_search_deferral")
+
         live_units_started = time.perf_counter()
         (
             router_payload,

headroom/proxy/helpers.py

@@ -14,6 +14,7 @@
 import logging
 import os
 import random
+import re
 import threading
 import time
 from collections import OrderedDict
@@ -3145,3 +3146,216 @@ def reset_tool_search_hint_state() -> None:
     global _tool_search_hint_emitted
     with _tool_search_hint_lock:
         _tool_search_hint_emitted = False
+
+
+# ---------------------------------------------------------------------------
+# Server-side Tool Search injection (opencode / non-Claude-Code clients).
+#
+# Clients that eagerly materialize every tool schema (opencode ships ~135 tool
+# defs ≈ 28k tokens on EVERY request) never opt into Anthropic's Tool Search
+# Tool themselves. Unlike the Claude Code case above — where the schemas are
+# already in the client's own context and the proxy can't reverse it — a plain
+# API client's tools live only in the request body, so the proxy CAN defer them:
+# mark the non-core tools ``defer_loading: true`` and inject a tool_search tool.
+# Anthropic then excludes deferred tools from the context window (they stop
+# counting as input tokens until the model searches for one), while every tool
+# stays callable. Deterministic output → the tools prefix still prompt-caches.
+# ---------------------------------------------------------------------------
+
+# Core coding tools kept non-deferred so routine edit/read/run loops never pay a
+# search round-trip. Everything else (Slack/Linear/Sentry/Notion/Snowflake/…) is
+# deferred and loaded on demand. Anthropic recommends keeping the 3–5 (here a few
+# more) most frequent tools resident.
+_TOOL_SEARCH_CORE_TOOLS = frozenset(
+    {
+        "bash",
+        "bash_background",
+        "bash_background_output",
+        "bash_background_wait",
+        "bash_background_kill",
+        "read",
+        "write",
+        "edit",
+        "multiedit",
+        "apply_patch",
+        "glob",
+        "grep",
+        "task",
+        "todowrite",
+        "todoread",
+        "webfetch",
+        "question",
+        "skill",
+    }
+)
+_TOOL_SEARCH_DEFAULT_TYPE = "tool_search_tool_regex_20251119"
+_TOOL_SEARCH_DEFAULT_NAME = "tool_search_tool_regex"
+# Below this many tools the ~search round-trip isn't worth it (Anthropic's own
+# guidance: standard calling is better under ~10 tools).
+_TOOL_SEARCH_MIN_TOOLS = 12
+
+
+def inject_tool_search_deferral(
+    tools: Any,
+    *,
+    core_tools: frozenset[str] = _TOOL_SEARCH_CORE_TOOLS,
+    search_type: str = _TOOL_SEARCH_DEFAULT_TYPE,
+    search_name: str = _TOOL_SEARCH_DEFAULT_NAME,
+) -> Any:
+    """Return a new ``tools`` list with non-core tools deferred + a search tool
+    injected, or the original list unchanged when injection doesn't apply.
+
+    No-op when: not a list, fewer than ``_TOOL_SEARCH_MIN_TOOLS``, a tool_search
+    tool is already present (client already defers), or nothing would be deferred.
+
+    Invariants enforced (else Anthropic 400s): the search tool is never deferred;
+    at least one tool stays non-deferred; a deferred tool never carries
+    ``cache_control`` — if the client's tools cache breakpoint sat on a now-deferred
+    tool, it is moved to the last non-deferred real tool so the (smaller) tools
+    prefix still caches.
+    """
+    if not isinstance(tools, list) or len(tools) < _TOOL_SEARCH_MIN_TOOLS:
+        return tools
+    for tool in tools:
+        if isinstance(tool, dict) and str(tool.get("type", "")).startswith(
+            _TOOL_SEARCH_TOOL_TYPE_PREFIX
+        ):
+            return tools  # client already uses tool search — leave it alone
+
+    search_tool = {"type": search_type, "name": search_name}
+    out: list[Any] = [search_tool]
+    deferred = 0
+    dropped_cache_control = False
+    last_resident_real: dict[str, Any] | None = None
+    resident_has_cache_control = False
+
+    for tool in tools:
+        if not isinstance(tool, dict) or tool.get("type") or tool.get("name") in core_tools:
+            # Non-dict, server/typed tools (web_search, computer, …), and core
+            # tools stay resident and unchanged.
+            out.append(tool)
+            if isinstance(tool, dict) and not tool.get("type"):
+                last_resident_real = tool
+                resident_has_cache_control = resident_has_cache_control or bool(
+                    tool.get("cache_control")
+                )
+            continue
+        new_tool = dict(tool)
+        new_tool["defer_loading"] = True
+        if new_tool.pop("cache_control", None) is not None:
+            dropped_cache_control = True
+        out.append(new_tool)
+        deferred += 1
+
+    if deferred == 0:
+        return tools  # nothing to defer → don't perturb the cache prefix
+    # Preserve a tools cache breakpoint: if we stripped cache_control off a
+    # deferred tool and no resident tool carries one, move it to the last
+    # resident real tool (never the search tool, to keep its shape canonical).
+    if dropped_cache_control and not resident_has_cache_control and last_resident_real is not None:
+        last_resident_real["cache_control"] = {"type": "ephemeral"}
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Server-side Tool Search injection — OpenAI Responses API (gpt-5.4+).
+#
+# The OpenAI-side analogue of inject_tool_search_deferral above. OpenAI shipped
+# the same idea for the Responses API on gpt-5.4+: mark a function/MCP tool
+# ``defer_loading: true`` and add a ``{"type": "tool_search"}`` tool, and OpenAI
+# keeps the deferred tools' heavy parameter schemas OUT of the model's context
+# (only name+description remain) until the model searches for one — while every
+# tool stays callable and the prompt cache is preserved. Same win as Anthropic
+# (~15-25k tool-schema tokens -> ~200) for clients that ship a big tool surface
+# and never opt into tool search themselves (opencode, plain API clients).
+#
+# Differences from the Anthropic path that require a separate function:
+#   * Responses function tools carry ``type: "function"`` (Anthropic real tools
+#     have no ``type``), so the resident/defer test is inverted — we defer
+#     ``function`` (non-core) and ``mcp`` tools and keep OTHER typed/hosted tools
+#     (web_search, file_search, code_interpreter, computer, image_generation, and
+#     the search tool itself) resident.
+#   * Model-gated: only gpt-5.4+ support it; older models 400 on the fields.
+#   * No ``cache_control`` (OpenAI caches automatically), so no breakpoint move.
+# ---------------------------------------------------------------------------
+
+_OPENAI_TOOL_SEARCH_TYPE = "tool_search"
+_OPENAI_TOOL_SEARCH_MIN_TOOLS = 12
+# gpt-5.4 is the first model with Responses tool_search (OpenAI docs). Version-
+# gated by default; overridable per deployment via a regex in
+# HEADROOM_OPENAI_TOOL_SEARCH_MODELS (matched against the model name) so new
+# model families can be enabled without a code edit + release.
+_OPENAI_TOOL_SEARCH_MIN_VERSION = (5, 4)
+
+
+def _model_supports_openai_tool_search(model: str | None) -> bool:
+    """True when an OpenAI model supports the Responses ``tool_search`` feature.
+
+    Default gate: ``gpt-<major>.<minor>`` >= 5.4. A regex in
+    ``HEADROOM_OPENAI_TOOL_SEARCH_MODELS`` (matched against the model name) wins
+    when set; a malformed pattern falls back to the version gate rather than
+    crashing.
+    """
+    if not model:
+        return False
+    override = os.environ.get("HEADROOM_OPENAI_TOOL_SEARCH_MODELS", "").strip()
+    if override:
+        try:
+            return re.search(override, model) is not None
+        except re.error:
+            pass  # malformed override → fall back to the version gate
+    match = re.match(r"gpt-(\d+)(?:\.(\d+))?", model.strip().lower())
+    if not match:
+        return False
+    major, minor = int(match.group(1)), int(match.group(2) or 0)
+    return (major, minor) >= _OPENAI_TOOL_SEARCH_MIN_VERSION
+
+
+def inject_tool_search_deferral_openai(
+    tools: Any,
+    model: str | None,
+    *,
+    core_tools: frozenset[str] = _TOOL_SEARCH_CORE_TOOLS,
+) -> Any:
+    """Return a new Responses ``tools`` list with non-core function/MCP tools
+    deferred + a ``{"type": "tool_search"}`` tool injected, or the original list
+    unchanged when injection doesn't apply.
+
+    No-op when: the model doesn't support tool search (gpt-5.4+ only), ``tools``
+    is not a list, there are fewer than ``_OPENAI_TOOL_SEARCH_MIN_TOOLS``, a
+    tool_search tool is already present (client already defers), or nothing would
+    be deferred. Core coding tools and hosted/typed tools (web_search,
+    file_search, code_interpreter, computer, …) stay resident and unchanged, so
+    routine edit/read/run loops never pay a search round-trip and the request
+    stays valid; the injected search tool is itself resident.
+    """
+    if not _model_supports_openai_tool_search(model):
+        return tools
+    if not isinstance(tools, list) or len(tools) < _OPENAI_TOOL_SEARCH_MIN_TOOLS:
+        return tools
+    for tool in tools:
+        if isinstance(tool, dict) and tool.get("type") == _OPENAI_TOOL_SEARCH_TYPE:
+            return tools  # client already uses tool search — leave it alone
+
+    out: list[Any] = [{"type": _OPENAI_TOOL_SEARCH_TYPE}]
+    deferred = 0
+    for tool in tools:
+        if not isinstance(tool, dict):
+            out.append(tool)
+            continue
+        ttype = tool.get("type")
+        # Deferrable: a non-core function, or an MCP server (OpenAI models are
+        # trained to search namespaces / MCP servers). Everything else — core
+        # coding tools and other hosted tools — stays resident.
+        deferrable = (ttype == "function" and tool.get("name") not in core_tools) or ttype == "mcp"
+        if deferrable and not tool.get("defer_loading"):
+            new_tool = dict(tool)
+            new_tool["defer_loading"] = True
+            out.append(new_tool)
+            deferred += 1
+        else:
+            out.append(tool)
+
+    if deferred == 0:
+        return tools  # nothing to defer → don't perturb the request / cache prefix
+    return out

pyproject.toml

@@ -323,6 +323,8 @@ constraint-dependencies = [
     "gitpython>=3.1.50",
     # GHSA-f4xh-w4cj-qxq8 (High) — transitive via langchain-core; fix at 0.8.18
     "langsmith>=0.9.0",
+    # CVE-2026-49825 (High, XSS) — transitive via lxml[html-clean]; fix at 0.4.5
+    "lxml-html-clean>=0.4.5",
 ]
 
 # Pin the project's package index to public PyPI. Without this, `uv lock`