fix(tool-choice): match named mcp entries by server label and let built-in entries allow their interception call

Backport of #50467 to lts/1.108.x.
Cherry-picked from 4755727276 (litellm_allowed_tools_flat_tool_choice). Resolved: kept only this pick's helper and import where the line lacks main's sandbox result types and core_helpers import. Test hunks that need main-only fixtures move to the line test commit.

litellm/integrations/code_interpreter_interception/handler.py

@@ -17,6 +17,7 @@ from pydantic import ValidationError
 import litellm
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.allowed_tools_tool_choice import allowed_tools_tool_choice_with_builtin_as_function
 from litellm.llms.base_llm.sandbox.transformation import (
     CodeExecutionResult,
     ContainerHandle,
@@ -141,6 +142,10 @@ class _SupportsOutput(Protocol):
 _CachedContainer: TypeAlias = tuple[ContainerHandle, SandboxToolParams | None, float, str | None]
 
 
+def _is_code_interpreter_tool(tool: Mapping[str, object]) -> bool:
+    return tool.get("type") == "code_interpreter"
+
+
 def _output_item_type(item: object) -> object:
     if isinstance(item, dict):
         item_mapping: Final[dict[str, object]] = item
@@ -301,8 +306,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
             (function_tool if isinstance(tool, dict) and tool.get("type") == "code_interpreter" else tool)
             for tool in tools
         ]
-        if self._tool_choice_targets_code_interpreter(kwargs.get("tool_choice")):
-            kwargs["tool_choice"] = self._get_function_tool_choice(call_type=call_type)
+        if "tool_choice" in kwargs:
+            kwargs["tool_choice"] = self._intercepted_tool_choice(kwargs["tool_choice"], call_type)
         return kwargs
 
     @staticmethod
@@ -386,6 +391,14 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
             or (isinstance(function, dict) and function.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME)
         )
 
+    @classmethod
+    def _intercepted_tool_choice(cls, tool_choice: object, call_type: CallTypes | None) -> object:
+        if cls._tool_choice_targets_code_interpreter(tool_choice):
+            return cls._get_function_tool_choice(call_type=call_type)
+        return allowed_tools_tool_choice_with_builtin_as_function(
+            tool_choice, _is_code_interpreter_tool, LITELLM_CODE_EXECUTION_TOOL_NAME
+        )
+
     def _resolve_provider(self, kwargs: dict[str, object]) -> str | None:
         provider: Final = kwargs.get("custom_llm_provider")
         if isinstance(provider, str) and provider:

litellm/integrations/websearch_interception/handler.py

@@ -9,7 +9,7 @@ server-side using litellm router's search tools.
 import asyncio
 import math
 import uuid
-from collections.abc import AsyncIterator, Mapping, Sequence
+from collections.abc import AsyncIterator, Callable, Mapping, Sequence
 from dataclasses import dataclass
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast
@@ -37,6 +37,7 @@ from litellm.integrations.websearch_interception.transformation import (
 from litellm.litellm_core_utils.agentic_loop_settings import (
     validated_max_agentic_loops,
 )
+from litellm.litellm_core_utils.allowed_tools_tool_choice import allowed_tools_tool_choice_with_builtin_as_function
 from litellm.llms.base_llm.search.transformation import SearchResponse
 from litellm.types.integrations.custom_logger import (
     CHAT_COMPLETION_AGENTIC_SURFACE,
@@ -246,6 +247,18 @@ class _ParentRequestCorrelation:
 _NO_ACOMPLETION_NAMED: Final[_AcompletionNamedParams] = {}
 
 
+def _tool_choice_naming_litellm_web_search(
+    kwargs: Mapping[str, object], is_search_tool: Callable[[Mapping[str, object]], bool]
+) -> Mapping[str, object]:
+    if "tool_choice" not in kwargs:
+        return MappingProxyType({})
+    return {
+        "tool_choice": allowed_tools_tool_choice_with_builtin_as_function(
+            kwargs["tool_choice"], is_search_tool, LITELLM_WEB_SEARCH_TOOL_NAME
+        )
+    }
+
+
 class WebSearchInterceptionLogger(CustomLogger):
     """
     CustomLogger that intercepts WebSearch tool calls for models that don't
@@ -493,13 +506,17 @@ class WebSearchInterceptionLogger(CustomLogger):
                 # Keep other tools as-is
                 converted_tools.append(tool)
 
-        kwargs["tools"] = converted_tools
+        converted_kwargs: Final[dict[str, object]] = {
+            **kwargs,
+            "tools": converted_tools,
+            **_tool_choice_naming_litellm_web_search(kwargs, lambda tool: is_search_tool(dict(tool))),
+        }
 
         if kwargs.get("stream"):
             verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False")
-            return as_converted_stream(kwargs, WEBSEARCH_INTERCEPTION_PREFIX)
+            return as_converted_stream(converted_kwargs, WEBSEARCH_INTERCEPTION_PREFIX)
 
-        return kwargs
+        return converted_kwargs
 
     def _convert_responses_tools(
         self, kwargs: Mapping[str, object], tools: Sequence[dict[str, object]]
@@ -514,7 +531,11 @@ class WebSearchInterceptionLogger(CustomLogger):
             get_litellm_web_search_tool_responses() if is_web_search_tool_responses(tool) else tool for tool in tools
         ]
 
-        converted_kwargs: Final = {**kwargs, "tools": converted_tools}
+        converted_kwargs: Final = {
+            **kwargs,
+            "tools": converted_tools,
+            **_tool_choice_naming_litellm_web_search(kwargs, is_web_search_tool_responses),
+        }
 
         if kwargs.get("stream"):
             verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False")

litellm/litellm_core_utils/allowed_tools_tool_choice.py

@@ -1,4 +1,4 @@
-from collections.abc import Iterable, Mapping
+from collections.abc import Callable, Iterable, Mapping
 from types import MappingProxyType
 from typing import Final, Literal
 
@@ -176,7 +176,7 @@ def _allowed_tool_names(tool_choice: object) -> frozenset[object]:
     if parsed is None:
         return frozenset()
     flattened: Final = tuple(_responses_allowed_tool(tool) for tool in parsed.allowed_tools.tools)
-    return frozenset(tool.get("name") for tool in flattened if tool is not None)
+    return frozenset(tool.get("name") for tool in flattened if tool is not None and _server_entry(tool) is None)
 
 
 def refuse_a_malformed_allowed_tools_tool_choice(tool_choice: object, model: str) -> None:
@@ -191,15 +191,22 @@ def refuse_a_malformed_allowed_tools_tool_choice(tool_choice: object, model: str
         )
 
 
-def _whole_server_label(tool: Mapping[str, object]) -> str | None:
+def _server_entry(tool: Mapping[str, object]) -> tuple[str, str | None] | None:
     label: Final = tool.get("server_label")
-    return label if tool.get("type") == "mcp" and tool.get("name") is None and isinstance(label, str) else None
+    name: Final = tool.get("name")
+    if tool.get("type") != "mcp" or not isinstance(label, str):
+        return None
+    return label, name if isinstance(name, str) else None
 
 
-def _whole_server_labels(tool_choice: object) -> frozenset[str]:
+def _server_entries(tool_choice: object) -> frozenset[tuple[str, str | None]]:
     parsed: Final = parse_allowed_tools_tool_choice(tool_choice)
-    labels: Final = () if parsed is None else tuple(map(_whole_server_label, parsed.allowed_tools.tools))
-    return frozenset(label for label in labels if label is not None)
+    entries: Final = () if parsed is None else tuple(map(_server_entry, parsed.allowed_tools.tools))
+    return frozenset(entry for entry in entries if entry is not None)
+
+
+def _a_server_entry_allows(entries: frozenset[tuple[str, str | None]], name: str, labels: frozenset[str]) -> bool:
+    return any(label in labels and entry_name in (None, name) for label, entry_name in entries)
 
 
 def tool_calls_the_choice_excludes(
@@ -210,11 +217,12 @@ def tool_calls_the_choice_excludes(
     if not is_allowed_tools_tool_choice(tool_choice):
         return ()
     allowed: Final = _allowed_tool_names(tool_choice)
-    whole_servers: Final = _whole_server_labels(tool_choice)
+    server_entries: Final = _server_entries(tool_choice)
     return tuple(
         name
         for name in tool_names
-        if name not in allowed and whole_servers.isdisjoint(server_labels_by_tool.get(name, frozenset()))
+        if name not in allowed
+        and not _a_server_entry_allows(server_entries, name, server_labels_by_tool.get(name, frozenset()))
     )
 
 
@@ -232,6 +240,38 @@ def refuse_tool_calls_the_choice_excludes(tool_choice: object, tool_names: Itera
         )
 
 
+def _with_builtin_as(
+    tool: Mapping[str, object], is_builtin: Callable[[Mapping[str, object]], bool], function_entry: Mapping[str, object]
+) -> Mapping[str, object]:
+    return function_entry if tool.get("type") not in ("function", "custom", "mcp") and is_builtin(tool) else tool
+
+
+def allowed_tools_tool_choice_with_builtin_as_function(
+    tool_choice: object, is_builtin: Callable[[Mapping[str, object]], bool], function_name: str
+) -> object:
+    try:
+        parsed: Final = _ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice)
+    except ValidationError:
+        return tool_choice
+    match parsed:
+        case ChatAllowedToolsToolChoice():
+            nested: Final = {"type": "function", "function": {"name": function_name}}
+            return {
+                "type": "allowed_tools",
+                "allowed_tools": {
+                    "mode": parsed.allowed_tools.mode,
+                    "tools": [dict(_with_builtin_as(tool, is_builtin, nested)) for tool in parsed.allowed_tools.tools],
+                },
+            }
+        case ResponsesAllowedToolsToolChoice():
+            flat: Final = {"type": "function", "name": function_name}
+            return {
+                "type": "allowed_tools",
+                "mode": parsed.mode,
+                "tools": [dict(_with_builtin_as(tool, is_builtin, flat)) for tool in parsed.tools],
+            }
+
+
 def chat_allowed_tools_tool_choice(choice: ChatAllowedToolsToolChoice) -> ChatCompletionAllowedToolChoiceParam:
     return ChatCompletionAllowedToolChoiceParam(
         type="allowed_tools",

tests/unit/integrations/code_interpreter_interception/test_handler.py

@@ -19,6 +19,7 @@ from litellm.integrations.code_interpreter_interception.handler import (
     LITELLM_CODE_EXECUTION_TOOL_NAME,
     CodeInterpreterInterceptionLogger,
 )
+from litellm.litellm_core_utils.allowed_tools_tool_choice import tool_calls_the_choice_excludes
 from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult
 from litellm.types.integrations.custom_logger import (
     CHAT_COMPLETION_AGENTIC_SURFACE,
@@ -1200,6 +1201,38 @@ async def test_extract_tool_calls_reads_object_attributes():
 # ---------------------------------------------------------------------------
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("call_type", "tool_choice"),
+    [
+        (CallTypes.aresponses, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "code_interpreter"}]}),
+        (
+            CallTypes.acompletion,
+            {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "code_interpreter"}]}},
+        ),
+    ],
+)
+async def test_pre_call_lets_an_allowed_tools_choice_listing_code_interpreter_allow_the_run(
+    call_type: CallTypes, tool_choice: dict[str, object]
+):
+    logger: Final = CodeInterpreterInterceptionLogger()
+
+    result: Final = TypeAdapter(dict[str, object]).validate_python(
+        await logger.async_pre_call_deployment_hook(
+            {
+                "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}],
+                "tool_choice": tool_choice,
+                "custom_llm_provider": "openai",
+            },
+            call_type,
+        )
+    )
+
+    assert tool_calls_the_choice_excludes(result["tool_choice"], (LITELLM_CODE_EXECUTION_TOOL_NAME, "get_weather")) == (
+        "get_weather",
+    )
+
+
 @pytest.mark.asyncio
 async def test_pre_call_uses_session_id_from_metadata_as_sandbox_key():
     """When session_id is in request metadata, it becomes the sandbox key so the

tests/unit/integrations/websearch_interception/test_websearch_responses.py

@@ -279,6 +279,39 @@ async def test_deployment_hook_converts_native_responses_web_search_tool():
     assert tool["parameters"]["required"] == ["query"]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("call_type", "tool_choice"),
+    [
+        (CallTypes.aresponses, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "web_search"}]}),
+        (CallTypes.acompletion, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "web_search"}]}),
+        (
+            CallTypes.acompletion,
+            {"type": "allowed_tools", "allowed_tools": {"mode": "required", "tools": [{"type": "web_search"}]}},
+        ),
+    ],
+)
+async def test_deployment_hook_lets_an_allowed_tools_choice_listing_web_search_allow_the_search(
+    call_type: CallTypes, tool_choice: dict[str, object]
+):
+    logger: Final = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI])
+
+    result: Final = await logger.async_pre_call_deployment_hook(
+        kwargs={
+            "model": "gpt-4o",
+            "custom_llm_provider": "openai",
+            "tools": [{"type": "web_search"}, {"type": "function", "name": "get_weather"}],
+            "tool_choice": tool_choice,
+        },
+        call_type=call_type,
+    )
+
+    assert result is not None
+    assert tool_calls_the_choice_excludes(result["tool_choice"], ("litellm_web_search", "get_weather")) == (
+        "get_weather",
+    )
+
+
 @pytest.mark.asyncio
 async def test_deployment_hook_responses_returns_none_without_web_search():
     """No web search tool in a responses request -> deployment hook makes no change."""

tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py

@@ -8,6 +8,7 @@ import litellm
 from litellm.exceptions import BadRequestError
 from litellm.litellm_core_utils.allowed_tools_tool_choice import (
     ChatAllowedToolsToolChoice,
+    allowed_tools_tool_choice_with_builtin_as_function,
     drops_unsupported_allowed_tools_tool_choice,
     is_allowed_tools_tool_choice,
     chat_allowed_tools_tool_choice,
@@ -297,6 +298,17 @@ def test_an_mcp_entry_with_a_null_name_is_a_valid_choice(tool_choice: object):
     assert refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") is None
 
 
+_DOCS_SEARCH: Final = {
+    "type": "allowed_tools",
+    "mode": "auto",
+    "tools": [{"type": "mcp", "server_label": "docs", "name": "search"}],
+}
+
+
+def test_a_named_mcp_entry_allows_its_call_on_a_server_its_label_serves():
+    assert tool_calls_the_choice_excludes(_DOCS_SEARCH, ("search",), {"search": frozenset({"docs", "other"})}) == ()
+
+
 def test_an_mcp_entry_without_a_name_allows_every_call_its_server_serves():
     assert tool_calls_the_choice_excludes(_WHOLE_DOCS_SERVER, ("search",), {"search": frozenset({"docs"})}) == ()
 
@@ -310,9 +322,59 @@ def test_an_mcp_entry_without_a_name_allows_every_call_its_server_serves():
             {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "docs", "name": "read"}]},
             frozenset({"docs"}),
         ),
+        (_DOCS_SEARCH, frozenset({"other"})),
+        (_DOCS_SEARCH, frozenset()),
     ],
 )
 def test_an_mcp_entry_allows_no_call_outside_its_server_or_its_name(
     tool_choice: object, server_labels: frozenset[str]
 ):
     assert tool_calls_the_choice_excludes(tool_choice, ("search",), {"search": server_labels}) == ("search",)
+
+
+def _is_web_search(tool: Mapping[str, object]) -> bool:
+    return tool.get("type") == "web_search"
+
+
+@pytest.mark.parametrize(
+    "tool_choice",
+    [
+        {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "web_search"}, {"type": "function", "name": "get_time"}]},
+        {
+            "type": "allowed_tools",
+            "allowed_tools": {
+                "mode": "auto",
+                "tools": [{"type": "web_search"}, {"type": "function", "function": {"name": "get_time"}}],
+            },
+        },
+    ],
+)
+def test_a_builtin_entry_allows_the_function_that_replaces_it_in_the_same_layout(tool_choice: dict[str, object]):
+    rewritten: Final = allowed_tools_tool_choice_with_builtin_as_function(
+        tool_choice, _is_web_search, "litellm_web_search"
+    )
+
+    assert tool_calls_the_choice_excludes(tool_choice, ("litellm_web_search", "get_time")) == ("litellm_web_search",)
+    assert tool_calls_the_choice_excludes(rewritten, ("litellm_web_search", "get_time", "get_weather")) == (
+        "get_weather",
+    )
+    assert isinstance(rewritten, dict) and rewritten.keys() == tool_choice.keys()
+
+
+def test_a_function_or_mcp_entry_is_never_taken_for_a_builtin():
+    tool_choice: Final = {
+        "type": "allowed_tools",
+        "mode": "auto",
+        "tools": [{"type": "function", "name": "web_search"}, {"type": "mcp", "server_label": "docs", "name": "web_search"}],
+    }
+
+    rewritten: Final = allowed_tools_tool_choice_with_builtin_as_function(
+        tool_choice, lambda tool: tool.get("name") == "web_search", "litellm_web_search"
+    )
+
+    assert parse_allowed_tools_tool_choice(rewritten) == parse_allowed_tools_tool_choice(tool_choice)
+
+
+@pytest.mark.parametrize("tool_choice", ["auto", None, {"type": "allowed_tools", "mode": "auto", "tools": []}])
+def test_any_other_or_malformed_choice_is_left_as_sent(tool_choice: object):
+    assert allowed_tools_tool_choice_with_builtin_as_function(tool_choice, _is_web_search, "litellm_web_search") is tool_choice