Merge 3610a3f62a into 5337c68dd3

litellm/integrations/websearch_interception/handler.py

@@ -20,6 +20,7 @@ from litellm.anthropic_interface import messages as anthropic_messages
 from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.integrations.websearch_interception.tools import (
+    collect_rewritten_tool_names,
     get_litellm_web_search_tool,
     get_litellm_web_search_tool_openai,
     get_litellm_web_search_tool_responses,
@@ -27,6 +28,7 @@ from litellm.integrations.websearch_interception.tools import (
     is_web_search_tool,
     is_web_search_tool_chat_completion,
     is_web_search_tool_responses,
+    rewrite_web_search_tool_choice,
 )
 from litellm.integrations.websearch_interception.transformation import (
     WebSearchTransformation,
@@ -367,6 +369,18 @@ class WebSearchInterceptionLogger(CustomLogger):
 
         kwargs["tools"] = converted_tools
 
+        # A request like Claude Code's web-search call carries
+        # ``tool_choice={"type": "tool", "name": "web_search"}``; renaming
+        # the tool above without renaming the choice leaves the provider
+        # with a dangling reference and a "Tool 'web_search' not found"
+        # 400 (issue #30822). Keep the two in lockstep, scoped to the
+        # exact names we just renamed in ``tools`` so we don't hijack a
+        # Cowork-style ``tool_choice={"name": "WebSearch"}`` that
+        # legitimately points at a client-side tool we left alone.
+        if "tool_choice" in kwargs:
+            rewritten_names = collect_rewritten_tool_names(tools)
+            kwargs["tool_choice"] = rewrite_web_search_tool_choice(kwargs["tool_choice"], rewritten_names)
+
         if kwargs.get("stream"):
             verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False")
             kwargs["stream"] = False
@@ -546,6 +560,13 @@ class WebSearchInterceptionLogger(CustomLogger):
         if "tool_choice" in kwargs:
             kwargs["tool_choice"] = self._sync_forced_tool_choice(kwargs.get("tool_choice"), converted_tools)
 
+        # Direct litellm.acompletion callers reach this path without the
+        # deployment hook running, so the same scoped tool_choice rewrite
+        # has to be applied here too (issue #30822).
+        if "tool_choice" in kwargs:
+            rewritten_names = collect_rewritten_tool_names(tools)
+            kwargs["tool_choice"] = rewrite_web_search_tool_choice(kwargs["tool_choice"], rewritten_names)
+
         # Also convert here for direct callers that bypass the deployment hook.
         if kwargs.get("stream"):
             verbose_logger.debug("WebSearchInterception: Converting stream=True to stream=False")

litellm/integrations/websearch_interception/tools.py

@@ -290,3 +290,74 @@ def is_web_search_tool(tool: dict[str, Any]) -> bool:
         return True
 
     return False
+
+
+def collect_rewritten_tool_names(original_tools: list[dict[str, Any]]) -> set[str]:
+    """Return the set of names from ``original_tools`` that the converter
+    rewrote to ``LITELLM_WEB_SEARCH_TOOL_NAME``. Used to scope
+    ``rewrite_web_search_tool_choice`` to only the names that were
+    actually renamed in this request — a bare ``WebSearch`` with no
+    ``input_schema`` is rewritten, but a Cowork client-side
+    ``WebSearch`` tool that ships with ``input_schema`` is not, and the
+    tool_choice rewriter must follow the same gate.
+    """
+    names: set[str] = set()
+    for tool in original_tools:
+        if not isinstance(tool, dict):
+            continue
+        if not is_web_search_tool(tool):
+            continue
+        name = tool.get("name")
+        if isinstance(name, str) and name:
+            names.add(name)
+        function = tool.get("function")
+        if isinstance(function, dict):
+            function_name = function.get("name")
+            if isinstance(function_name, str) and function_name:
+                names.add(function_name)
+    return names
+
+
+def rewrite_web_search_tool_choice(tool_choice: Any, rewritten_names: set[str]) -> Any:
+    """Rename ``tool_choice`` entries that point at a tool the converter
+    just renamed to ``LITELLM_WEB_SEARCH_TOOL_NAME`` so the choice
+    keeps resolving against the new tools array.
+
+    Without this, ``tool_choice={"type": "tool", "name": "web_search"}``
+    (Anthropic style) or
+    ``tool_choice={"type": "function", "function": {"name": "web_search"}}``
+    (OpenAI style) lands on Bedrock after we've already renamed the tool
+    entry, and the provider rejects with
+    ``"Tool 'web_search' not found in provided tools"`` (issue #30822).
+
+    Scoped to ``rewritten_names`` (the output of
+    :func:`collect_rewritten_tool_names`) so a Cowork-style request that
+    forces ``tool_choice={"type": "tool", "name": "WebSearch"}`` while
+    keeping its client-side ``WebSearch`` tool (which carries an
+    ``input_schema`` and is not converted) does not get rewritten — that
+    would produce the inverse of the original bug.
+
+    Returns the (possibly rewritten) tool_choice. Non-dict values, plain
+    strings (``"auto"``/``"none"``/``"any"``/``"required"``), and entries
+    naming a tool we did not rewrite are returned unchanged.
+    """
+    if not isinstance(tool_choice, dict) or not rewritten_names:
+        return tool_choice
+
+    # Anthropic shape: {"type": "tool", "name": "..."}
+    name = tool_choice.get("name")
+    if isinstance(name, str) and name in rewritten_names:
+        rewritten = dict(tool_choice)
+        rewritten["name"] = LITELLM_WEB_SEARCH_TOOL_NAME
+        return rewritten
+
+    # OpenAI shape: {"type": "function", "function": {"name": "..."}}
+    function = tool_choice.get("function")
+    if isinstance(function, dict):
+        function_name = function.get("name")
+        if isinstance(function_name, str) and function_name in rewritten_names:
+            rewritten = dict(tool_choice)
+            rewritten["function"] = {**function, "name": LITELLM_WEB_SEARCH_TOOL_NAME}
+            return rewritten
+
+    return tool_choice

tests/test_litellm/integrations/websearch_interception/test_tool_choice_rewrite.py

@@ -0,0 +1,157 @@
+"""Regression coverage for issue #30822.
+
+`websearch_interception` converts the `tools` array but, before this fix,
+left the `tool_choice` field untouched. A Claude Code-style request that
+forced the web search tool via
+``tool_choice={"type": "tool", "name": "web_search"}`` was sent to the
+provider after the tool entry had been renamed to ``litellm_web_search``,
+and the provider returned ``Tool 'web_search' not found in provided tools``.
+
+The rewriter is scoped to names we actually renamed in ``tools`` so a
+Cowork-style request that points ``tool_choice`` at a client-side
+``WebSearch`` tool (which carries an ``input_schema`` and is not
+converted) is not hijacked the other way around.
+"""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../../../.."))
+
+from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME
+from litellm.integrations.websearch_interception.tools import (
+    collect_rewritten_tool_names,
+    rewrite_web_search_tool_choice,
+)
+
+
+def _rewritable(tool_name: str):
+    """Helper: build a ``rewritten_names`` set as if a tool named
+    ``tool_name`` had been converted by the handler."""
+    return {tool_name}
+
+
+def test_rewrites_anthropic_style_tool_choice_naming_web_search():
+    out = rewrite_web_search_tool_choice(
+        {"type": "tool", "name": "web_search"}, _rewritable("web_search")
+    )
+    assert out == {"type": "tool", "name": LITELLM_WEB_SEARCH_TOOL_NAME}
+
+
+def test_rewrites_anthropic_style_tool_choice_naming_legacy_WebSearch():
+    out = rewrite_web_search_tool_choice(
+        {"type": "tool", "name": "WebSearch"}, _rewritable("WebSearch")
+    )
+    assert out == {"type": "tool", "name": LITELLM_WEB_SEARCH_TOOL_NAME}
+
+
+def test_rewrites_openai_style_tool_choice_naming_web_search():
+    out = rewrite_web_search_tool_choice(
+        {"type": "function", "function": {"name": "web_search"}},
+        _rewritable("web_search"),
+    )
+    assert out == {
+        "type": "function",
+        "function": {"name": LITELLM_WEB_SEARCH_TOOL_NAME},
+    }
+
+
+def test_leaves_unrelated_tool_choice_unchanged():
+    original = {"type": "tool", "name": "my_other_tool"}
+    assert (
+        rewrite_web_search_tool_choice(original, _rewritable("web_search")) == original
+    )
+
+
+def test_leaves_string_modes_unchanged():
+    for mode in ("auto", "none", "any", "required"):
+        assert (
+            rewrite_web_search_tool_choice(mode, _rewritable("web_search")) == mode
+        )
+
+
+def test_leaves_none_unchanged():
+    assert rewrite_web_search_tool_choice(None, _rewritable("web_search")) is None
+
+
+def test_preserves_extra_fields_on_anthropic_style():
+    out = rewrite_web_search_tool_choice(
+        {"type": "tool", "name": "web_search", "disable_parallel_tool_use": True},
+        _rewritable("web_search"),
+    )
+    assert out == {
+        "type": "tool",
+        "name": LITELLM_WEB_SEARCH_TOOL_NAME,
+        "disable_parallel_tool_use": True,
+    }
+
+
+def test_preserves_extra_fields_on_openai_style():
+    out = rewrite_web_search_tool_choice(
+        {
+            "type": "function",
+            "function": {"name": "web_search", "extra": "x"},
+            "outer": "y",
+        },
+        _rewritable("web_search"),
+    )
+    assert out == {
+        "type": "function",
+        "function": {"name": LITELLM_WEB_SEARCH_TOOL_NAME, "extra": "x"},
+        "outer": "y",
+    }
+
+
+def test_does_not_mutate_input():
+    original = {"type": "tool", "name": "web_search"}
+    snapshot = dict(original)
+    rewrite_web_search_tool_choice(original, _rewritable("web_search"))
+    assert original == snapshot
+
+
+def test_empty_rewritten_names_leaves_choice_unchanged():
+    """If no tool was actually renamed, the rewriter is a no-op."""
+    original = {"type": "tool", "name": "web_search"}
+    assert rewrite_web_search_tool_choice(original, set()) is original
+
+
+def test_cowork_tool_choice_pointing_at_client_side_WebSearch_not_rewritten():
+    """Greptile-flagged guard for #30872: a request that ships its own
+    client-side ``WebSearch`` tool (with ``input_schema``) keeps
+    ``WebSearch`` in the tools array because ``is_web_search_tool`` skips
+    it. The rewriter must follow the same gate and leave a forcing
+    ``tool_choice={"name": "WebSearch"}`` alone — otherwise we would
+    produce the inverse of the original bug.
+    """
+    cowork_tools = [
+        {
+            "name": "WebSearch",
+            "input_schema": {"type": "object", "properties": {}, "required": []},
+        }
+    ]
+    rewritten = collect_rewritten_tool_names(cowork_tools)
+    assert rewritten == set()
+
+    tool_choice = {"type": "tool", "name": "WebSearch"}
+    out = rewrite_web_search_tool_choice(tool_choice, rewritten)
+    assert out == tool_choice
+
+
+def test_collect_rewritten_tool_names_picks_up_converted_tools():
+    """Names that ``is_web_search_tool`` matches end up in the returned set
+    so the rewriter knows which ``tool_choice`` references to rename.
+    Covers Anthropic-native (``web_search_20250305`` → ``web_search``),
+    Claude Code (bare ``web_search`` + a type), and legacy bare
+    ``WebSearch`` (no schema) — but NOT a Cowork client-side
+    ``WebSearch`` carrying an ``input_schema``.
+    """
+    tools = [
+        {"type": "web_search_20250305", "name": "web_search"},
+        {"name": "WebSearch"},  # legacy interception marker
+        {"name": "calculator"},  # untouched
+        {
+            "name": "WebSearch",
+            "input_schema": {"type": "object", "properties": {}, "required": []},
+        },  # Cowork client-side tool — not rewritten
+    ]
+    assert collect_rewritten_tool_names(tools) == {"web_search", "WebSearch"}