From 6be4f0b9eadb487a2c8fc3042fae6098cfa309a3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 14:45:40 -0800 Subject: [PATCH] fix(tool-choice): resolve mcp entries like the gateway and keep the choice on every hook path 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 --- .../compression_interception/handler.py | 5 + .../websearch_interception/handler.py | 46 +++-- .../agentic_loop_refusal.py | 10 + .../allowed_tools_tool_choice.py | 151 +++++++++++---- .../chat_completion_agentic_loop.py | 27 ++- litellm/llms/custom_httpx/llm_http_handler.py | 16 +- litellm/main.py | 7 +- litellm/responses/main.py | 24 ++- .../responses/mcp/chat_completions_handler.py | 13 +- .../mcp/litellm_proxy_mcp_handler.py | 146 ++++++++++++-- .../responses/mcp/mcp_streaming_iterator.py | 10 +- litellm/responses/mcp/tool_precheck.py | 15 -- litellm/tool_loop.py | 23 +-- litellm/types/integrations/custom_logger.py | 5 + litellm/utils.py | 23 +-- .../test_compression_interception_handler.py | 47 +++++ .../test_websearch_chat_completion.py | 79 ++++++++ .../test_allowed_tools_tool_choice.py | 181 +++++++++++++++++- .../test_chat_completion_agentic_loop.py | 59 +++++- .../custom_httpx/test_llm_http_handler.py | 89 +++++++++ .../test_handler.py | 34 ++++ .../mcp/test_aresponses_api_with_mcp.py | 167 ++++++++++++++-- .../mcp/test_chat_completions_handler.py | 181 +++++++++++++++++- .../mcp/test_litellm_proxy_mcp_handler.py | 92 ++++++++- tests/unit/test_main.py | 3 +- tests/unit/test_tool_loop.py | 34 +++- 26 files changed, 1303 insertions(+), 184 deletions(-) diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 5352ce6b6a0..6c04d16f817 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -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()) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 632864d05c9..4b2773dbaa3 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -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 diff --git a/litellm/litellm_core_utils/agentic_loop_refusal.py b/litellm/litellm_core_utils/agentic_loop_refusal.py index f3ee15dffdd..828d5fa03ad 100644 --- a/litellm/litellm_core_utils/agentic_loop_refusal.py +++ b/litellm/litellm_core_utils/agentic_loop_refusal.py @@ -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"))) diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index 1d254aeb559..ee6a4fa5cb8 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -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: diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 42a8b3013ee..53451b04515 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -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, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a10db31606..843be4fca62 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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, diff --git a/litellm/main.py b/litellm/main.py index 77220aba5ac..d61b9270b63 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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 {} ), }, diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f70f9758cb5..f346ac84baa 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -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: diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index f622362dce4..f956459cd1b 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -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: diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index d1188e2ecbc..9863b448fd4 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -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 diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 2f366837b3b..8e8baf7667c 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -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 () diff --git a/litellm/responses/mcp/tool_precheck.py b/litellm/responses/mcp/tool_precheck.py index cbd8a36d049..a67d81e70ae 100644 --- a/litellm/responses/mcp/tool_precheck.py +++ b/litellm/responses/mcp/tool_precheck.py @@ -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) diff --git a/litellm/tool_loop.py b/litellm/tool_loop.py index 9fc9af3a1c3..8b19471055b 100644 --- a/litellm/tool_loop.py +++ b/litellm/tool_loop.py @@ -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, diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 14ed719a880..127b3d61753 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -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. diff --git a/litellm/utils.py b/litellm/utils.py index 6a960a32a84..7b2464f3489 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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: diff --git a/tests/unit/integrations/compression_interception/test_compression_interception_handler.py b/tests/unit/integrations/compression_interception/test_compression_interception_handler.py index d7a1d6f14e1..549c0363942 100644 --- a/tests/unit/integrations/compression_interception/test_compression_interception_handler.py +++ b/tests/unit/integrations/compression_interception/test_compression_interception_handler.py @@ -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"]) diff --git a/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py b/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py index fc415374096..89ad8b72070 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py @@ -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"]) diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index 5a87f0e6434..ec398e8731a 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -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 diff --git a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py index 20bd02811fc..7d41cbd6628 100644 --- a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -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") diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index dd0e5762845..71abfd84496 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -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() diff --git a/tests/unit/responses/litellm_completion_transformation/test_handler.py b/tests/unit/responses/litellm_completion_transformation/test_handler.py index 4794f3701c5..1611225dfe9 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_handler.py @@ -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( diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 01b223f5e26..c48433989ea 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -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)) diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index ea013184a15..b0e8af97e3e 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -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() diff --git a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py index a16661c0174..220d46d90ea 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -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", diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 7313a1600f9..37c7a4cd422 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -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", ) diff --git a/tests/unit/test_tool_loop.py b/tests/unit/test_tool_loop.py index e679f9817c1..f56a88e8899 100644 --- a/tests/unit/test_tool_loop.py +++ b/tests/unit/test_tool_loop.py @@ -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: