scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(tool-choice): resolve mcp entries like the gateway and keep the choice on every hook path

Browse source 
Under drop_params a malformed allowed_tools choice now acts as if it was never sent, and only a choice that parses is stored and enforced. Without drop_params it answers 400 before any provider request, on /v1/responses with gateway tools before MCP listing

An mcp entry resolves its label the way the gateway picks servers to list (server name in any case, alias, access group, toolset), and a named entry matches the tool's own name or the gateway's prefixed name. The gateway rewrites mcp entries into function entries in the caller's layout and sends "none" when nothing is left

The web search follow-up keeps an allowed_tools choice, compression's litellm_content_retrieve is never excluded and joins the provider-facing choice, and the native /v1/responses hook loop applies the exclusion check and stores the choice
This commit is contained in:
mateo-berri 2026-11-06 14:45:40 -08:00
parent cd076fb13a
commit 6be4f0b9ea
26 changed files with 1303 additions and 184 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/integrations/compression_interception/handler.py
View file

@ -15,6 +15,7 @@ from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.compression import compress
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.allowed_tools_tool_choice import allowed_tools_tool_choice_with_function
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.types.integrations.compression_interception import (
CompressionInterceptionConfig,
@ -180,6 +181,10 @@ class CompressionInterceptionLogger(CustomLogger):
existing_tools=cast(list[dict[str, object]] | None, kwargs.get("tools")),
compressed_tools=compressed_tools,
)
if "tool_choice" in kwargs:
kwargs["tool_choice"] = allowed_tools_tool_choice_with_function(
kwargs["tool_choice"], LITELLM_CONTENT_RETRIEVE_TOOL_NAME
)
call_id = cast(str | None, kwargs.get("litellm_call_id"))
if not call_id:
call_id = str(uuid.uuid4())

46
litellm/integrations/websearch_interception/handler.py
View file

@ -14,6 +14,7 @@ from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast
from pydantic import TypeAdapter
from typing_extensions import Never, ReadOnly
import litellm
@ -38,7 +39,10 @@ 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.litellm_core_utils.allowed_tools_tool_choice import (
allowed_tools_tool_choice_with_builtin_as_function,
is_allowed_tools_tool_choice,
)
from litellm.litellm_core_utils.core_helpers import qualify_provider_stripped_model
from litellm.llms.base_llm.search.transformation import SearchResponse
from litellm.types.integrations.custom_logger import (
@ -101,6 +105,8 @@ WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks"
_RESPONSE_CONTENT_FIELD: Final = "content"
_OPTIONAL_PARAMS: Final = TypeAdapter(Mapping[str, object])
_CLAUDE_CODE_WEB_SEARCH_PROMPT_PREFIX: Final = "Perform a web search for the query:"
_ResponseT: Final = TypeVar("_ResponseT")
@ -276,6 +282,11 @@ def _tool_choice_naming_litellm_web_search(
}
def _kept_allowed_tools_tool_choice(optional_params: Mapping[str, object]) -> Mapping[str, object]:
tool_choice: Final = optional_params.get("tool_choice")
return {"tool_choice": tool_choice} if is_allowed_tools_tool_choice(tool_choice) else MappingProxyType({})
class WebSearchInterceptionLogger(CustomLogger):
"""
CustomLogger that intercepts WebSearch tool calls for models that don't
@ -1863,9 +1874,11 @@ class WebSearchInterceptionLogger(CustomLogger):
)
if request_patch.messages is None:
raise ValueError("WebSearchInterception: missing follow-up messages")
params: Final = dict(optional_params)
params.update(request_patch.optional_params)
params.pop("tool_choice", None)
merged_params: Final = {**optional_params, **request_patch.optional_params}
params: Final = {
**{key: value for key, value in merged_params.items() if key != "tool_choice"},
**_kept_allowed_tools_tool_choice(merged_params),
}
patch_kwargs: Final = dict[str, object](request_patch.kwargs)
return await litellm.acompletion(
model=request_patch.model or model,
@ -1978,17 +1991,20 @@ class WebSearchInterceptionLogger(CustomLogger):
tools_param: Final = optional_params.get("tools")
optional_params_clean: Final = {
k: v
for k, v in optional_params.items()
if k
not in {
"tools",
"tool_choice",
"extra_body",
"model_alias_map",
"stream_response",
"custom_prompt_dict",
}
**{
k: v
for k, v in optional_params.items()
if k
not in {
"tools",
"tool_choice",
"extra_body",
"model_alias_map",
"stream_response",
"custom_prompt_dict",
}
},
**_kept_allowed_tools_tool_choice(_OPTIONAL_PARAMS.validate_python(optional_params)),
}
if tools_param is not None:
optional_params_clean["tools"] = tools_param

10
litellm/litellm_core_utils/agentic_loop_refusal.py
View file

@ -40,6 +40,16 @@ async def every_claimed_tool_call(
return MappingProxyType({"tool_calls": (*_claimed_calls(refused_tool_calls), *chain.from_iterable(claims))})
_TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call", "tool_use"})
def called_tool_names(response: object) -> frozenset[str]:
items: Final = (*parts_of(message_field(response, "output")), *parts_of(message_field(response, "content")))
return _strings(
message_field(item, "name") for item in items if message_field(item, "type") in _TOOL_CALL_ITEM_TYPES
)
def _call_ids(call: object) -> frozenset[str]:
return _strings((message_field(call, "id"), message_field(call, "call_id")))

151
litellm/litellm_core_utils/allowed_tools_tool_choice.py
View file

@ -1,4 +1,5 @@
from collections.abc import Callable, Iterable, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal
@ -179,16 +180,26 @@ def _allowed_tool_names(tool_choice: object) -> frozenset[object]:
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:
if is_allowed_tools_tool_choice(tool_choice) and parse_allowed_tools_tool_choice(tool_choice) is None:
raise BadRequestError(
message=(
f"Invalid tool choice, tool_choice={tool_choice}. An allowed_tools tool_choice must be "
f"{CHAT_ALLOWED_TOOLS_LAYOUT} or {RESPONSES_ALLOWED_TOOLS_LAYOUT}"
),
model=model,
llm_provider="",
)
def parse_allowed_tools_tool_choice_or_refuse(
tool_choice: object, model: str, drop_params: bool
) -> ChatAllowedToolsToolChoice | None:
parsed: Final = parse_allowed_tools_tool_choice(tool_choice)
if parsed is not None or drop_params:
return parsed
raise BadRequestError(
message=(
f"Invalid tool choice, tool_choice={tool_choice}. An allowed_tools tool_choice must be "
f"{CHAT_ALLOWED_TOOLS_LAYOUT} or {RESPONSES_ALLOWED_TOOLS_LAYOUT}"
),
model=model,
llm_provider="",
)
def allowed_tools_tool_choice_in_effect(tool_choice: object, model: str, drop_params: bool) -> object:
if not is_allowed_tools_tool_choice(tool_choice):
return tool_choice
return None if parse_allowed_tools_tool_choice_or_refuse(tool_choice, model, drop_params) is None else tool_choice
def _server_entry(tool: Mapping[str, object]) -> tuple[str, str | None] | None:
@ -205,14 +216,32 @@ def _server_entries(tool_choice: object) -> frozenset[tuple[str, str | None]]:
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)
@dataclass(frozen=True, slots=True)
class GatewayTool:
names: frozenset[str]
server_labels: frozenset[str]
@dataclass(frozen=True, slots=True)
class GatewayTools:
server_labels: frozenset[str]
tools: Mapping[str, GatewayTool]
NO_GATEWAY_TOOLS: Final = GatewayTools(server_labels=frozenset(), tools=MappingProxyType({}))
def _server_entry_allows(entry: tuple[str, str | None], tool: GatewayTool | None) -> bool:
label, name = entry
return tool is not None and label in tool.server_labels and (name is None or name in tool.names)
def _a_server_entry_allows(entries: Iterable[tuple[str, str | None]], tool: GatewayTool | None) -> bool:
return any(_server_entry_allows(entry, tool) for entry in entries)
def tool_calls_the_choice_excludes(
tool_choice: object,
tool_names: Iterable[str],
server_labels_by_tool: Mapping[str, frozenset[str]] = MappingProxyType({}),
tool_choice: object, tool_names: Iterable[str], gateway: GatewayTools = NO_GATEWAY_TOOLS
) -> tuple[str, ...]:
if not is_allowed_tools_tool_choice(tool_choice):
return ()
@ -221,8 +250,7 @@ def tool_calls_the_choice_excludes(
return tuple(
name
for name in tool_names
if name not in allowed
and not _a_server_entry_allows(server_entries, name, server_labels_by_tool.get(name, frozenset()))
if name not in allowed and not _a_server_entry_allows(server_entries, gateway.tools.get(name))
)
@ -246,29 +274,80 @@ def _with_builtin_as(
return function_entry if tool.get("type") not in ("function", "custom", "mcp") and is_builtin(tool) else tool
def _function_entry(
layout: ChatAllowedToolsToolChoice | ResponsesAllowedToolsToolChoice, name: str
) -> Mapping[str, object]:
match layout:
case ChatAllowedToolsToolChoice():
return {"type": "function", "function": {"name": name}}
case ResponsesAllowedToolsToolChoice():
return {"type": "function", "name": name}
def _entries_of(
layout: ChatAllowedToolsToolChoice | ResponsesAllowedToolsToolChoice,
) -> tuple[Mapping[str, object], ...]:
match layout:
case ChatAllowedToolsToolChoice():
return layout.allowed_tools.tools
case ResponsesAllowedToolsToolChoice():
return layout.tools
def _with_entries(
layout: ChatAllowedToolsToolChoice | ResponsesAllowedToolsToolChoice, tools: Iterable[Mapping[str, object]]
) -> Mapping[str, object]:
entries: Final = [dict(tool) for tool in tools]
match layout:
case ChatAllowedToolsToolChoice():
return {"type": "allowed_tools", "allowed_tools": {"mode": layout.allowed_tools.mode, "tools": entries}}
case ResponsesAllowedToolsToolChoice():
return {"type": "allowed_tools", "mode": layout.mode, "tools": entries}
def _callers_layout(tool_choice: object) -> ChatAllowedToolsToolChoice | ResponsesAllowedToolsToolChoice | None:
if parse_allowed_tools_tool_choice(tool_choice) is None:
return None
return _ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice)
def allowed_tools_tool_choice_with_builtin_as_function(
tool_choice: object, is_builtin: Callable[[Mapping[str, object]], bool], function_name: str
) -> object:
if parse_allowed_tools_tool_choice(tool_choice) is None:
layout: Final = _callers_layout(tool_choice)
if layout is None:
return tool_choice
parsed: Final = _ALLOWED_TOOLS_TOOL_CHOICE.validate_python(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],
}
entry: Final = _function_entry(layout, function_name)
return _with_entries(layout, (_with_builtin_as(tool, is_builtin, entry) for tool in _entries_of(layout)))
def allowed_tools_tool_choice_with_function(tool_choice: object, function_name: str) -> object:
layout: Final = _callers_layout(tool_choice)
if layout is None or function_name in _allowed_tool_names(tool_choice):
return tool_choice
return _with_entries(layout, (*_entries_of(layout), _function_entry(layout, function_name)))
def _gateway_entry(tool: Mapping[str, object], gateway_labels: frozenset[str]) -> tuple[str, str | None] | None:
entry: Final = _server_entry(tool)
return entry if entry is not None and entry[0] in gateway_labels else None
def allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice: object, gateway: GatewayTools) -> object:
layout: Final = _callers_layout(tool_choice)
entries: Final = () if layout is None else _entries_of(layout)
gateway_entries: Final = tuple(
entry for entry in (_gateway_entry(tool, gateway.server_labels) for tool in entries) if entry is not None
)
if layout is None or not gateway_entries:
return tool_choice
kept: Final = tuple(tool for tool in entries if _gateway_entry(tool, gateway.server_labels) is None)
allowed: Final = tuple(
name for name, tool in gateway.tools.items() if _a_server_entry_allows(gateway_entries, tool)
)
if not kept and not allowed:
return "none"
return _with_entries(layout, (*kept, *(_function_entry(layout, name) for name in allowed)))
def chat_allowed_tools_tool_choice(choice: ChatAllowedToolsToolChoice) -> ChatCompletionAllowedToolChoiceParam:

27
litellm/litellm_core_utils/chat_completion_agentic_loop.py
View file

@ -1,7 +1,7 @@
# this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api
import json
from collections.abc import Awaitable, Iterator, Mapping, Sequence
from collections.abc import Awaitable, Iterable, Iterator, Mapping, Sequence
from itertools import chain
from types import MappingProxyType
from typing import Final, TypeVar
@ -19,7 +19,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import (
validated_max_agentic_loops,
)
from litellm.litellm_core_utils.allowed_tools_tool_choice import (
is_allowed_tools_tool_choice,
parse_allowed_tools_tool_choice,
tool_calls_the_choice_excludes,
)
from litellm.litellm_core_utils.core_helpers import qualify_agentic_followup_model
@ -203,15 +203,23 @@ def _called_tool_names(response: ModelResponse) -> tuple[str, ...]:
return tuple(map(_tool_call_name, calls))
def _holds_a_call_the_choice_excludes(response: ModelResponse, kwargs: Mapping[str, object]) -> bool:
return bool(tool_calls_the_choice_excludes(kwargs.get(AGENTIC_LOOP_TOOL_CHOICE_KEY), _called_tool_names(response)))
def _server_fulfilled_tool_names(callbacks: Iterable[CustomLogger]) -> frozenset[str]:
return frozenset(chain.from_iterable(callback.server_fulfilled_tool_names for callback in callbacks))
def _with_caller_tool_choice(
def holds_a_call_the_choice_excludes(
tool_names: Iterable[str], kwargs: Mapping[str, object], callbacks: Iterable[CustomLogger]
) -> bool:
server_fulfilled: Final = _server_fulfilled_tool_names(callbacks)
client_tool_names: Final = (name for name in tool_names if name not in server_fulfilled)
return bool(tool_calls_the_choice_excludes(kwargs.get(AGENTIC_LOOP_TOOL_CHOICE_KEY), client_tool_names))
def with_caller_tool_choice(
kwargs: Mapping[str, object], optional_params: Mapping[str, object]
) -> Mapping[str, object]:
tool_choice: Final = optional_params.get("tool_choice")
if AGENTIC_LOOP_TOOL_CHOICE_KEY in kwargs or not is_allowed_tools_tool_choice(tool_choice):
if AGENTIC_LOOP_TOOL_CHOICE_KEY in kwargs or parse_allowed_tools_tool_choice(tool_choice) is None:
return kwargs
return {**kwargs, AGENTIC_LOOP_TOOL_CHOICE_KEY: tool_choice}
@ -277,6 +285,7 @@ def _fingerprint_of_a_turn_to_run(
*,
response: ModelResponse,
request_kwargs: Mapping[str, object],
callbacks: Sequence[CustomLogger],
tool_calls: object,
fingerprints: Sequence[str],
depth: int,
@ -284,7 +293,7 @@ def _fingerprint_of_a_turn_to_run(
model: str,
stream: bool,
) -> str | None:
if _holds_a_call_the_choice_excludes(response, request_kwargs):
if holds_a_call_the_choice_excludes(_called_tool_names(response), request_kwargs, callbacks):
return None
try:
return _check_agentic_loop_safety(
@ -434,7 +443,7 @@ async def maybe_run_chat_completion_agentic_loop(
stream: bool,
) -> ModelResponse | CustomStreamWrapper | None:
callbacks: Final = _custom_loggers(logging_obj)
request_kwargs: Final = _with_caller_tool_choice(kwargs, optional_params)
request_kwargs: Final = with_caller_tool_choice(kwargs, optional_params)
depth, max_loops, fingerprints = _agentic_loop_settings(request_kwargs)
tools: Final = _request_tools(optional_params)
@ -470,6 +479,7 @@ async def maybe_run_chat_completion_agentic_loop(
fingerprint = _fingerprint_of_a_turn_to_run(
response=response,
request_kwargs=request_kwargs,
callbacks=callbacks,
tool_calls=tool_calls,
fingerprints=fingerprints,
depth=depth,
@ -592,6 +602,7 @@ async def run_chat_completion_agentic_hooks(
fingerprint = _fingerprint_of_a_turn_to_run(
response=response,
request_kwargs=kwargs,
callbacks=callbacks,
tool_calls=tool_calls,
fingerprints=fingerprints,
depth=depth,

16
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -48,6 +48,7 @@ from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAG
from litellm.files.types import FileContentStreamingResult
from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs
from litellm.litellm_core_utils.agentic_loop_refusal import (
called_tool_names,
every_claimed_tool_call,
finalize_refused_responses_turn,
is_refused_item,
@ -64,7 +65,9 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import (
)
from litellm.litellm_core_utils.chat_completion_agentic_loop import (
bill_first_round_on_failure,
holds_a_call_the_choice_excludes,
run_chat_completion_agentic_hooks,
with_caller_tool_choice,
)
from litellm.litellm_core_utils.core_helpers import as_str_mapping, set_provider_response_headers_in_hidden_params
from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value
@ -160,6 +163,7 @@ from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
AgenticLoopRequestPatch,
AgenticLoopSafetyError,
AgenticLoopToolChoiceError,
converted_stream_requested,
is_interception_internal_key,
)
@ -6383,7 +6387,13 @@ class BaseLLMHTTPHandler:
tools: Final = anthropic_messages_optional_request_params.get("tools", [])
depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs)
hook_kwargs: Final = {**kwargs, "_agentic_loop_api_surface": api_surface}
hook_kwargs: Final = {
**with_caller_tool_choice(
_OPTIONAL_PARAMS.validate_python(kwargs),
_OPTIONAL_PARAMS.validate_python(anthropic_messages_optional_request_params),
),
"_agentic_loop_api_surface": api_surface,
}
for callback in callbacks:
if not isinstance(callback, CustomLogger):
@ -6422,6 +6432,10 @@ class BaseLLMHTTPHandler:
# Safety guards must run OUTSIDE the callback try/except — they are
# bounded-loop / cycle-break rails, not callback bugs.
try:
if holds_a_call_the_choice_excludes(
called_tool_names(response), hook_kwargs, (cb for cb in callbacks if isinstance(cb, CustomLogger))
):
raise AgenticLoopToolChoiceError(model)
fingerprint = self._check_agentic_loop_safety(
tool_calls=tool_calls,
fingerprints=fingerprints,

7
litellm/main.py
View file

@ -74,7 +74,7 @@ from litellm.exceptions import LiteLLMUnknownProvider
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.allowed_tools_tool_choice import (
drops_unsupported_allowed_tools_tool_choice,
is_allowed_tools_tool_choice,
parse_allowed_tools_tool_choice,
)
from litellm.litellm_core_utils.asyncify import asyncify, run_async_function
from litellm.litellm_core_utils.audio_utils.utils import (
@ -5594,8 +5594,7 @@ def completion(
model=model,
drop_params=drops_unsupported_params(
cast("Mapping[str, object]", kwargs).get("drop_params") # cast-ok: kwargs is the untyped request dict
)
and not may_name_litellm_mcp_gateway(tools), # pyright: ignore[reportUnknownArgumentType] # untyped tools list
),
)
# validate optional params
stop = validate_openai_optional_params(stop=stop)
@ -6212,7 +6211,7 @@ def completion(
**(
{AGENTIC_LOOP_TOOL_CHOICE_KEY: requested_tool_choice}
if AGENTIC_LOOP_TOOL_CHOICE_KEY not in kwargs
and is_allowed_tools_tool_choice(requested_tool_choice)
and parse_allowed_tools_tool_choice(requested_tool_choice) is not None
else {}
),
},

24
litellm/responses/main.py
View file

@ -29,6 +29,10 @@ from litellm.integrations.anthropic_cache_control_hook import (
configured_injection_points,
routed_model_info,
)
from litellm.litellm_core_utils.allowed_tools_tool_choice import (
allowed_tools_tool_choice_in_effect,
allowed_tools_tool_choice_with_gateway_tools_as_functions,
)
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params
@ -312,6 +316,11 @@ async def aresponses_api_with_mcp(
LiteLLM_Proxy_MCP_Handler,
)
request_kwargs: Final = _STR_MAPPING.validate_python(kwargs)
tool_choice_in_effect: Final = allowed_tools_tool_choice_in_effect(
tool_choice, model, drops_unsupported_params(request_kwargs.get("drop_params"))
)
# Parse MCP tools and separate from other tools
(
mcp_tools_with_litellm_proxy,
@ -330,14 +339,13 @@ async def aresponses_api_with_mcp(
else (None, None, None, None)
)
request_kwargs: Final = _STR_MAPPING.validate_python(kwargs)
logging_obj: Final = request_kwargs.get("litellm_logging_obj")
joins_request_trace: Final = isinstance(logging_obj, LiteLLMLoggingObj)
bridged_first_turn: Final = (
joins_request_trace
and not stream
and previous_response_id is None
and not provider_stores_the_conversation(model, custom_llm_provider, kwargs)
and not provider_stores_the_conversation(model, custom_llm_provider, request_kwargs)
)
listing_trace_id: Final = (
@ -361,6 +369,9 @@ async def aresponses_api_with_mcp(
raw_headers=discovery_raw_headers,
)
openai_tools: Final = LiteLLM_Proxy_MCP_Handler.transform_mcp_tools_to_openai(original_mcp_tools)
gateway: Final = await LiteLLM_Proxy_MCP_Handler.gateway_tools(
tool_choice_in_effect, tool_server_map, mcp_tools_with_litellm_proxy, user_api_key_auth
)
# Combine with other tools
all_tools: Final = openai_tools + other_tools if (openai_tools or other_tools) else None
@ -379,7 +390,7 @@ async def aresponses_api_with_mcp(
"stream": stream,
"temperature": temperature,
"text": text,
"tool_choice": tool_choice,
"tool_choice": allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice_in_effect, gateway),
"top_p": top_p,
"truncation": truncation,
"user": user,
@ -418,6 +429,8 @@ async def aresponses_api_with_mcp(
previous_response_id=previous_response_id,
tool_server_map=tool_server_map,
served_tools=original_mcp_tools,
tool_choice=tool_choice_in_effect,
gateway=gateway,
**kwargs,
)
await mcp_streaming_response.start_initial_response()
@ -458,9 +471,8 @@ async def aresponses_api_with_mcp(
with first_round_billed_on_failure(logging_obj, response):
tool_calls: Final = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run(
LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response),
tool_choice,
tool_server_map,
mcp_tools_with_litellm_proxy,
tool_choice_in_effect,
gateway,
)
if tool_calls:

13
litellm/responses/mcp/chat_completions_handler.py
View file

@ -11,6 +11,9 @@ from pydantic import TypeAdapter
from typing_extensions import TypedDict, Unpack
from litellm._internal_context import internal_sub_call
from litellm.litellm_core_utils.allowed_tools_tool_choice import (
allowed_tools_tool_choice_with_gateway_tools_as_functions,
)
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.responses.mcp.chat_stream_round_billing import ChatStreamRoundBilling, engaged_round_billing
@ -163,6 +166,9 @@ async def acompletion_with_mcp(
deduplicated_mcp_tools,
target_format="chat",
)
gateway: Final = await LiteLLM_Proxy_MCP_Handler.gateway_tools(
kwargs.get("tool_choice"), tool_server_map, mcp_tools_with_litellm_proxy, user_api_key_auth
)
# Combine with other tools
all_tools: Final = openai_tools + other_tools if (openai_tools or other_tools) else None
@ -182,6 +188,7 @@ async def acompletion_with_mcp(
"tools": all_tools,
"_skip_mcp_handler": True, # Prevent recursion
**clean_kwargs,
"tool_choice": allowed_tools_tool_choice_with_gateway_tools_as_functions(kwargs.get("tool_choice"), gateway),
}
# If not auto-executing, just make the call with transformed tools
@ -494,8 +501,7 @@ async def acompletion_with_mcp(
self.tool_calls = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run(
LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response=complete_response),
kwargs.get("tool_choice"),
self.tool_server_map,
mcp_tools_with_litellm_proxy,
gateway,
)
if self.tool_calls:
@ -677,8 +683,7 @@ async def acompletion_with_mcp(
tool_calls: Final = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run(
LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response=initial_response),
kwargs.get("tool_choice"),
tool_server_map,
mcp_tools_with_litellm_proxy,
gateway,
)
if not tool_calls:

146
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -1,5 +1,6 @@
import traceback
from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload
@ -11,12 +12,20 @@ from typing_extensions import Unpack
from litellm._logging import verbose_logger
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
from litellm.litellm_core_utils.allowed_tools_tool_choice import tool_calls_the_choice_excludes
from litellm.litellm_core_utils.allowed_tools_tool_choice import (
NO_GATEWAY_TOOLS,
GatewayTool,
GatewayTools,
parse_allowed_tools_tool_choice,
tool_calls_the_choice_excludes,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.mcp_server.catalog import catalog_operation, global_manager
from litellm.proxy.mcp_server.utils import (
iter_known_server_prefixes,
iter_known_tool_name_spellings,
logging_safe_mcp_headers,
match_known_tool_name,
split_server_prefix_from_name,
strip_known_server_prefix,
)
@ -27,7 +36,6 @@ from litellm.responses.mcp.tool_precheck import (
names_gateway_explicitly,
names_whole_gateway,
proxy_path_mcp_name,
server_labels_serving,
)
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
from litellm.types.llms.openai import (
@ -120,6 +128,47 @@ async def _served_mcp_path_names(
return await served_names(names) if names else frozenset[str]()
def _served_tool(tool_name: str, server_name: str) -> tuple[MCPServer | None, str]:
from litellm.proxy.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy.mcp_server.server import (
_resolve_display_name_to_original,
)
mcp_server: Final = global_mcp_server_manager.get_mcp_server_by_name(
server_name
) or global_mcp_server_manager.get_mcp_server_from_tool_name(tool_name)
resolved_tool_name: Final = _resolve_display_name_to_original(tool_name, [mcp_server]) if mcp_server else tool_name
return mcp_server, strip_known_server_prefix(resolved_tool_name, mcp_server)
@dataclass(frozen=True, slots=True)
class _LabelScope:
label: str
whole_gateway: bool = False
server_ids: frozenset[str] = frozenset()
toolset_tools: Mapping[str, tuple[str, ...]] = MappingProxyType({})
def serves(self, server: MCPServer | None, tool_name: str) -> bool:
if self.whole_gateway:
return True
if server is None:
return False
toolset_names: Final = self.toolset_tools.get(server.server_id)
return server.server_id in self.server_ids or (
toolset_names is not None and match_known_tool_name(tool_name, server, toolset_names) is not None
)
def _gateway_tool(name: str, server: MCPServer | None, tool_name: str, scopes: Iterable[_LabelScope]) -> GatewayTool:
spellings: Final = () if server is None else tuple(iter_known_tool_name_spellings(tool_name, server))
return GatewayTool(
names=frozenset((name, tool_name, *spellings)),
server_labels=frozenset(scope.label for scope in scopes if scope.serves(server, tool_name)),
)
class LiteLLM_Proxy_MCP_Handler:
"""
Helper class with static methods for MCP integration with Responses API.
@ -612,20 +661,80 @@ class LiteLLM_Proxy_MCP_Handler:
_should_auto_execute_tools = should_auto_execute_tools
@staticmethod
def tool_calls_to_run(
tool_calls: Sequence[object],
async def _label_scope(
tool: Mapping[str, object],
servers: Sequence[MCPServer],
user_api_key_auth: "UserAPIKeyAuth | None",
granted_toolsets: "GrantedToolsetIds | None",
) -> _LabelScope | None:
from litellm.proxy.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy.mcp_server.server import (
_get_allowed_mcp_servers_from_mcp_server_names,
)
label: Final = tool.get("server_label")
server_url: Final = tool.get("server_url")
if not isinstance(label, str) or not isinstance(server_url, str):
return None
if names_whole_gateway(tool):
return _LabelScope(label=label, whole_gateway=True)
kind, value = await LiteLLM_Proxy_MCP_Handler._resolve_proxy_mcp_name(
server_url.split("/")[-1], user_api_key_auth, granted_toolsets
)
match kind:
case "server":
resolved: Final = await _get_allowed_mcp_servers_from_mcp_server_names([value], list(servers))
return _LabelScope(label=label, server_ids=frozenset(server.server_id for server in resolved))
case "toolset":
permissions: Final = await global_mcp_server_manager.resolve_toolset_tool_permissions(
toolset_ids=[value]
)
return _LabelScope(
label=label,
toolset_tools=MappingProxyType({key: tuple(names) for key, names in permissions.items()}),
)
case "denied_toolset":
return _LabelScope(label=label)
@staticmethod
async def gateway_tools(
tool_choice: object,
tool_server_map: Mapping[str, str],
mcp_tools: Iterable[object],
) -> Sequence[object]:
mcp_tools_with_litellm_proxy: Iterable[Mapping[str, object]],
user_api_key_auth: "UserAPIKeyAuth | None",
granted_toolsets: "GrantedToolsetIds | None" = None,
) -> GatewayTools:
if parse_allowed_tools_tool_choice(tool_choice) is None:
return NO_GATEWAY_TOOLS
served: Final = {name: _served_tool(name, server_name) for name, server_name in tool_server_map.items()}
servers: Final = {server.server_id: server for server, _ in served.values() if server is not None}
scopes: Final = tuple(
[
await LiteLLM_Proxy_MCP_Handler._label_scope(
tool, tuple(servers.values()), user_api_key_auth, granted_toolsets
)
for tool in mcp_tools_with_litellm_proxy
]
)
label_scopes: Final = tuple(scope for scope in scopes if scope is not None)
return GatewayTools(
server_labels=frozenset(scope.label for scope in label_scopes),
tools=MappingProxyType(
{
name: _gateway_tool(name, server, tool_name, label_scopes)
for name, (server, tool_name) in served.items()
}
),
)
@staticmethod
def tool_calls_to_run(tool_calls: Sequence[object], tool_choice: object, gateway: GatewayTools) -> Sequence[object]:
names: Final = tuple(
LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call)[0] or "" for tool_call in tool_calls
)
references: Final = tuple(mcp_tools)
labels_by_tool: Final = MappingProxyType(
{name: server_labels_serving(references, tool_server_map.get(name)) for name in names}
)
return () if tool_calls_the_choice_excludes(tool_choice, names, labels_by_tool) else tool_calls
return () if tool_calls_the_choice_excludes(tool_choice, names, gateway) else tool_calls
@staticmethod
def extract_tool_calls_from_response(response: ResponsesAPIResponse) -> list[object]:
@ -768,9 +877,6 @@ class LiteLLM_Proxy_MCP_Handler:
from litellm.proxy.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy.mcp_server.server import (
_resolve_display_name_to_original,
)
from litellm.proxy.proxy_server import proxy_logging_obj
tool_results: Final[list[MCPToolResult]] = []
@ -799,13 +905,7 @@ class LiteLLM_Proxy_MCP_Handler:
server_name = tool_server_map[tool_name]
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(
server_name
) or global_mcp_server_manager.get_mcp_server_from_tool_name(tool_name)
resolved_tool_name = (
_resolve_display_name_to_original(tool_name, [mcp_server]) if mcp_server else tool_name
)
sanitized_tool_name = strip_known_server_prefix(resolved_tool_name, mcp_server)
mcp_server, sanitized_tool_name = _served_tool(tool_name, server_name)
start_time = datetime.now()
logging_input: Sequence[Mapping[str, object]] = [
@ -1235,6 +1335,8 @@ class LiteLLM_Proxy_MCP_Handler:
previous_response_id: str | None,
tool_server_map: dict[str, str],
served_tools: Sequence[MCPTool] | None = None,
tool_choice: object = None,
gateway: GatewayTools = NO_GATEWAY_TOOLS,
**kwargs,
) -> "MCPEnhancedStreamingIterator":
"""
@ -1269,6 +1371,8 @@ class LiteLLM_Proxy_MCP_Handler:
user_api_key_auth=kwargs.get("user_api_key_auth")
or kwargs.get("litellm_metadata", {}).get("user_api_key_auth"),
original_request_params=request_params,
tool_choice=tool_choice,
gateway=gateway,
)
@staticmethod

10
litellm/responses/mcp/mcp_streaming_iterator.py
View file

@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final, cast
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.agentic_loop_settings import validated_max_agentic_loops
from litellm.litellm_core_utils.allowed_tools_tool_choice import NO_GATEWAY_TOOLS, GatewayTools
from litellm.responses.litellm_completion_transformation.streaming_iterator import LiteLLMCompletionStreamingIterator
from litellm.responses.mcp.request_context import MCPRequestContext
from litellm.responses.mcp.stream_round_billing import StreamedRoundBilling
@ -286,6 +287,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
user_api_key_auth: "UserAPIKeyAuth | None" = None,
original_request_params: dict[str, Any] | None = None,
served_tools: Sequence[MCPTool] | None = None,
tool_choice: object = None,
gateway: GatewayTools = NO_GATEWAY_TOOLS,
):
# MCP setup
self.mcp_tools_with_litellm_proxy = mcp_tools_with_litellm_proxy or []
@ -310,6 +313,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
self.mcp_discovery_generated = True # Events are already generated
self.mcp_events = mcp_events # Store the initial MCP events for backward compatibility
self.tool_server_map = tool_server_map
self.tool_choice: Final = tool_choice
self.gateway: Final = gateway
self.served_tools = tuple(served_tools) if served_tools is not None else None
# Iterator references
@ -841,9 +846,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
tool_calls: Final = (
LiteLLM_Proxy_MCP_Handler.tool_calls_to_run(
LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(self.collected_response),
self.original_request_params.get("tool_choice"),
self.tool_server_map,
self.mcp_tools_with_litellm_proxy,
self.tool_choice,
self.gateway,
)
if self.collected_response is not None
else ()

15
litellm/responses/mcp/tool_precheck.py
View file

@ -29,21 +29,6 @@ def names_whole_gateway(tool: object) -> bool:
)
def _server_label_serving(tool: object, server_name: str) -> str | None:
match tool:
case {"type": "mcp", "server_label": str() as label} if (
names_whole_gateway(tool) or _mcp_server_url(tool) == f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{server_name}"
):
return label
case _:
return None
def server_labels_serving(mcp_tools: Iterable[object], server_name: str | None) -> frozenset[str]:
labels: Final = () if server_name is None else tuple(_server_label_serving(tool, server_name) for tool in mcp_tools)
return frozenset(label for label in labels if label is not None)
def proxy_path_mcp_name(tool: object) -> str | None:
server_url: Final = _mcp_server_url(tool)
match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url)

23
litellm/tool_loop.py
View file

@ -7,9 +7,10 @@ from typing_extensions import TypedDict, Unpack
from litellm.constants import DEFAULT_TOOL_LOOP_MAX_ROUNDS
from litellm.litellm_core_utils.allowed_tools_tool_choice import (
refuse_a_malformed_allowed_tools_tool_choice,
allowed_tools_tool_choice_in_effect,
refuse_tool_calls_the_choice_excludes,
)
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantMessage,
@ -36,12 +37,14 @@ class ToolLoopMaxRoundsExceeded(RuntimeError):
super().__init__(f"model still requested tool calls on round {max_rounds} of {max_rounds}")
def _validate_tool_loop_args(max_rounds: int, model: str, completion_kwargs: _ToolLoopCompletionKwargs) -> None:
def _tool_choice_in_effect(max_rounds: int, model: str, completion_kwargs: _ToolLoopCompletionKwargs) -> object:
if max_rounds < 1:
raise ValueError(f"max_rounds must be >= 1, got {max_rounds}")
if completion_kwargs.get("stream"):
raise ValueError("run_tool_loop requires whole responses; stream=True is not supported")
refuse_a_malformed_allowed_tools_tool_choice(completion_kwargs.get("tool_choice"), model)
return allowed_tools_tool_choice_in_effect(
completion_kwargs.get("tool_choice"), model, drops_unsupported_params(completion_kwargs.get("drop_params"))
)
def _expect_model_response(response: object) -> ModelResponse:
@ -120,12 +123,10 @@ def _client_tool_calls(
def _refuse_tool_calls_the_choice_excludes(
tool_calls: tuple[ChatCompletionMessageToolCall, ...], model: str, completion_kwargs: _ToolLoopCompletionKwargs
tool_calls: tuple[ChatCompletionMessageToolCall, ...], model: str, tool_choice: object
) -> None:
refuse_tool_calls_the_choice_excludes(
completion_kwargs.get("tool_choice"),
tuple(tool_call.function.name or "" for tool_call in tool_calls),
model,
tool_choice, tuple(tool_call.function.name or "" for tool_call in tool_calls), model
)
@ -163,7 +164,7 @@ def run_tool_loop(
"""
import litellm
_validate_tool_loop_args(max_rounds, model, completion_kwargs)
tool_choice: Final = _tool_choice_in_effect(max_rounds, model, completion_kwargs)
history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns
for round_number in range(1, max_rounds + 1):
message = (
@ -181,7 +182,7 @@ def run_tool_loop(
return message.content
if round_number == max_rounds:
break
_refuse_tool_calls_the_choice_excludes(client_tool_calls, model, completion_kwargs)
_refuse_tool_calls_the_choice_excludes(client_tool_calls, model, tool_choice)
tool_results = tuple(execute_tool(tool_call) for tool_call in client_tool_calls)
history = (
*history,
@ -215,7 +216,7 @@ async def arun_tool_loop(
"""
import litellm
_validate_tool_loop_args(max_rounds, model, completion_kwargs)
tool_choice: Final = _tool_choice_in_effect(max_rounds, model, completion_kwargs)
history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns
for round_number in range(1, max_rounds + 1):
message = (
@ -233,7 +234,7 @@ async def arun_tool_loop(
return message.content
if round_number == max_rounds:
break
_refuse_tool_calls_the_choice_excludes(client_tool_calls, model, completion_kwargs)
_refuse_tool_calls_the_choice_excludes(client_tool_calls, model, tool_choice)
tool_results = tuple([await execute_tool(tool_call) for tool_call in client_tool_calls])
history = (
*history,

5
litellm/types/integrations/custom_logger.py
View file

@ -81,6 +81,11 @@ class AgenticLoopLimitError(AgenticLoopSafetyError):
pass
class AgenticLoopToolChoiceError(AgenticLoopSafetyError):
def __init__(self, model: str) -> None:
super().__init__(f"The allowed_tools tool_choice excludes a tool call the model made, model={model}")
class StandardCustomLoggerInitParams(LiteLLMBaseModel):
"""
Params for initializing a CustomLogger.

23
litellm/utils.py
View file

@ -85,10 +85,8 @@ from litellm.constants import (
)
from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure
from litellm.litellm_core_utils.allowed_tools_tool_choice import (
CHAT_ALLOWED_TOOLS_LAYOUT,
RESPONSES_ALLOWED_TOOLS_LAYOUT,
chat_allowed_tools_tool_choice,
parse_allowed_tools_tool_choice,
parse_allowed_tools_tool_choice_or_refuse,
)
from litellm.litellm_core_utils.core_helpers import (
bind_budget_reservation_to_callbacks,
@ -9669,23 +9667,14 @@ def validate_chat_completion_tool_choice(
if tool_choice_type in ("auto", "none", "required") and "function" not in tool_choice:
return tool_choice_type
if tool_choice_type == "allowed_tools":
allowed_tools: Final = parse_allowed_tools_tool_choice(
allowed_tools: Final = parse_allowed_tools_tool_choice_or_refuse(
cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic
object, tool_choice
)
),
model,
drop_params,
)
if allowed_tools is None and drop_params:
return None
if allowed_tools is None:
raise BadRequestError(
message=(
f"Invalid tool choice, tool_choice={tool_choice}. An allowed_tools tool_choice must be "
f"{CHAT_ALLOWED_TOOLS_LAYOUT} or {RESPONSES_ALLOWED_TOOLS_LAYOUT}"
),
model=model,
llm_provider="",
)
return dict(chat_allowed_tools_tool_choice(allowed_tools))
return None if allowed_tools is None else dict(chat_allowed_tools_tool_choice(allowed_tools))
# Standard OpenAI format: {"type": "function", "function": {...}}
if tool_choice_type is None or tool_choice.get("function") is None:

47
tests/unit/integrations/compression_interception/test_compression_interception_handler.py
View file

@ -548,3 +548,50 @@ async def test_pre_call_hook_counts_tokens_off_the_event_loop():
assert result["messages"] is messages
assert "tools" not in result
assert_loop_stayed_free(took, lags)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("tool_choice", "retrieval_entry"),
[
(
{"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "existing_tool"}]},
{"type": "function", "name": "litellm_content_retrieve"},
),
(
{
"type": "allowed_tools",
"allowed_tools": {"mode": "auto", "tools": [{"type": "function", "function": {"name": "existing_tool"}}]},
},
{"type": "function", "function": {"name": "litellm_content_retrieve"}},
),
],
)
async def test_pre_call_hook_lets_an_allowed_tools_choice_allow_the_retrieval_tool_it_adds(
monkeypatch, tool_choice, retrieval_entry
):
monkeypatch.setattr(
"litellm.integrations.compression_interception.handler.compress",
lambda **_: {
"messages": [{"role": "user", "content": "stubbed"}],
"original_tokens": 12000,
"compressed_tokens": 5000,
"cache": {"auth.py": "full file content"},
"tools": [{"type": "function", "function": {"name": "litellm_content_retrieve", "parameters": {}}}],
},
)
result = await CompressionInterceptionLogger().async_pre_call_deployment_hook(
kwargs={
"model": "anthropic/claude-sonnet-4-5",
"messages": [{"role": "user", "content": "very large context"}],
"tools": [{"type": "function", "function": {"name": "existing_tool", "parameters": {"type": "object"}}}],
"tool_choice": tool_choice,
},
call_type=CallTypes.anthropic_messages,
)
assert result is not None
entries = result["tool_choice"].get("tools") or result["tool_choice"]["allowed_tools"]["tools"]
assert entries[-1] == retrieval_entry
assert entries[:-1] == (tool_choice.get("tools") or tool_choice["allowed_tools"]["tools"])

79
tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py
View file

@ -422,6 +422,85 @@ async def test_execute_chat_completion_agentic_loop_strips_tool_choice():
)
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ["auto", "required"])
async def test_execute_chat_completion_agentic_loop_keeps_an_allowed_tools_choice(mode):
websearch_logger = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI])
tool_choice = {
"type": "allowed_tools",
"allowed_tools": {
"mode": mode,
"tools": [{"type": "function", "function": {"name": "litellm_web_search"}}],
},
}
captured_kwargs: dict = {}
async def fake_acompletion(**kwargs):
captured_kwargs.update(kwargs)
return ModelResponse(id="followup", model="gpt-4o", object="chat.completion")
async def fake_search(query, **_):
return ("Bitcoin price is $60,000", None)
with patch.object(websearch_logger, "_execute_search", side_effect=fake_search):
with patch("litellm.acompletion", side_effect=fake_acompletion):
await websearch_logger._execute_chat_completion_agentic_loop(
model="gpt-4o",
messages=[{"role": "user", "content": "What is Bitcoin price?"}],
tool_calls=[{"id": "call_1", "name": "litellm_web_search", "input": {"query": "bitcoin price"}}],
optional_params={
"tools": [
{"type": "function", "function": {"name": "litellm_web_search"}},
{"type": "function", "function": {"name": "get_time"}},
],
"tool_choice": tool_choice,
},
logging_obj=MagicMock(),
stream=False,
kwargs={},
)
assert captured_kwargs["tool_choice"] == tool_choice
@pytest.mark.asyncio
@pytest.mark.parametrize(
("tool_choice", "kept"),
[
(
{
"type": "allowed_tools",
"allowed_tools": {"mode": "required", "tools": [{"type": "function", "function": {"name": "get_time"}}]},
},
True,
),
({"type": "function", "function": {"name": "litellm_web_search"}}, False),
("required", False),
],
)
async def test_chat_followup_plan_keeps_only_an_allowed_tools_choice(
tool_choice: object, kept: bool, perplexity_search: respx.Route
) -> None:
plan: Final = await WebSearchInterceptionLogger().async_build_chat_completion_agentic_loop_plan(
tools={
"tool_calls": [{"id": "call_abc", "name": "litellm_web_search", "input": {"query": "latest ai news"}}],
"tool_type": "websearch",
"provider": "openai",
"response_format": "openai",
},
model="gpt-4o",
messages=[{"role": "user", "content": "What's the latest AI news?"}],
response=MagicMock(),
optional_params={"tool_choice": tool_choice, "max_tokens": 512},
logging_obj=None,
stream=False,
kwargs={"custom_llm_provider": "openai"},
)
assert plan.request_patch is not None
assert plan.request_patch.optional_params.get("tool_choice") == (tool_choice if kept else None)
if __name__ == "__main__":
# Run with: pytest test_websearch_chat_completion.py -v -s
pytest.main([__file__, "-v", "-s"])

181
tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py
View file

@ -8,12 +8,16 @@ import litellm
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.allowed_tools_tool_choice import (
ChatAllowedToolsToolChoice,
GatewayTool,
GatewayTools,
allowed_tools_tool_choice_in_effect,
allowed_tools_tool_choice_with_builtin_as_function,
allowed_tools_tool_choice_with_function,
allowed_tools_tool_choice_with_gateway_tools_as_functions,
drops_unsupported_allowed_tools_tool_choice,
is_allowed_tools_tool_choice,
chat_allowed_tools_tool_choice,
parse_allowed_tools_tool_choice,
refuse_a_malformed_allowed_tools_tool_choice,
refuse_tool_calls_the_choice_excludes,
tool_calls_the_choice_excludes,
unenforced_allowed_tools_tool_choice,
@ -266,11 +270,12 @@ def test_a_choice_the_provider_cannot_take_answers_400_without_drop_params():
*_CHOICES_WITH_A_NAME_THAT_IS_NOT_A_STRING,
],
)
def test_a_malformed_choice_is_refused(tool_choice: object):
def test_a_malformed_choice_is_refused_without_drop_params_and_dropped_with_it(tool_choice: object):
with pytest.raises(BadRequestError, match="Invalid tool choice") as exc:
refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model")
allowed_tools_tool_choice_in_effect(tool_choice, "the-model", False)
assert exc.value.status_code == 400
assert allowed_tools_tool_choice_in_effect(tool_choice, "the-model", True) is None
@pytest.mark.parametrize("tool_choice", _CHOICES_WITH_A_NAME_THAT_IS_NOT_A_STRING)
@ -279,8 +284,9 @@ def test_a_name_that_is_not_a_string_lists_no_tool(tool_choice: object):
@pytest.mark.parametrize("tool_choice", [_FLAT_CHOICE, _CHAT_LAYOUT, "required", None])
def test_a_valid_choice_or_any_other_choice_is_not_called_malformed(tool_choice: object):
assert refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") is None
@pytest.mark.parametrize("drop_params", [True, False])
def test_a_valid_choice_or_any_other_choice_stays_in_effect(tool_choice: object, drop_params: bool):
assert allowed_tools_tool_choice_in_effect(tool_choice, "the-model", drop_params) is tool_choice
@pytest.mark.parametrize(
@ -295,7 +301,7 @@ def test_a_valid_choice_or_any_other_choice_is_not_called_malformed(tool_choice:
)
def test_an_mcp_entry_with_a_null_name_is_a_valid_choice(tool_choice: object):
assert parse_allowed_tools_tool_choice(tool_choice) is not None
assert refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") is None
assert allowed_tools_tool_choice_in_effect(tool_choice, "the-model", False) is tool_choice
_DOCS_SEARCH: Final = {
@ -305,12 +311,36 @@ _DOCS_SEARCH: Final = {
}
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 _gateway_serving_search(*server_labels: str) -> GatewayTools:
return GatewayTools(
server_labels=frozenset({"docs", "other"}),
tools={
"docs_server-search": GatewayTool(
names=frozenset({"docs_server-search", "search"}), server_labels=frozenset(server_labels)
)
},
)
@pytest.mark.parametrize("name", ["search", "docs_server-search"])
def test_a_named_mcp_entry_allows_its_call_on_a_server_its_label_serves(name: str):
tool_choice: Final = {
"type": "allowed_tools",
"mode": "auto",
"tools": [{"type": "mcp", "server_label": "docs", "name": name}],
}
assert (
tool_calls_the_choice_excludes(tool_choice, ("docs_server-search",), _gateway_serving_search("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"})}) == ()
assert (
tool_calls_the_choice_excludes(_WHOLE_DOCS_SERVER, ("docs_server-search",), _gateway_serving_search("docs"))
== ()
)
@pytest.mark.parametrize(
@ -329,7 +359,9 @@ def test_an_mcp_entry_without_a_name_allows_every_call_its_server_serves():
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",)
assert tool_calls_the_choice_excludes(
tool_choice, ("docs_server-search",), _gateway_serving_search(*server_labels)
) == ("docs_server-search",)
def _is_web_search(tool: Mapping[str, object]) -> bool:
@ -388,3 +420,132 @@ def test_a_function_or_mcp_entry_is_never_taken_for_a_builtin():
)
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
_GATEWAY_WITH_TWO_SERVERS: Final = GatewayTools(
server_labels=frozenset({"docs", "weather"}),
tools={
"docs_server-search": GatewayTool(
names=frozenset({"docs_server-search", "search"}), server_labels=frozenset({"docs"})
),
"docs_server-read": GatewayTool(names=frozenset({"docs_server-read", "read"}), server_labels=frozenset({"docs"})),
"weather_server-forecast": GatewayTool(
names=frozenset({"weather_server-forecast", "forecast"}), server_labels=frozenset({"weather"})
),
},
)
@pytest.mark.parametrize(
("tool_choice", "expected"),
[
(
{
"type": "allowed_tools",
"mode": "required",
"tools": [{"type": "mcp", "server_label": "docs"}, {"type": "function", "name": "get_time"}],
},
{
"type": "allowed_tools",
"mode": "required",
"tools": [
{"type": "function", "name": "get_time"},
{"type": "function", "name": "docs_server-search"},
{"type": "function", "name": "docs_server-read"},
],
},
),
(
{
"type": "allowed_tools",
"allowed_tools": {
"mode": "auto",
"tools": [{"type": "mcp", "server_label": "weather", "name": "forecast"}],
},
},
{
"type": "allowed_tools",
"allowed_tools": {
"mode": "auto",
"tools": [{"type": "function", "function": {"name": "weather_server-forecast"}}],
},
},
),
],
)
def test_gateway_mcp_entries_become_the_functions_they_allow_in_the_callers_layout(
tool_choice: dict[str, object], expected: dict[str, object]
):
assert allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice, _GATEWAY_WITH_TWO_SERVERS) == expected
def test_an_mcp_entry_of_a_label_the_gateway_does_not_serve_is_forwarded_as_sent():
tool_choice: Final = {
"type": "allowed_tools",
"mode": "auto",
"tools": [{"type": "mcp", "server_label": "remote"}, {"type": "mcp", "server_label": "docs", "name": "read"}],
}
assert allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice, _GATEWAY_WITH_TWO_SERVERS) == {
"type": "allowed_tools",
"mode": "auto",
"tools": [{"type": "mcp", "server_label": "remote"}, {"type": "function", "name": "docs_server-read"}],
}
@pytest.mark.parametrize("mode", ["auto", "required"])
def test_gateway_entries_that_allow_no_converted_tool_send_none_instead_of_an_empty_list(mode: str):
tool_choice: Final = {
"type": "allowed_tools",
"mode": mode,
"tools": [{"type": "mcp", "server_label": "docs", "name": "delete"}],
}
assert allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice, _GATEWAY_WITH_TWO_SERVERS) == "none"
@pytest.mark.parametrize("tool_choice", ["auto", None, _FLAT_CHOICE, {"type": "allowed_tools", "mode": "auto", "tools": []}])
def test_a_choice_without_gateway_entries_is_forwarded_as_sent(tool_choice: object):
assert allowed_tools_tool_choice_with_gateway_tools_as_functions(tool_choice, _GATEWAY_WITH_TWO_SERVERS) is tool_choice
@pytest.mark.parametrize(
("tool_choice", "expected"),
[
(
{"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]},
{
"type": "allowed_tools",
"mode": "auto",
"tools": [{"type": "function", "name": "get_time"}, {"type": "function", "name": "retrieve"}],
},
),
(
{"type": "allowed_tools", "allowed_tools": {"mode": "required", "tools": [{"type": "mcp", "server_label": "x"}]}},
{
"type": "allowed_tools",
"allowed_tools": {
"mode": "required",
"tools": [{"type": "mcp", "server_label": "x"}, {"type": "function", "function": {"name": "retrieve"}}],
},
},
),
],
)
def test_an_added_function_joins_the_choice_in_the_callers_layout(
tool_choice: dict[str, object], expected: dict[str, object]
):
assert allowed_tools_tool_choice_with_function(tool_choice, "retrieve") == expected
@pytest.mark.parametrize(
"tool_choice",
[
"auto",
None,
{"type": "function", "function": {"name": "get_time"}},
{"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "retrieve"}]},
],
)
def test_a_choice_that_is_not_allowed_tools_or_already_lists_the_function_is_left_as_sent(tool_choice: object):
assert allowed_tools_tool_choice_with_function(tool_choice, "retrieve") is tool_choice

59
tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py
View file

@ -23,7 +23,7 @@ they stay out of the body even without it.
import datetime
import json
from collections.abc import Callable, Mapping, Sequence
from typing import Any, Dict, Final, List, Optional, Tuple
from typing import Any, ClassVar, Dict, Final, List, Optional, Tuple
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -1284,3 +1284,60 @@ async def test_acompletion_hands_the_callers_choice_to_the_chat_hooks(tool_choic
)
assert (hook.runs, mock_client.post.await_count, response.choices[0].message.tool_calls) == (hook_runs, 1, None)
class _RetrievalHookLogger(_ChatHookLogger):
server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({"litellm_code_execution"})
@pytest.mark.asyncio
@pytest.mark.parametrize("model_response", [_tool_call_model_response, _custom_tool_call_model_response])
async def test_a_chat_hook_runs_a_tool_litellm_serves_itself_that_no_choice_can_list(
model_response: Callable[[], ModelResponse],
):
hook: Final = _RetrievalHookLogger()
result: Final = await run_chat_completion_agentic_hooks(
response=model_response(),
model="gpt-4o-mini",
messages=[{"role": "user", "content": "what is 6*7?"}],
optional_params={},
kwargs={"_agentic_loop_tool_choice": _ALLOWS_ONLY_GET_WEATHER},
logging_obj=_logging_with(hook),
custom_llm_provider="openai",
stream=False,
)
assert hook.runs == 1
assert isinstance(result, ModelResponse)
assert result.choices[0].message.content == "ran the hook"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"tool_choice",
[
{"type": "allowed_tools", "mode": "auto", "tools": []},
{"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["get_weather"]}]},
],
)
async def test_a_malformed_choice_dropped_by_drop_params_never_holds_back_the_chat_hooks(
tool_choice: dict[str, object],
):
hook: Final = _ChatHookLogger()
mock_client: Final = MagicMock()
mock_client.post = AsyncMock(return_value=_raw_response_for(_tool_call_model_response()))
response: Final = await litellm.acompletion(
model="openai/gpt-4o-mini",
messages=[{"role": "user", "content": "what is 6*7?"}],
tools=[{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}],
tool_choice=tool_choice,
drop_params=True,
api_key="sk-test",
client=mock_client,
success_callback=[hook],
)
assert (hook.runs, response.choices[0].message.content) == (1, "ran the hook")

89
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -455,6 +455,95 @@ def test_response_api_handler_sends_no_stream_options_when_web_search_forces_a_n
assert "stream_options" not in recorder.seen_kwargs
class _LookupLoop(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.runs = 0
async def async_should_run_agentic_loop(
self,
response: object,
model: str,
messages: list[dict[str, object]],
tools: list[dict[str, object]] | None,
stream: bool,
custom_llm_provider: str,
kwargs: dict[str, object],
) -> tuple[bool, dict[str, object]]:
return True, {"tool_calls": [{"id": "call_lookup", "name": "lookup"}]}
async def async_run_agentic_loop(
self,
tools: dict,
model: str,
messages: list[dict],
response: object,
anthropic_messages_provider_config: object,
anthropic_messages_optional_request_params: dict,
logging_obj: object,
stream: bool,
kwargs: dict,
) -> object:
self.runs += 1
return response
def _a_lookup_call(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
json={
"id": "resp_1",
"object": "response",
"created_at": 0,
"status": "completed",
"model": "gpt-4o",
"output": [
{
"type": "function_call",
"id": "fc_1",
"call_id": "call_lookup",
"name": "lookup",
"arguments": "{}",
"status": "completed",
}
],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
},
)
@pytest.mark.parametrize("allowed_name,runs", [("get_time", 0), ("lookup", 1)])
def test_response_api_handler_runs_no_hook_call_the_allowed_tools_choice_excludes(allowed_name: str, runs: int):
hook: Final = _LookupLoop()
response: Final = BaseLLMHTTPHandler().response_api_handler(
model="gpt-4o",
input="what time is it?",
responses_api_provider_config=AzureOpenAIResponsesAPIConfig(),
response_api_optional_request_params={
"tools": [
{"type": "function", "name": "get_time", "parameters": {"type": "object", "properties": {}}},
{"type": "function", "name": "lookup", "parameters": {"type": "object", "properties": {}}},
],
"tool_choice": {
"type": "allowed_tools",
"mode": "auto",
"tools": [{"type": "function", "name": allowed_name}],
},
},
custom_llm_provider="azure",
litellm_params=GenericLiteLLMParams(
api_key="sk-test", api_base="https://example.openai.azure.com", api_version="2025-04-01-preview"
),
logging_obj=Mock(litellm_call_id="call-1", dynamic_success_callbacks=[hook], model_call_details={}),
client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_a_lookup_call))),
)
assert hook.runs == runs
assert isinstance(response, ResponsesAPIResponse)
assert len(response.output) == runs
@pytest.mark.asyncio
async def test_async_response_api_handler_streams_when_provider_transform_adds_stream():
handler = BaseLLMHTTPHandler()

34
tests/unit/responses/litellm_completion_transformation/test_handler.py
View file

@ -332,6 +332,40 @@ async def test_bridged_next_turn_replays_a_turn_stored_under_the_id_the_client_h
]
@pytest.mark.asyncio
async def test_bridged_follow_up_after_a_client_tool_search_reports_auto_for_an_allowed_tools_choice_it_dropped():
previous_response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="azure_ai", model_id="deployment-a", response_id="chatcmpl-turn-one"
)
provider: Final = _RecordingAnthropicHandler(_ANTHROPIC_MESSAGE_PAYLOAD)
client: Final = AsyncHTTPHandler()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(provider))
with patch(
"litellm.proxy.proxy_server.prisma_client",
SimpleNamespace(db=_StoredTurns((_stored_mcp_turn(previous_response_id),))),
):
response = await litellm.aresponses(
model="azure_ai/claude-sonnet-4-6",
api_base="https://fake-foundry-resource.services.ai.azure.com",
api_key="fake-api-key",
input="Double it",
tools=[
{"type": "tool_search", "execution": "client"},
{"type": "function", "name": "get_time", "parameters": {"type": "object", "properties": {}}},
],
tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]},
drop_params=True,
previous_response_id=previous_response_id,
client=client,
)
assert provider.request_body is not None, "the bridged turn never reached the provider"
assert "tool_choice" not in provider.request_body
assert isinstance(response, ResponsesAPIResponse)
assert response.tool_choice == "auto"
@pytest.mark.asyncio
@pytest.mark.parametrize("client_field", ["caller_input", "names_non_string_fields"])
async def test_bridged_follow_up_after_a_client_tool_search_with_a_client_field_named_like_a_bridge_keyword_reaches_the_provider(

167
tests/unit/responses/mcp/test_aresponses_api_with_mcp.py
View file

@ -36,6 +36,7 @@ from litellm.types.llms.openai import (
ResponsesAPIStreamingResponse,
ToolParam,
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.utils import ModelResponse, StandardLoggingPayload, Usage
@ -1704,6 +1705,15 @@ async def _auto_execute_echo(
return await _auto_execute(logging_obj, _ECHO_MCP_TOOL, "PING")
_DEMO_MCP_SERVER: Final = MCPServer(
server_id="demo", name="demo", server_name="demo", transport="http", url="https://demo.example/mcp", auth_type="none"
)
def _demo_server_named(name: str) -> MCPServer | None:
return _DEMO_MCP_SERVER if name == "demo" else None
def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock) -> None:
monkeypatch.setitem(
sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None)
@ -1715,7 +1725,7 @@ def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool
get_registry=MagicMock(return_value={}),
call_tool=call_tool,
get_mcp_server_from_tool_name=MagicMock(return_value=None),
get_mcp_server_by_name=MagicMock(return_value=None),
get_mcp_server_by_name=MagicMock(side_effect=_demo_server_named),
),
)
@ -1806,26 +1816,13 @@ def _anthropic_call_to_echo() -> httpx.Response:
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"tool_choice",
[
_ALLOWED_ONLY_ANOTHER_TOOL,
{"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["echo"]}]},
],
)
async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes(
tool_choice: dict[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
async def _call_echo_on_anthropic_under_drop_params(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[httpx.Response], tool_choice: object
) -> tuple[object, respx.Route]:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
litellm.in_memory_llm_clients_cache.flush_cache()
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=_anthropic_call_to_echo()
)
call_tool: Final = AsyncMock()
_served_by_a_recording_mcp_server(monkeypatch, call_tool)
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=list(rounds))
with _echo_tool_listed():
response: Final = await litellm.aresponses(
model="anthropic/claude-haiku-4-5",
@ -1835,6 +1832,19 @@ async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes(
drop_params=True,
api_key="sk-ant-fake",
)
return response, route
@pytest.mark.asyncio
async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
call_tool: Final = AsyncMock()
_served_by_a_recording_mcp_server(monkeypatch, call_tool)
response, route = await _call_echo_on_anthropic_under_drop_params(
respx_mock, monkeypatch, (_anthropic_call_to_echo(),), _ALLOWED_ONLY_ANOTHER_TOOL
)
call_tool.assert_not_awaited()
assert "tool_choice" not in json.loads(route.calls[0].request.content)
@ -1843,6 +1853,127 @@ async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes(
assert _output_types(response.output) == ("function_call",)
def _anthropic_answer() -> httpx.Response:
return httpx.Response(
200,
json={
"id": "msg_answer",
"type": "message",
"role": "assistant",
"model": "claude-haiku-4-5",
"content": [{"type": "text", "text": "PING back"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 5},
},
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"tool_choice",
[
{"type": "allowed_tools", "mode": "auto", "tools": []},
{"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["echo"]}]},
],
)
async def test_a_malformed_choice_under_drop_params_runs_an_mcp_call_as_if_never_sent(
tool_choice: dict[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")]))
_served_by_a_recording_mcp_server(monkeypatch, call_tool)
response, route = await _call_echo_on_anthropic_under_drop_params(
respx_mock, monkeypatch, (_anthropic_call_to_echo(), _anthropic_answer()), tool_choice
)
assert call_tool.await_count == 1
assert all("tool_choice" not in json.loads(call.request.content) for call in route.calls)
assert isinstance(response, ResponsesAPIResponse)
assert "message" in _output_types(response.output)
@pytest.mark.asyncio
async def test_a_malformed_choice_without_drop_params_answers_400_before_listing_or_calling_a_model(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
listing: Final = AsyncMock()
_served_by_a_recording_mcp_server(monkeypatch, AsyncMock())
monkeypatch.setattr("litellm.proxy.mcp_server.server._get_tools_from_mcp_servers", listing)
with pytest.raises(litellm.BadRequestError, match="Invalid tool choice"):
await litellm.aresponses(
model="openai/gpt-5.6",
input="Call echo with PING",
tools=[dict(_ECHO_MCP_TOOL)],
tool_choice={"type": "allowed_tools", "mode": "auto", "tools": []},
api_key="sk-mcp-rounds",
api_base="https://mcp-rounds.example.com/v1",
)
listing.assert_not_awaited()
assert not respx_mock.calls
_DEMO_LABELED_ECHO_TOOL: Final = MappingProxyType({**_ECHO_MCP_TOOL, "server_label": "demo"})
async def _call_echo_with_the_demo_label(tool_choice: object, stream: bool) -> tuple[object, ...]:
with _echo_tool_listed():
response: Final = await litellm.aresponses(
model="openai/gpt-5.6",
input="Call echo with PING",
tools=[dict(_DEMO_LABELED_ECHO_TOOL)],
tool_choice=tool_choice,
api_key="sk-mcp-rounds",
api_base="https://mcp-rounds.example.com/v1",
stream=stream,
)
return await _drain(response) if stream else (response,)
@pytest.mark.asyncio
async def test_the_provider_gets_the_converted_function_in_place_of_a_gateway_mcp_entry(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
route: Final = _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY))
call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")]))
_served_by_a_recording_mcp_server(monkeypatch, call_tool)
await _call_echo_with_the_demo_label(
{"type": "allowed_tools", "mode": "required", "tools": [{"type": "mcp", "server_label": "demo", "name": "echo"}]},
stream=False,
)
assert json.loads(route.calls[0].request.content)["tool_choice"] == {
"type": "allowed_tools",
"mode": "required",
"tools": [{"type": "function", "name": "echo"}],
}
assert call_tool.await_count == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
async def test_a_gateway_mcp_entry_that_allows_no_converted_tool_sends_none_and_runs_nothing(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool
) -> None:
route: Final = _serve_responses(
respx_mock,
monkeypatch,
(_streamed(_ROUND_ONE_BODY) if stream else httpx.Response(200, content=json.dumps(_ROUND_ONE_BODY, default=dict)),),
)
call_tool: Final = AsyncMock()
_served_by_a_recording_mcp_server(monkeypatch, call_tool)
await _call_echo_with_the_demo_label(
{"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "demo", "name": "reverse"}]},
stream=stream,
)
assert json.loads(route.calls[0].request.content)["tool_choice"] == "none"
call_tool.assert_not_awaited()
@pytest.mark.asyncio
async def test_an_mcp_tool_the_choice_lists_still_runs(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch):
_serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY))

181
tests/unit/responses/mcp/test_chat_completions_handler.py
View file

@ -2773,10 +2773,19 @@ _DEMO_MCP_SERVER: Final = MCPServer(
)
def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock) -> None:
def _demo_server_named(name: str) -> MCPServer | None:
return _DEMO_MCP_SERVER if name.lower() == "demo" else None
def _served_by_a_recording_mcp_server(
monkeypatch: pytest.MonkeyPatch,
call_tool: AsyncMock,
listed: MCPTool = _MCP_ECHO_TOOL,
catalog: SimpleNamespace = SimpleNamespace(operation=nullcontext),
) -> None:
monkeypatch.setattr(
"litellm.proxy.mcp_server.server._get_tools_from_mcp_servers",
AsyncMock(return_value=AggregateToolListing(tools=[_MCP_ECHO_TOOL], outcomes={})),
AsyncMock(return_value=AggregateToolListing(tools=[listed], outcomes={})),
)
monkeypatch.setitem(
sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None)
@ -2784,11 +2793,11 @@ def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool
monkeypatch.setattr(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager",
SimpleNamespace(
catalog=SimpleNamespace(operation=nullcontext),
catalog=catalog,
get_registry=MagicMock(return_value={}),
call_tool=call_tool,
get_mcp_server_from_tool_name=MagicMock(return_value=None),
get_mcp_server_by_name=MagicMock(return_value=None),
get_mcp_server_by_name=MagicMock(side_effect=_demo_server_named),
get_allowed_mcp_servers=AsyncMock(return_value=[_DEMO_MCP_SERVER.server_id]),
get_mcp_servers_from_ids=MagicMock(return_value=[_DEMO_MCP_SERVER]),
),
@ -2824,7 +2833,7 @@ async def test_an_mcp_tool_the_choice_excludes_never_runs_when_the_model_ignored
call_tool.assert_not_awaited()
def _anthropic_call_to_echo() -> Response:
def _anthropic_call_to_echo(name: str = "echo") -> Response:
return Response(
200,
json={
@ -2832,7 +2841,7 @@ def _anthropic_call_to_echo() -> Response:
"type": "message",
"role": "assistant",
"model": "claude-haiku-4-5",
"content": [{"type": "tool_use", "id": "toolu_echo", "name": "echo", "input": {"value": "PING"}}],
"content": [{"type": "tool_use", "id": "toolu_echo", "name": name, "input": {"value": "PING"}}],
"stop_reason": "tool_use",
"usage": {"input_tokens": 10, "output_tokens": 5},
},
@ -2918,3 +2927,163 @@ async def test_a_dropped_choice_still_runs_an_mcp_call_it_lists(
assert call_tool.await_count == 1
assert response.choices[0].message.content == "PING back"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"tool_choice",
[
{"type": "allowed_tools", "mode": "auto", "tools": []},
{"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["echo"]}]},
],
)
async def test_a_malformed_choice_under_drop_params_runs_an_mcp_call_as_if_never_sent(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tool_choice: object
):
call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")]))
_served_by_a_recording_mcp_server(monkeypatch, call_tool)
response, route = await _call_echo_on_anthropic_under_drop_params(
respx_mock, monkeypatch, (_anthropic_call_to_echo(), _anthropic_answer()), tool_choice
)
assert call_tool.await_count == 1
assert response.choices[0].message.content == "PING back"
assert all("tool_choice" not in json.loads(call.request.content) for call in route.calls)
_DEMO_ECHO_TOOL: Final = MCPTool.model_validate(
{"name": "demo-echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}},
by_name=False,
)
_DEMO_IN_AN_ACCESS_GROUP: Final = _DEMO_MCP_SERVER.model_copy(update={"access_groups": ["demo_group"]})
def _choice_of(*tools: Mapping[str, object]) -> Mapping[str, object]:
return {"type": "allowed_tools", "mode": "auto", "tools": list(tools)}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("tool_choice", "server_url", "listed"),
[
(_choice_of({"type": "mcp", "server_label": "demo", "name": "echo"}), "litellm_proxy/mcp/demo", _DEMO_ECHO_TOOL),
(
_choice_of({"type": "mcp", "server_label": "demo", "name": "demo-echo"}),
"litellm_proxy/mcp/demo",
_DEMO_ECHO_TOOL,
),
(_WHOLE_DEMO_SERVER, "litellm_proxy/mcp/DEMO", _MCP_ECHO_TOOL),
(_WHOLE_DEMO_SERVER, "litellm_proxy/mcp/demo_group", _MCP_ECHO_TOOL),
],
)
async def test_an_mcp_entry_allows_the_tool_the_gateway_resolves_for_its_label_and_name(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
tool_choice: object,
server_url: str,
listed: MCPTool,
):
call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")]))
_served_by_a_recording_mcp_server(
monkeypatch,
call_tool,
listed,
SimpleNamespace(
operation=nullcontext,
current=lambda: SimpleNamespace(servers={_DEMO_MCP_SERVER.server_id: _DEMO_IN_AN_ACCESS_GROUP}),
),
)
response, _route = await _call_echo_on_anthropic_under_drop_params(
respx_mock, monkeypatch, (_anthropic_call_to_echo(listed.name), _anthropic_answer()), tool_choice, server_url
)
assert call_tool.await_count == 1
assert response.choices[0].message.content == "PING back"
@pytest.mark.asyncio
async def test_a_named_mcp_entry_allows_no_other_tool_on_its_server(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
call_tool: Final = AsyncMock()
_served_by_a_recording_mcp_server(monkeypatch, call_tool, _DEMO_ECHO_TOOL)
response, _route = await _call_echo_on_anthropic_under_drop_params(
respx_mock,
monkeypatch,
(_anthropic_call_to_echo("demo-echo"),),
_choice_of({"type": "mcp", "server_label": "demo", "name": "reverse"}),
)
call_tool.assert_not_awaited()
assert [call.function.name for call in response.choices[0].message.tool_calls] == ["demo-echo"]
async def _call_echo_on_openai_with_the_demo_server(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], tool_choice: object
) -> tuple[object, ...]:
_serve_chat_rounds(respx_mock, monkeypatch, rounds)
await litellm.acompletion(
model="openai/gpt-5.6",
messages=[{"role": "user", "content": "Call echo with PING"}],
tools=[
{
"type": "mcp",
"server_label": "demo",
"server_url": "litellm_proxy/mcp/demo",
"require_approval": "never",
},
{"type": "function", "function": {"name": "get_time", "parameters": {"type": "object"}}},
],
tool_choice=tool_choice,
api_key="sk-mcp-chat-rounds",
api_base="https://mcp-chat-rounds.example.com/v1",
)
return tuple(json.loads(call.request.content).get("tool_choice") for call in respx_mock.calls)
@pytest.mark.asyncio
async def test_the_provider_gets_the_converted_function_in_place_of_a_gateway_mcp_entry(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")]))
_served_by_a_recording_mcp_server(monkeypatch, call_tool)
sent: Final = await _call_echo_on_openai_with_the_demo_server(
respx_mock,
monkeypatch,
(_ok(_CHAT_ROUND_ONE), _ok(_CHAT_ROUND_TWO)),
_choice_of({"type": "mcp", "server_label": "demo"}, {"type": "function", "name": "get_time"}),
)
assert sent[0] == {
"type": "allowed_tools",
"allowed_tools": {
"mode": "auto",
"tools": [
{"type": "function", "function": {"name": "get_time"}},
{"type": "function", "function": {"name": "echo"}},
],
},
}
assert call_tool.await_count == 1
@pytest.mark.asyncio
async def test_a_gateway_mcp_entry_that_allows_no_converted_tool_sends_none_and_runs_nothing(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
call_tool: Final = AsyncMock()
_served_by_a_recording_mcp_server(monkeypatch, call_tool)
sent: Final = await _call_echo_on_openai_with_the_demo_server(
respx_mock,
monkeypatch,
(_ok(_CHAT_ROUND_ONE),),
_choice_of({"type": "mcp", "server_label": "demo", "name": "reverse"}),
)
assert sent == ("none",)
call_tool.assert_not_awaited()

92
tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py
View file

@ -1632,14 +1632,14 @@ async def _no_toolset_grants(context: UserAPIKeyAuth) -> frozenset[str]:
return frozenset()
async def _tools_listed_through_the_gateway(
async def _listing_through_the_gateway(
monkeypatch: pytest.MonkeyPatch,
user_api_key_auth: UserAPIKeyAuth,
server_urls: tuple[str, ...],
mcp_tools: Sequence[Mapping[str, object]],
granted_toolsets: GrantedToolsetIds,
toolset_table: types.SimpleNamespace | None = None,
agent_permission: LiteLLM_ObjectPermissionTable | None = None,
) -> list[str]:
) -> tuple[list[MCPTool], list[str]]:
async def upstream_listing(_client: object, server_name: str) -> list[MCPTool]:
return list(UPSTREAM_TOOLS[server_name])
@ -1665,17 +1665,99 @@ async def _tools_listed_through_the_gateway(
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"supported_db_objects": []})
try:
tools, _server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager(
return await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager(
user_api_key_auth=user_api_key_auth,
mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": url} for url in server_urls],
mcp_tools_with_litellm_proxy=mcp_tools,
granted_toolsets=granted_toolsets,
)
finally:
for server in GATEWAY_SERVERS:
manager._drop_listed_tools(server.server_id)
async def _tools_listed_through_the_gateway(
monkeypatch: pytest.MonkeyPatch,
user_api_key_auth: UserAPIKeyAuth,
server_urls: tuple[str, ...],
granted_toolsets: GrantedToolsetIds,
toolset_table: types.SimpleNamespace | None = None,
agent_permission: LiteLLM_ObjectPermissionTable | None = None,
) -> list[str]:
tools, _server_names = await _listing_through_the_gateway(
monkeypatch,
user_api_key_auth,
[{"type": "mcp", "server_url": url} for url in server_urls],
granted_toolsets,
toolset_table,
agent_permission,
)
return sorted(tool.name for tool in tools)
async def _calls_the_gateway_runs(
monkeypatch: pytest.MonkeyPatch,
labeled_server_urls: Mapping[str, str],
allowed_label: str,
called: tuple[str, ...],
unresolvable_tools: Mapping[str, str] = MappingProxyType({}),
) -> tuple[str, ...]:
user: Final = UserAPIKeyAuth(api_key="sk-team", team_id="team-1", user_role=LitellmUserRoles.INTERNAL_USER)
mcp_tools: Final = [
{"type": "mcp", "server_label": label, "server_url": url} for label, url in labeled_server_urls.items()
]
tools, server_names = await _listing_through_the_gateway(monkeypatch, user, mcp_tools, _team_grants("ts-granted"))
_deduplicated, tool_server_map = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools(tools, server_names)
tool_choice: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": allowed_label}]}
gateway: Final = await LiteLLM_Proxy_MCP_Handler.gateway_tools(
tool_choice, {**tool_server_map, **unresolvable_tools}, mcp_tools, user, _team_grants("ts-granted")
)
return tuple(
name
for name in called
if LiteLLM_Proxy_MCP_Handler.tool_calls_to_run(
({"type": "function_call", "name": name, "call_id": f"call_{name}", "arguments": "{}"},),
tool_choice,
gateway,
)
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("labeled_server_urls", "allowed_label", "runs"),
[
({"team": "litellm_proxy/mcp/team-toolset", "docs": "litellm_proxy/mcp/docs"}, "team", ("calc-add",)),
({"team": "litellm_proxy/mcp/team-toolset", "docs": "litellm_proxy/mcp/docs"}, "docs", ("docs-lookup",)),
({"team": "litellm_proxy/mcp/team-toolset", "other": "litellm_proxy/mcp/other-toolset"}, "other", ()),
],
ids=["granted_toolset", "server_beside_a_toolset", "ungranted_toolset"],
)
async def test_an_mcp_entry_allows_only_the_tools_its_labels_toolset_or_server_serves(
monkeypatch: pytest.MonkeyPatch,
labeled_server_urls: Mapping[str, str],
allowed_label: str,
runs: tuple[str, ...],
):
ran: Final = await _calls_the_gateway_runs(
monkeypatch, labeled_server_urls, allowed_label, ("calc-add", "docs-lookup")
)
assert ran == runs
@pytest.mark.asyncio
async def test_an_mcp_entry_allows_no_call_on_a_server_the_gateway_cannot_resolve(monkeypatch: pytest.MonkeyPatch):
ran: Final = await _calls_the_gateway_runs(
monkeypatch,
{"team": "litellm_proxy/mcp/team-toolset"},
"team",
("calc-add", "ghost-add"),
MappingProxyType({"ghost-add": "ghost"}),
)
assert ran == ("calc-add",)
def _key_holding(*server_ids: str, role: LitellmUserRoles = LitellmUserRoles.INTERNAL_USER) -> UserAPIKeyAuth:
return UserAPIKeyAuth(
api_key="sk-team",

3
tests/unit/test_main.py
View file

@ -656,7 +656,7 @@ def test_completion_drop_params_drops_an_allowed_tools_tool_choice_it_cannot_sen
{"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "mcp", "name": ["get_time"]}]}},
],
)
async def test_a_malformed_allowed_tools_tool_choice_stays_a_400_on_mcp_gateway_tools(
async def test_a_malformed_allowed_tools_tool_choice_is_a_400_on_mcp_gateway_tools_without_drop_params(
tool_choice: dict[str, object], respx_mock: respx.MockRouter
):
with pytest.raises(litellm.BadRequestError, match="Invalid tool choice") as exc:
@ -665,7 +665,6 @@ async def test_a_malformed_allowed_tools_tool_choice_stays_a_400_on_mcp_gateway_
messages=[{"role": "user", "content": "What time is it?"}],
tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}],
tool_choice=tool_choice,
drop_params=True,
api_base="http://localhost:12346/v1",
api_key="fake_api_key",
)

34
tests/unit/test_tool_loop.py
View file

@ -350,13 +350,45 @@ def test_a_malformed_choice_is_refused_before_any_request(
tools=WEATHER_TOOLS,
execute_tool=_tool_result,
tool_choice=tool_choice,
drop_params=True,
api_key="sk-test",
)
assert not respx_mock.calls
@pytest.mark.parametrize(
"tool_choice",
[_EMPTY_ALLOWED_TOOLS, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["get_time"]}]}],
)
def test_a_malformed_choice_under_drop_params_runs_the_loop_as_if_never_sent(
tool_choice: dict[str, object], respx_mock: respx.MockRouter
) -> None:
respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock(
side_effect=[
httpx.Response(200, json=_openai_response(None, [_TIME_CALL])),
httpx.Response(200, json=_openai_response("done")),
]
)
executed: Final[list[str | None]] = []
def execute(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage:
executed.append(tool_call.function.name)
return _tool_result(tool_call)
answer: Final = litellm.run_tool_loop(
model="openai/gpt-5-mini",
messages=[{"role": "user", "content": "hi"}],
tools=WEATHER_TOOLS,
execute_tool=execute,
tool_choice=tool_choice,
drop_params=True,
api_key="sk-test",
)
assert (answer, executed) == ("done", ["get_time"])
assert all("tool_choice" not in body for body in _request_bodies(respx_mock))
def test_a_choice_on_a_model_alias_reaches_the_loop(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None: