From e0686d791e7dcfdba3258098f369e1132d3acd16 Mon Sep 17 00:00:00 2001 From: nate-berri Date: Sat, 10 Oct 2026 09:22:06 +0000 Subject: [PATCH 1/9] refactor(mcp): relocate litellm_proxy_mcp_handler under proxy mcp_server Part of LIT-9423 --- .../mcp_server/litellm_proxy_mcp_handler.py | 1408 ++++++++++++++++ .../mcp/litellm_proxy_mcp_handler.py | 1411 +---------------- .../test_litellm_proxy_mcp_handler.py | 24 +- .../mcp/test_chat_completions_handler.py | 4 +- 4 files changed, 1432 insertions(+), 1415 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/litellm_proxy_mcp_handler.py rename tests/unit/{responses/mcp => proxy/_experimental/mcp_server}/test_litellm_proxy_mcp_handler.py (97%) diff --git a/litellm/proxy/_experimental/mcp_server/litellm_proxy_mcp_handler.py b/litellm/proxy/_experimental/mcp_server/litellm_proxy_mcp_handler.py new file mode 100644 index 00000000000..d3e3509b8f0 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/litellm_proxy_mcp_handler.py @@ -0,0 +1,1408 @@ +import re +import traceback +from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence +from datetime import datetime +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload + +from openai.types.chat import ChatCompletionToolParam +from openai.types.responses.function_tool_param import FunctionToolParam + +from litellm._logging import verbose_logger +from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager +from litellm.proxy._experimental.mcp_server.utils import ( + iter_known_server_prefixes, + logging_safe_mcp_headers, + split_server_prefix_from_name, + strip_known_server_prefix, +) +from litellm.responses.main import aresponses # noqa: TID251 # inner call must skip the MCP gateway that invoked it +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator +from litellm.types.llms.openai import ( + ResponseInputParam, + ResponsesAPIResponse, + ResponsesAPIStreamingResponse, +) +from litellm.types.llms.openai import ToolParam as ResponsesToolParam +from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.types.utils import ( + CallTypes, + ChatCompletionMessageCustomToolCall, + ChatCompletionMessageToolCall, + Choices, + Message, + ModelResponse, + StandardLoggingMCPToolCall, +) +from litellm.utils import Rules, function_setup + +if TYPE_CHECKING: + from mcp.types import CallToolResult + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.ui_session_utils import GrantedToolsetIds + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging +else: + MCPTool = Any + +# NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling +ToolParam: TypeAlias = Mapping[str, object] +SplitTools: TypeAlias = tuple[list[ToolParam], list[Any]] + + +class MCPToolResult(TypedDict): + tool_call_id: str | None + result: str + name: str | None + + +LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" +LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" + +_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") + + +def _mcp_server_url(tool: ToolParam) -> str | None: + if not isinstance(tool, dict) or tool.get("type") != "mcp": + return None + server_url: Final = tool.get("server_url") + return server_url if isinstance(server_url, str) else None + + +def _names_gateway_explicitly(tool: ToolParam) -> bool: + return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) + + +def _proxy_path_mcp_name(tool: ToolParam) -> 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) + return None if match is None else match.group(1) + + +def _registered_mcp_servers() -> Collection[MCPServer]: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + return global_mcp_server_manager.get_registry().values() + + +def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool: + requested: Final = name.lower() + return any( + requested in (known.lower() for known in (*iter_known_server_prefixes(server), server.name)) + or name in (server.access_groups or ()) + for server in servers + ) + + +async def _toolset_exists(name: str) -> bool: + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return False + return await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is not None + except Exception as e: + verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, e) + return False + + +@catalog_operation(global_manager) +async def _gateway_served_names( + names: Collection[str], + servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers, + toolset_exists: Callable[[str], Awaitable[bool]] = _toolset_exists, +) -> frozenset[str]: + registered: Final = tuple(servers()) if names else () + return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)]) + + +async def _served_mcp_path_names( + tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] +) -> frozenset[str]: + names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None) + return await served_names(names) if names else frozenset[str]() + + +class LiteLLM_Proxy_MCP_Handler: + """ + Helper class with static methods for MCP integration with Responses API. + + This handles when a user passes mcp server_url="litellm_proxy" in their tools. + """ + + @staticmethod + def get_parent_request_tags(kwargs: dict[str, Any] | None) -> list[str]: + """Tags from the parent LLM request, using the same extraction logic as standard logging (incl. User-Agent).""" + if not kwargs: + return [] + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + litellm_params: Final = kwargs.get("litellm_params") or kwargs + proxy_server_request = litellm_params.get("proxy_server_request") or kwargs.get("proxy_server_request") or {} + return StandardLoggingPayloadSetup.get_request_tags( + litellm_params=litellm_params, + proxy_server_request=proxy_server_request, + ) + + _get_parent_request_tags = get_parent_request_tags + + @staticmethod + def should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool: + """True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in + /mcp/. `_split_mcp_tools` then settles which of the latter the gateway actually serves.""" + return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ()) + + _should_use_litellm_mcp_gateway = should_use_litellm_mcp_gateway + + @staticmethod + def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools: + items: Final = tuple(tools or ()) + gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)] + other_tools: Final[list[ToolParam]] = [tool for tool in items if not _names_gateway_explicitly(tool)] + return gateway_tools, other_tools + + @staticmethod + async def split_mcp_tools( + tools: Iterable[Mapping[str, object]] | None, + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, + ) -> SplitTools: + items: Final = tuple(tools or ()) + served: Final = await _served_mcp_path_names(items, served_names) + return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools( + [ + {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} + if (name := _proxy_path_mcp_name(tool)) in served + else tool + for tool in items + ] + ) + + _split_mcp_tools = split_mcp_tools + + @staticmethod + @catalog_operation(global_manager) + async def routes_through_gateway( + tools: Iterable[Mapping[str, object]] | None, + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, + ) -> tuple[bool, ...]: + items: Final = tuple(tools or ()) + served: Final = await _served_mcp_path_names(items, served_names) + return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items) + + @staticmethod + async def _apply_toolset_permissions( + resolved_toolset_ids: list[str], + resolved_mcp_servers: list[str], + user_api_key_auth: "UserAPIKeyAuth", + ) -> "UserAPIKeyAuth": + """Apply resolved toolset permissions to user_api_key_auth and return updated auth.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + tool_permissions: Final = await global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=resolved_toolset_ids + ) + all_server_ids: Final = list(set(tool_permissions.keys()) | set(resolved_mcp_servers)) + existing_op: Final = user_api_key_auth.object_permission + if existing_op is not None: + merged_tool_perms: Final = dict(existing_op.mcp_tool_permissions or {}) + for server_id, tool_names in tool_permissions.items(): + existing_tools = merged_tool_perms.get(server_id, []) + merged_tool_perms[server_id] = list(set(existing_tools) | set(tool_names)) + updated_op = existing_op.model_copy( + update={ + "mcp_servers": all_server_ids, + "mcp_tool_permissions": merged_tool_perms, + "mcp_toolsets": [], + } + ) + else: + updated_op = LiteLLM_ObjectPermissionTable( + object_permission_id="toolset-scope", + mcp_servers=all_server_ids, + mcp_tool_permissions=tool_permissions, + ) + return user_api_key_auth.model_copy( + update={"object_permission": updated_op, "mcp_explicit_grants_only": True} + ) + except Exception as _e: + verbose_logger.debug("Could not apply toolset permissions: %s", _e) + return user_api_key_auth + + @staticmethod + @catalog_operation(global_manager) + async def _get_mcp_tools_from_manager( + user_api_key_auth: "UserAPIKeyAuth | None", + mcp_tools_with_litellm_proxy: Iterable[Mapping[str, object]] | None, + litellm_trace_id: str | None = None, + mcp_auth_header: str | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + request_tags: list[str] | None = None, + raw_headers: dict[str, str] | None = None, + granted_toolsets: "GrantedToolsetIds | None" = None, + ) -> tuple[list[MCPTool], list[str]]: + """ + Get available tools from the MCP server manager. + + Args: + user_api_key_auth: User authentication info for access control + mcp_tools_with_litellm_proxy: ToolParam objects with server_url starting with "litellm_proxy" + mcp_auth_header: Optional deprecated auth header for MCP servers + mcp_server_auth_headers: Optional server-specific auth headers (e.g. from x-mcp-{alias}-*) + + Returns: + List of MCP tools + List names of allowed MCP servers + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.server import ( + _get_allowed_mcp_servers_from_mcp_server_names, + _get_tools_from_mcp_servers, + ) + + mcp_servers: Final = [ + server_url.split("/")[-1] + for _tool in (mcp_tools_with_litellm_proxy or ()) + for server_url in (_tool.get("server_url", "") if isinstance(_tool, dict) else "",) + if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL_PREFIX) + ] + + # Resolve toolset names: collect all toolset IDs first, then apply their + # combined permissions in a single pass so multiple toolsets are unioned + # rather than the last one overwriting the others. + resolved_mcp_servers: Final[list[str]] = [] + resolved_toolset_ids: Final[list[str]] = [] + for name in mcp_servers: + if not global_mcp_server_manager.get_mcp_server_by_name(name): + try: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is not None: + toolset = await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) + if toolset is not None: + if user_api_key_auth is not None: + from litellm.proxy._experimental.mcp_server.ui_session_utils import ( + granted_toolset_ids, + ) + from litellm.proxy.management_endpoints.common_utils import ( + user_api_key_has_admin_view, + ) + + if not user_api_key_has_admin_view(user_api_key_auth) and toolset.toolset_id not in ( + await (granted_toolsets or granted_toolset_ids)(user_api_key_auth) + ): + verbose_logger.debug("Key does not have access to toolset '%s', skipping.", name) + continue + resolved_toolset_ids.append(toolset.toolset_id) + # Don't add to resolved_mcp_servers — toolset scope + # restricts via object_permission, not server name filter. + continue + except Exception as _e: + verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, _e) + resolved_mcp_servers.append(name) + + # Apply all resolved toolsets at once (union), avoiding permission overwrite. + if resolved_toolset_ids and user_api_key_auth is not None: + user_api_key_auth = await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions( + resolved_toolset_ids=resolved_toolset_ids, + resolved_mcp_servers=resolved_mcp_servers, + user_api_key_auth=user_api_key_auth, + ) + + # When toolsets were resolved we updated object_permission.mcp_servers to the + # full union (toolset server IDs + direct server names). Passing a name-based + # filter here would exclude those toolset server IDs (which are UUIDs, not + # names), so use None and let the auth object's mcp_servers do the filtering. + effective_server_filter: Final = None if resolved_toolset_ids else (resolved_mcp_servers or None) + + listing: Final = await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=effective_server_filter, + mcp_server_auth_headers=mcp_server_auth_headers, + log_list_tools_to_spendlogs=True, + list_tools_log_source="responses", + litellm_trace_id=litellm_trace_id, + request_tags=request_tags, + raw_headers=raw_headers, + ) + tools: Final = listing.tools + + allowed_mcp_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) + allowed_mcp_servers = global_mcp_server_manager.get_mcp_servers_from_ids(allowed_mcp_server_ids) + + allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( + mcp_servers=effective_server_filter, + allowed_mcp_servers=allowed_mcp_servers, + ) + + server_names: Final = [ + server_name + for server in allowed_mcp_servers + if server is not None + for server_name in ( + getattr(server, "server_name", None) or getattr(server, "alias", None) or getattr(server, "name", None), + ) + if isinstance(server_name, str) + ] + + return tools, server_names + + @staticmethod + def _deduplicate_mcp_tools( + mcp_tools: list[MCPTool], allowed_mcp_servers: list[str] + ) -> tuple[list[MCPTool], dict[str, str]]: + """ + Deduplicate MCP tools by name, keeping the first occurrence of each tool. + + Args: + mcp_tools: List of MCP tools that may contain duplicates + + Returns: + List of deduplicated MCP tools + The returned dictionary maps each tool_name to the server_name + """ + seen_names: Final[set[str]] = set() + deduplicated_tools: Final[list[MCPTool]] = [] + tool_server_map: Final[dict[str, str]] = {} + + for tool in mcp_tools: + if isinstance(tool, dict): + tool_name = tool.get("name") + else: + tool_name = getattr(tool, "name", None) + + if tool_name and tool_name not in seen_names: + seen_names.add(tool_name) + deduplicated_tools.append(tool) + if len(allowed_mcp_servers) == 1: + tool_server_map[tool_name] = allowed_mcp_servers[0] + else: + _, tool_server_map[tool_name] = split_server_prefix_from_name(tool_name) + + return deduplicated_tools, tool_server_map + + @staticmethod + def _filter_mcp_tools_by_allowed_tools( + mcp_tools: list[MCPTool], mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]] + ) -> list[MCPTool]: + """Filter MCP tools based on allowed_tools parameter from the original tool configs.""" + # Collect all allowed tool names from all MCP tool configs + allowed_tool_names: Final[set[str]] = set() + for tool_config in mcp_tools_with_litellm_proxy: + if isinstance(tool_config, dict) and "allowed_tools" in tool_config: + allowed_tools = tool_config.get("allowed_tools", []) + if isinstance(allowed_tools, list): + allowed_tool_names.update(allowed_tools) + + # If no allowed_tools specified, return all tools + if not allowed_tool_names: + return mcp_tools + + # Filter tools based on allowed names + filtered_tools: Final = [ + mcp_tool + for mcp_tool in mcp_tools + for tool_name in (mcp_tool.get("name") if isinstance(mcp_tool, dict) else getattr(mcp_tool, "name", None),) + if tool_name + and (tool_name in allowed_tool_names or split_server_prefix_from_name(tool_name)[0] in allowed_tool_names) + ] + + return filtered_tools + + @staticmethod + async def process_mcp_tools_to_openai_format( + user_api_key_auth: "UserAPIKeyAuth | None", + mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], + litellm_trace_id: str | None = None, + request_tags: list[str] | None = None, + ) -> tuple[list[Any], dict[str, str]]: + """ + Centralized method to process MCP tools through the complete pipeline. + + Args: + user_api_key_auth: User authentication info for access control + mcp_tools_with_litellm_proxy: ToolParam objects with server_url starting with "litellm_proxy" + litellm_trace_id: Optional trace ID for linking list_mcp_tools spend logs to parent request + + Returns: + List of tools in OpenAI format ready to be sent to the LLM + The returned dictionary maps each tool_name to the server_name + """ + ( + deduplicated_mcp_tools, + tool_server_map, + ) = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_without_openai_transform( + user_api_key_auth, + mcp_tools_with_litellm_proxy, + litellm_trace_id=litellm_trace_id, + request_tags=request_tags, + ) + + openai_tools: Final = LiteLLM_Proxy_MCP_Handler.transform_mcp_tools_to_openai(deduplicated_mcp_tools) + + return openai_tools, tool_server_map + + _process_mcp_tools_to_openai_format = process_mcp_tools_to_openai_format + + @staticmethod + async def process_mcp_tools_without_openai_transform( + user_api_key_auth: "UserAPIKeyAuth | None", + mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], + litellm_trace_id: str | None = None, + mcp_auth_header: str | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + request_tags: list[str] | None = None, + raw_headers: dict[str, str] | None = None, + ) -> tuple[list[MCPTool], dict[str, str]]: + """ + Process MCP tools through filtering and deduplication pipeline without OpenAI transformation. + This is useful for cases where we need the original MCP tool objects (e.g., for events). + + Args: + user_api_key_auth: User authentication info for access control + mcp_tools_with_litellm_proxy: ToolParam objects with server_url starting with "litellm_proxy" + mcp_auth_header: Optional deprecated auth header for MCP servers + mcp_server_auth_headers: Optional server-specific auth headers (e.g. from x-mcp-{alias}-*) + + Returns: + List of filtered and deduplicated MCP tools in their original format + """ + if not mcp_tools_with_litellm_proxy: + return [], {} + + typed_user_api_key_auth: Final[UserAPIKeyAuth | None] = user_api_key_auth + + # Step 1: Fetch MCP tools from manager + ( + mcp_tools_fetched, + allowed_mcp_servers, + ) = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( + user_api_key_auth=typed_user_api_key_auth, + mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, + litellm_trace_id=litellm_trace_id, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + request_tags=request_tags, + raw_headers=raw_headers, + ) + + # Step 2: Filter tools based on allowed_tools parameter + filtered_mcp_tools: Final = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools( + mcp_tools=mcp_tools_fetched, + mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, + ) + + # Step 3: Deduplicate tools after filtering + ( + deduplicated_mcp_tools, + tool_server_map, + ) = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools(filtered_mcp_tools, allowed_mcp_servers) + + return deduplicated_mcp_tools, tool_server_map + + _process_mcp_tools_without_openai_transform = process_mcp_tools_without_openai_transform + + @overload + @staticmethod + def transform_mcp_tools_to_openai( + mcp_tools: Sequence[MCPTool], + target_format: Literal["responses"] = ..., + ) -> list[FunctionToolParam]: ... + + @overload + @staticmethod + def transform_mcp_tools_to_openai( + mcp_tools: Sequence[MCPTool], + target_format: Literal["chat"], + ) -> list[ChatCompletionToolParam]: ... + + @staticmethod + def transform_mcp_tools_to_openai( + mcp_tools: Sequence[MCPTool], + target_format: Literal["responses", "chat"] = "responses", + ) -> Sequence[FunctionToolParam | ChatCompletionToolParam]: + """Transform MCP tools to OpenAI-compatible format.""" + from litellm.experimental_mcp_client.tools import ( + transform_mcp_tool_to_openai_responses_api_tool, + transform_mcp_tool_to_openai_tool, + ) + + openai_tools: Final = [ + transform_mcp_tool_to_openai_tool(mcp_tool) + if target_format == "chat" + else transform_mcp_tool_to_openai_responses_api_tool(mcp_tool) + for mcp_tool in mcp_tools + ] + + return openai_tools + + _transform_mcp_tools_to_openai = transform_mcp_tools_to_openai + + @staticmethod + def should_auto_execute_tools( + mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], + ) -> bool: + """Check if we should auto-execute tool calls. + + Auto-execution requires EVERY MCP reference to opt in with + ``require_approval="never"``. A single reference that requires approval + ("always", "manual", the object form, or an unset value) disables + auto-execution for the whole request. This fails closed: when an + approval-required reference shares a request with a "never" one, the + model's tool calls are returned to the caller instead of being run, so + an approval-gated tool can never be invoked without approval. Returns + False for an empty list. + """ + references: Final = list(mcp_tools_with_litellm_proxy or []) + if not references: + return False + for tool in references: + approval = ( + tool.get("require_approval") if isinstance(tool, dict) else getattr(tool, "require_approval", None) + ) + if approval != "never": + return False + return True + + _should_auto_execute_tools = should_auto_execute_tools + + @staticmethod + def extract_tool_calls_from_response(response: ResponsesAPIResponse) -> list[object]: + """Extract tool calls from the response output.""" + tool_calls: Final[list[object]] = [] + for output_item in response.output: + # Check if this is a function call output item + if isinstance(output_item, dict) and output_item.get("type") == "function_call": + tool_calls.append(output_item) + elif hasattr(output_item, "type") and getattr(output_item, "type") == "function_call": + # Handle pydantic model case + tool_calls.append(output_item) + + return tool_calls + + _extract_tool_calls_from_response = extract_tool_calls_from_response + + @staticmethod + def extract_tool_calls_from_chat_response(response: ModelResponse) -> list[object]: + """Extract tool calls from a chat completion response.""" + tool_calls: Final[list[object]] = [] + + try: + for choice in response.choices: + message: Message | None = getattr(choice, "message", None) + if message is None: + continue + tool_call_entries: ( + Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None + ) = getattr(message, "tool_calls", None) + if tool_call_entries: + for tool_call in tool_call_entries: + if hasattr(tool_call, "model_dump"): + tool_calls.append(tool_call.model_dump()) + else: + tool_calls.append(tool_call) + except Exception: + verbose_logger.exception("Failed to extract tool calls from chat completion response") + + return tool_calls + + _extract_tool_calls_from_chat_response = extract_tool_calls_from_chat_response + + @staticmethod + def extract_tool_call_details( + tool_call: object, + ) -> tuple[str | None, str | None, str | None]: + """Extract tool name, arguments, and call_id from a tool call.""" + if isinstance(tool_call, dict): + tool_call_id = tool_call.get("call_id") or tool_call.get("id") + + # OpenAI chat completions wrap tool info under a `function` block + function_block: Final = tool_call.get("function") + if isinstance(function_block, dict): + tool_name = function_block.get("name") + tool_arguments = function_block.get("arguments") + else: + tool_name = tool_call.get("name") + # Anthropic tool_use blocks carry the arguments under `input` + tool_arguments = tool_call.get("arguments") + if tool_arguments is None: + tool_arguments = tool_call.get("input") + else: + tool_call_id = getattr(tool_call, "call_id", None) or getattr(tool_call, "id", None) + + function_obj: Final[object] = getattr(tool_call, "function", None) + if function_obj is not None: + tool_name = getattr(function_obj, "name", None) + tool_arguments = getattr(function_obj, "arguments", None) + else: + tool_name = getattr(tool_call, "name", None) + tool_arguments = getattr(tool_call, "arguments", None) + if tool_arguments is None: + tool_arguments = getattr(tool_call, "input", None) + + return tool_name, tool_arguments, tool_call_id + + _extract_tool_call_details = extract_tool_call_details + + @staticmethod + def _parse_tool_arguments(tool_arguments: str | None) -> dict[str, object]: + """Parse tool arguments, handling both string and dict formats.""" + import json + + if isinstance(tool_arguments, str): + try: + return json.loads(tool_arguments) + except json.JSONDecodeError: + return {} + else: + return tool_arguments or {} + + @staticmethod + def _parse_mcp_result(result: "CallToolResult") -> str: + """Parse MCP tool call result and extract meaningful content.""" + if not result or not hasattr(result, "content") or not result.content: + return "Tool executed successfully" + + # Import MCP content types for isinstance checks + try: + from mcp.types import EmbeddedResource, ImageContent, TextContent + except ImportError: + # Fallback to generic handling if MCP types not available + return "Tool executed successfully" + + text_parts: Final = [ + str(content_item.text) for content_item in result.content if isinstance(content_item, TextContent) + ] + other_content_types: Final = [ + "Image" + if isinstance(content_item, ImageContent) + else "EmbeddedResource" + if isinstance(content_item, EmbeddedResource) + else type(content_item).__name__ + for content_item in result.content + if not isinstance(content_item, TextContent) + ] + + # Combine text parts if any + result_text = " ".join(text_parts) if text_parts else "" + + # Add info about other content types + if other_content_types: + other_info: Final = f"[Generated {', '.join(other_content_types)}]" + result_text = f"{result_text} {other_info}".strip() + + return result_text or "Tool executed successfully" + + @staticmethod + @catalog_operation(global_manager) + async def execute_tool_calls( + tool_server_map: dict[str, str], + tool_calls: Sequence[object], + user_api_key_auth: "UserAPIKeyAuth | None", + mcp_auth_header: str | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + litellm_call_id: str | None = None, + litellm_trace_id: str | None = None, + request_tags: list[str] | None = None, + guardrail_context: Mapping[str, object] | None = None, + served_tools: Sequence[MCPTool] | None = None, + ) -> list[MCPToolResult]: + """Execute tool calls and return results.""" + from fastapi import HTTPException + + from litellm._uuid import uuid + from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.server import ( + _resolve_display_name_to_original, + ) + from litellm.proxy.proxy_server import proxy_logging_obj + + tool_results: Final[list[MCPToolResult]] = [] + tool_call_id: str | None = None + rules_obj: Final = Rules() + logging_safe_headers: Final = logging_safe_mcp_headers(raw_headers) + typed_user_api_key_auth: Final[UserAPIKeyAuth | None] = user_api_key_auth + for tool_call in tool_calls: + logging_request_data: dict[str, object] = {} + tool_name: str | None = None + try: + ( + tool_name, + tool_arguments, + tool_call_id, + ) = LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call) + + if not tool_name: + verbose_logger.warning("Tool call missing name: %s", tool_call) + continue + + parsed_arguments = LiteLLM_Proxy_MCP_Handler._parse_tool_arguments(tool_arguments) + + # Import here to avoid circular import + from litellm.proxy.proxy_server import proxy_logging_obj + + 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) + + start_time = datetime.now() + logging_input: Sequence[Mapping[str, object]] = [ + { + "role": "tool", + "content": { + "tool_name": sanitized_tool_name, + "arguments": parsed_arguments, + }, + } + ] + tool_logging_call_id = litellm_call_id or str(uuid.uuid4()) + logging_metadata: dict[str, object] = { + "tool_call_id": tool_call_id, + "tool_name": sanitized_tool_name, + "server_name": server_name, + "headers": logging_safe_headers, + } + logging_request_data = { + "model": f"MCP: {tool_name}", + "metadata": logging_metadata, + "input": logging_input, + "call_type": CallTypes.call_mcp_tool.value, + "litellm_call_id": tool_logging_call_id, + # Add proxy_server_request with arguments for callback logging + "proxy_server_request": { + "url": "/mcp/tools/call", + "method": "POST", + "headers": logging_safe_headers, + "body": { + "name": sanitized_tool_name, + "arguments": parsed_arguments, + }, + }, + } + if litellm_trace_id: + logging_request_data["litellm_trace_id"] = litellm_trace_id + if request_tags: + logging_metadata["tags"] = request_tags + if typed_user_api_key_auth is not None: + from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + ) + + LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data=logging_request_data, + user_api_key_dict=typed_user_api_key_auth, + _metadata_variable_name="metadata", + ) + user_identifier = getattr(typed_user_api_key_auth, "end_user_id", None) or getattr( + typed_user_api_key_auth, "user_id", None + ) + if user_identifier: + logging_request_data["user"] = user_identifier + + litellm_logging_obj: LiteLLMLoggingObj | None = None + try: + litellm_logging_obj, _ = function_setup( + original_function="call_mcp_tool", + rules_obj=rules_obj, + start_time=start_time, + **logging_request_data, + ) + except Exception as logging_error: + verbose_logger.debug( + "Failed to initialize logging for MCP tool call %s: %s", + tool_name, + logging_error, + ) + litellm_logging_obj = None + + logging_request_data["litellm_logging_obj"] = litellm_logging_obj + logging_request_data["arguments"] = parsed_arguments + + if litellm_logging_obj: + try: + litellm_logging_obj.pre_call( + input=logging_input, + api_key="", + ) + except Exception: + verbose_logger.exception("Failed to run pre_call for MCP tool logging") + + standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = { + "name": sanitized_tool_name, + "arguments": parsed_arguments, + "namespaced_tool_name": tool_name, + } + if mcp_server: + mcp_info = mcp_server.mcp_info or {} + standard_logging_mcp_tool_call["mcp_server_name"] = ( + mcp_info.get("server_name") or getattr(mcp_server, "server_name", None) or server_name + ) + logo_url = mcp_info.get("logo_url") + if logo_url: + standard_logging_mcp_tool_call["mcp_server_logo_url"] = logo_url + cost_info = mcp_info.get("mcp_server_cost_info") + if cost_info: + standard_logging_mcp_tool_call["mcp_server_cost_info"] = cost_info + + if litellm_logging_obj: + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call + litellm_logging_obj.model = f"MCP: {tool_name}" + litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value + + result = await global_mcp_server_manager.call_tool( + server_name=server_name, + name=sanitized_tool_name, + arguments=parsed_arguments, + user_api_key_auth=typed_user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + proxy_logging_obj=proxy_logging_obj, + litellm_logging_obj=litellm_logging_obj, + guardrail_context=guardrail_context, + listed_tool=( + next((tool for tool in served_tools if tool.name == tool_name), None) + if served_tools is not None + else ... + ), + ) + + if proxy_logging_obj: + result = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data=( + litellm_logging_obj.model_call_details + if litellm_logging_obj + else {"mcp_tool_name": tool_name} + ), + user_api_key_dict=typed_user_api_key_auth, + ) + + if litellm_logging_obj: + try: + litellm_logging_obj.post_call(original_response=result) + result = await litellm_logging_obj.async_post_mcp_tool_call_hook( + kwargs=litellm_logging_obj.model_call_details, + response_obj=result, + start_time=start_time, + end_time=datetime.now(), + ) + except Exception: + verbose_logger.exception("Failed to run post-call logging for MCP tool call %s", tool_name) + try: + await litellm_logging_obj.async_success_handler( + result=result, + start_time=start_time, + end_time=datetime.now(), + ) + except Exception: + verbose_logger.exception("Failed to log MCP tool call success for %s", tool_name) + + # Format result for inclusion in response + result_text = LiteLLM_Proxy_MCP_Handler._parse_mcp_result(result) + tool_results.append( + { + "tool_call_id": tool_call_id, + "result": result_text, + "name": tool_name, + } + ) + + except BlockedPiiEntityError as e: + await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( + proxy_logging_obj=proxy_logging_obj, + user_api_key_auth=typed_user_api_key_auth, + request_data=logging_request_data, + error=e, + ) + verbose_logger.error("BlockedPiiEntityError in MCP tool call: %s", e) + error_message = f"Tool call blocked: PII entity '{getattr(e, 'entity_type', 'unknown')}' detected by guardrail '{getattr(e, 'guardrail_name', 'unknown')}'. {e}" + tool_results.append( + { + "tool_call_id": tool_call_id, + "result": error_message, + "name": tool_name, + } + ) + except GuardrailRaisedException as e: + await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( + proxy_logging_obj=proxy_logging_obj, + user_api_key_auth=typed_user_api_key_auth, + request_data=logging_request_data, + error=e, + ) + verbose_logger.error("GuardrailRaisedException in MCP tool call: %s", e) + error_message = ( + f"Tool call blocked: Guardrail '{getattr(e, 'guardrail_name', 'unknown')}' violation. {e}" + ) + tool_results.append( + { + "tool_call_id": tool_call_id, + "result": error_message, + "name": tool_name, + } + ) + except HTTPException as e: + await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( + proxy_logging_obj=proxy_logging_obj, + user_api_key_auth=typed_user_api_key_auth, + request_data=logging_request_data, + error=e, + ) + verbose_logger.error("HTTPException in MCP tool call: %s", e) + error_message = f"Tool call failed: {str(e.detail) if hasattr(e, 'detail') else str(e)}" + tool_results.append( + { + "tool_call_id": tool_call_id, + "result": error_message, + "name": tool_name, + } + ) + except Exception as e: + await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( + proxy_logging_obj=proxy_logging_obj, + user_api_key_auth=typed_user_api_key_auth, + request_data=logging_request_data, + error=e, + ) + verbose_logger.exception("Error executing MCP tool call: %s", e) + tool_results.append( + { + "tool_call_id": tool_call_id, + "result": f"Error executing tool: {e}", + "name": tool_name, + } + ) + + return tool_results + + _execute_tool_calls = execute_tool_calls + + @staticmethod + def create_follow_up_messages_for_chat( + original_messages: list[object], + response: ModelResponse, + tool_results: Sequence[Mapping[str, object]], + ) -> Sequence[Mapping[str, object]]: + """Create follow-up chat messages that include tool execution results.""" + from copy import deepcopy + + from litellm.utils import convert_list_message_to_dict + + follow_up_messages: list[dict[str, object]] = convert_list_message_to_dict(deepcopy(original_messages)) + + if not follow_up_messages: + follow_up_messages = [] + + message_to_append: dict | None = None + try: + first_choice: Final = response.choices[0] + if isinstance(first_choice, Choices) and getattr(first_choice, "message", None): + message_to_append = first_choice.message.model_dump(exclude_none=True) + # Ensure tool_calls have arguments field (required by OpenAI API) + if message_to_append.get("tool_calls"): + for tool_call in message_to_append["tool_calls"]: + if isinstance(tool_call, dict) and "function" in tool_call: + if "arguments" not in tool_call["function"]: + tool_call["function"]["arguments"] = "{}" + except Exception: + verbose_logger.exception("Failed to convert assistant message for MCP flow") + + if message_to_append: + follow_up_messages.append(message_to_append) + + for tool_result in tool_results: + follow_up_messages.append( + { + "role": "tool", + "tool_call_id": tool_result.get("tool_call_id"), + "name": tool_result.get("name"), + "content": tool_result.get("result", ""), + } + ) + + return follow_up_messages + + _create_follow_up_messages_for_chat = create_follow_up_messages_for_chat + + @staticmethod + def is_persistence_disabled(call_params: Mapping[str, object]) -> bool: + """store=false means the provider kept nothing, so the follow-up call cannot chain on a response id.""" + return call_params.get("store") is False + + _is_persistence_disabled = is_persistence_disabled + + @staticmethod + def create_follow_up_input( + response: ResponsesAPIResponse, + tool_results: Sequence[Mapping[str, object]], + original_input: str | ResponseInputParam | None = None, + preserve_reasoning: bool = False, + ) -> list[object]: + """Create follow-up input with tool results in proper format.""" + follow_up_input: Final[list[object]] = [] + + # Add original user input if available to maintain conversation context + if original_input: + if isinstance(original_input, str): + follow_up_input.append({"type": "message", "role": "user", "content": original_input}) + elif isinstance(original_input, list): + follow_up_input.extend(original_input) + else: + follow_up_input.append(original_input) + + # Add the assistant message with function calls + assistant_message_content: Final[list[object]] = [] + turn_items: Final[list[Mapping[str, object]]] = [] + + for output_item in response.output: + if not isinstance(output_item, dict) and hasattr(output_item, "model_dump"): + output_item = output_item.model_dump(exclude_none=True) + + if isinstance(output_item, dict): + if output_item.get("type") == "function_call": + call_id = output_item.get("call_id") or output_item.get("id") + name = output_item.get("name") + arguments = output_item.get("arguments") + + # Only add if we have required fields + if call_id and name: + turn_items.append( + { + "type": "function_call", + "call_id": call_id, + "name": name, + "arguments": arguments, + } + ) + elif output_item.get("type") == "reasoning" and preserve_reasoning: + turn_items.append(output_item) + elif output_item.get("type") == "message": + # Extract content from message + content = output_item.get("content", []) + if isinstance(content, list): + assistant_message_content.extend(content) + else: + assistant_message_content.append(content) + + # Add assistant message only if there's actual content (not empty) + # For example, gemini requires that function call turns come immediately after user turns, + # so we should not add empty assistant messages + if assistant_message_content: + follow_up_input.append( + { + "type": "message", + "role": "assistant", + "content": assistant_message_content, + } + ) + + follow_up_input.extend(turn_items) + + # Add tool results (function call outputs) + for tool_result in tool_results: + follow_up_input.append( + { + "type": "function_call_output", + "call_id": tool_result["tool_call_id"], + "output": tool_result["result"], + } + ) + + return follow_up_input + + _create_follow_up_input = create_follow_up_input + + @staticmethod + async def make_follow_up_call( + follow_up_input: list[Any], + model: str, + all_tools: Sequence[ResponsesToolParam] | None, + response_id: str | None, + **call_params: Any, + ) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: + """Make follow-up response API call with tool results.""" + return await aresponses( + input=follow_up_input, + model=model, + tools=all_tools, # Keep tools for potential future calls + previous_response_id=response_id, # Link to previous response + **call_params, + ) + + _make_follow_up_call = make_follow_up_call + + @staticmethod + async def _log_mcp_tool_failure( + *, + proxy_logging_obj: Optional["ProxyLogging"], + user_api_key_auth: "UserAPIKeyAuth | None", + request_data: dict[str, object], + error: Exception, + ) -> None: + """Log MCP tool failures via proxy logging hooks.""" + + if proxy_logging_obj is None or user_api_key_auth is None: + return + + try: + traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + await proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=error, + user_api_key_dict=user_api_key_auth, + route="/responses/mcp/call_tool", + traceback_str=traceback_str, + ) + except Exception: + verbose_logger.exception("Failed to log MCP tool call failure") + + @staticmethod + def _create_mcp_streaming_response( + input: str | ResponseInputParam, + model: str, + all_tools: Sequence[object] | None, + mcp_tools_with_litellm_proxy: list[Mapping[str, object]], + mcp_discovery_events: list[ResponsesAPIStreamingResponse], + call_params: Mapping[str, object], + previous_response_id: str | None, + tool_server_map: dict[str, str], + served_tools: Sequence[MCPTool] | None = None, + **kwargs, + ) -> Any: + """ + Create MCP enhanced streaming response that handles the full MCP workflow. + + This creates a streaming iterator that: + 1. Immediately emits MCP discovery events + 2. Makes the LLM call and streams the response + 3. Handles tool execution and follow-up calls + """ + from litellm.responses.mcp.mcp_streaming_iterator import ( + MCPEnhancedStreamingIterator, + ) + + # Build the complete request parameters by merging all sources + request_params: Final = LiteLLM_Proxy_MCP_Handler._build_request_params( + input=input, + model=model, + all_tools=all_tools, + call_params=call_params, + previous_response_id=previous_response_id, + **kwargs, + ) + + # Create the enhanced streaming iterator that will handle everything + return MCPEnhancedStreamingIterator( + base_iterator=None, # Will be created internally + mcp_events=mcp_discovery_events, # Pre-generated MCP discovery events + tool_server_map=tool_server_map, + served_tools=served_tools, + mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, + 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, + ) + + @staticmethod + def _build_request_params( + input: str | ResponseInputParam, + model: str, + all_tools: Sequence[object] | None, + call_params: Mapping[str, object], + previous_response_id: str | None, + **kwargs: object, + ) -> dict[str, object]: + """ + Build a clean request parameters dictionary for MCP streaming. + + Combines input, model, tools with call_params and additional kwargs + in a clean, maintainable way. + """ + # Start with the core required parameters + request_params: Final[dict[str, object]] = { + "input": input, + "model": model, + "tools": all_tools, + } + + # Add previous_response_id if provided + if previous_response_id is not None: + request_params["previous_response_id"] = previous_response_id + + # Merge in all call_params (which contains most of the API parameters) + request_params.update(call_params) + + # Merge in any additional kwargs + request_params.update(kwargs) + + return request_params + + @staticmethod + def create_tool_execution_events( + tool_calls: Sequence[object], tool_results: Sequence[MCPToolResult] + ) -> list[ResponsesAPIStreamingResponse]: + """ + Create MCP tool execution events for streaming. + + Args: + tool_calls: List of tool calls from the LLM response + tool_results: List of tool execution results + + Returns: + List of MCP tool execution events for streaming + """ + from litellm._uuid import uuid + from litellm.responses.mcp.mcp_streaming_iterator import create_mcp_call_events + + tool_execution_events: Final[list[ResponsesAPIStreamingResponse]] = [] + + # Create events for each tool execution + for tool_result in tool_results: + tool_call_id = tool_result.get("tool_call_id", "unknown") + result_text = tool_result.get("result", "") + + # Extract tool name and arguments from tool calls + tool_name = "unknown" + tool_arguments = "{}" + for tool_call in tool_calls: + ( + name, + args, + call_id, + ) = LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call) + if call_id == tool_call_id: + tool_name = name or "unknown" + tool_arguments = args or "{}" + break + + execution_events = create_mcp_call_events( + tool_name=tool_name, + tool_call_id=tool_call_id, + arguments=tool_arguments, # Use actual arguments + result=result_text, + base_item_id=f"mcp_{uuid.uuid4().hex[:8]}", # Unique ID for each tool call + sequence_start=len(tool_execution_events) + 1, + ) + tool_execution_events.extend(execution_events) + + return tool_execution_events + + _create_tool_execution_events = create_tool_execution_events + + @staticmethod + def prepare_initial_call_params(call_params: Mapping[str, object], should_auto_execute: bool) -> dict[str, Any]: + """ + Prepare call parameters for the initial LLM call. + + For auto-execute scenarios, we need to disable streaming for the initial call + so we can process the tool calls before streaming the final response. + """ + initial_params: Final = dict(call_params) + + if should_auto_execute: + # Disable streaming for initial call when auto-executing tools + initial_params["stream"] = False + + return initial_params + + _prepare_initial_call_params = prepare_initial_call_params + + @staticmethod + def prepare_follow_up_call_params( + call_params: Mapping[str, object], original_stream_setting: bool + ) -> dict[str, Any]: + """ + Prepare call parameters for the follow-up LLM call after tool execution. + + Restores the original streaming setting and removes tool_choice since + we're now providing tool results, not requesting tool calls. + """ + follow_up_params: Final = dict(call_params) + + # Restore original streaming setting for follow-up call + follow_up_params["stream"] = original_stream_setting + + # Remove tool_choice since we're providing results, not requesting tool calls + follow_up_params.pop("tool_choice", None) + + return follow_up_params + + _prepare_follow_up_call_params = prepare_follow_up_call_params + + @staticmethod + def add_mcp_output_elements_to_response( + response: ResponsesAPIResponse, + mcp_tools_fetched: Sequence[object], + tool_results: Sequence[Mapping[str, object]], + ) -> ResponsesAPIResponse: + """Add custom output elements to the final response for MCP tool execution.""" + # Import the required classes for creating output items + import json + + from litellm._uuid import uuid + from litellm.types.responses.main import GenericResponseOutputItem, OutputText + + # Create output element for initial MCP tools + mcp_tools_output: Final = GenericResponseOutputItem( + type="mcp_tools_fetched", + id=f"mcp_tools_{uuid.uuid4().hex[:8]}", + status="completed", + role="system", + content=[ + OutputText( + type="output_text", + text=json.dumps(mcp_tools_fetched, indent=2, default=str), + annotations=[], + ) + ], + ) + + # Create output element for tool execution results + tool_results_output: Final = GenericResponseOutputItem( + type="tool_execution_results", + id=f"tool_results_{uuid.uuid4().hex[:8]}", + status="completed", + role="system", + content=[ + OutputText( + type="output_text", + text=json.dumps(tool_results, indent=2, default=str), + annotations=[], + ) + ], + ) + + # Add the new output elements to the response + response.output.append(mcp_tools_output.model_dump()) + response.output.append(tool_results_output.model_dump()) + + return response + + _add_mcp_output_elements_to_response = add_mcp_output_elements_to_response diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index d3e3509b8f0..1aba4dee970 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,1408 +1,17 @@ -import re -import traceback -from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence -from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload - -from openai.types.chat import ChatCompletionToolParam -from openai.types.responses.function_tool_param import FunctionToolParam - -from litellm._logging import verbose_logger -from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager -from litellm.proxy._experimental.mcp_server.utils import ( - iter_known_server_prefixes, - logging_safe_mcp_headers, - split_server_prefix_from_name, - strip_known_server_prefix, -) -from litellm.responses.main import aresponses # noqa: TID251 # inner call must skip the MCP gateway that invoked it -from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator -from litellm.types.llms.openai import ( - ResponseInputParam, - ResponsesAPIResponse, - ResponsesAPIStreamingResponse, -) -from litellm.types.llms.openai import ToolParam as ResponsesToolParam -from litellm.types.mcp_server.mcp_server_manager import MCPServer -from litellm.types.utils import ( - CallTypes, - ChatCompletionMessageCustomToolCall, - ChatCompletionMessageToolCall, - Choices, - Message, - ModelResponse, - StandardLoggingMCPToolCall, -) -from litellm.utils import Rules, function_setup +import importlib +from typing import TYPE_CHECKING, Final if TYPE_CHECKING: - from mcp.types import CallToolResult - from mcp.types import Tool as MCPTool - - from litellm.proxy._experimental.mcp_server.ui_session_utils import GrantedToolsetIds - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.utils import ProxyLogging -else: - MCPTool = Any - -# NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling -ToolParam: TypeAlias = Mapping[str, object] -SplitTools: TypeAlias = tuple[list[ToolParam], list[Any]] - - -class MCPToolResult(TypedDict): - tool_call_id: str | None - result: str - name: str | None - - -LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" -LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" - -_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") - - -def _mcp_server_url(tool: ToolParam) -> str | None: - if not isinstance(tool, dict) or tool.get("type") != "mcp": - return None - server_url: Final = tool.get("server_url") - return server_url if isinstance(server_url, str) else None - - -def _names_gateway_explicitly(tool: ToolParam) -> bool: - return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) - - -def _proxy_path_mcp_name(tool: ToolParam) -> 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) - return None if match is None else match.group(1) - - -def _registered_mcp_servers() -> Collection[MCPServer]: - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, + from litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler import ( + LITELLM_PROXY_MCP_SERVER_URL_PREFIX, + LiteLLM_Proxy_MCP_Handler, + MCPToolResult, ) - return global_mcp_server_manager.get_registry().values() +__all__ = ("LITELLM_PROXY_MCP_SERVER_URL_PREFIX", "LiteLLM_Proxy_MCP_Handler", "MCPToolResult") +_RELOCATED_MODULE: Final = "litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler" -def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool: - requested: Final = name.lower() - return any( - requested in (known.lower() for known in (*iter_known_server_prefixes(server), server.name)) - or name in (server.access_groups or ()) - for server in servers - ) - -async def _toolset_exists(name: str) -> bool: - try: - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - return False - return await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is not None - except Exception as e: - verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, e) - return False - - -@catalog_operation(global_manager) -async def _gateway_served_names( - names: Collection[str], - servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers, - toolset_exists: Callable[[str], Awaitable[bool]] = _toolset_exists, -) -> frozenset[str]: - registered: Final = tuple(servers()) if names else () - return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)]) - - -async def _served_mcp_path_names( - tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] -) -> frozenset[str]: - names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None) - return await served_names(names) if names else frozenset[str]() - - -class LiteLLM_Proxy_MCP_Handler: - """ - Helper class with static methods for MCP integration with Responses API. - - This handles when a user passes mcp server_url="litellm_proxy" in their tools. - """ - - @staticmethod - def get_parent_request_tags(kwargs: dict[str, Any] | None) -> list[str]: - """Tags from the parent LLM request, using the same extraction logic as standard logging (incl. User-Agent).""" - if not kwargs: - return [] - from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup - - litellm_params: Final = kwargs.get("litellm_params") or kwargs - proxy_server_request = litellm_params.get("proxy_server_request") or kwargs.get("proxy_server_request") or {} - return StandardLoggingPayloadSetup.get_request_tags( - litellm_params=litellm_params, - proxy_server_request=proxy_server_request, - ) - - _get_parent_request_tags = get_parent_request_tags - - @staticmethod - def should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool: - """True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in - /mcp/. `_split_mcp_tools` then settles which of the latter the gateway actually serves.""" - return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ()) - - _should_use_litellm_mcp_gateway = should_use_litellm_mcp_gateway - - @staticmethod - def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools: - items: Final = tuple(tools or ()) - gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)] - other_tools: Final[list[ToolParam]] = [tool for tool in items if not _names_gateway_explicitly(tool)] - return gateway_tools, other_tools - - @staticmethod - async def split_mcp_tools( - tools: Iterable[Mapping[str, object]] | None, - served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, - ) -> SplitTools: - items: Final = tuple(tools or ()) - served: Final = await _served_mcp_path_names(items, served_names) - return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools( - [ - {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} - if (name := _proxy_path_mcp_name(tool)) in served - else tool - for tool in items - ] - ) - - _split_mcp_tools = split_mcp_tools - - @staticmethod - @catalog_operation(global_manager) - async def routes_through_gateway( - tools: Iterable[Mapping[str, object]] | None, - served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, - ) -> tuple[bool, ...]: - items: Final = tuple(tools or ()) - served: Final = await _served_mcp_path_names(items, served_names) - return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items) - - @staticmethod - async def _apply_toolset_permissions( - resolved_toolset_ids: list[str], - resolved_mcp_servers: list[str], - user_api_key_auth: "UserAPIKeyAuth", - ) -> "UserAPIKeyAuth": - """Apply resolved toolset permissions to user_api_key_auth and return updated auth.""" - from litellm.proxy._types import LiteLLM_ObjectPermissionTable - - try: - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - - tool_permissions: Final = await global_mcp_server_manager.resolve_toolset_tool_permissions( - toolset_ids=resolved_toolset_ids - ) - all_server_ids: Final = list(set(tool_permissions.keys()) | set(resolved_mcp_servers)) - existing_op: Final = user_api_key_auth.object_permission - if existing_op is not None: - merged_tool_perms: Final = dict(existing_op.mcp_tool_permissions or {}) - for server_id, tool_names in tool_permissions.items(): - existing_tools = merged_tool_perms.get(server_id, []) - merged_tool_perms[server_id] = list(set(existing_tools) | set(tool_names)) - updated_op = existing_op.model_copy( - update={ - "mcp_servers": all_server_ids, - "mcp_tool_permissions": merged_tool_perms, - "mcp_toolsets": [], - } - ) - else: - updated_op = LiteLLM_ObjectPermissionTable( - object_permission_id="toolset-scope", - mcp_servers=all_server_ids, - mcp_tool_permissions=tool_permissions, - ) - return user_api_key_auth.model_copy( - update={"object_permission": updated_op, "mcp_explicit_grants_only": True} - ) - except Exception as _e: - verbose_logger.debug("Could not apply toolset permissions: %s", _e) - return user_api_key_auth - - @staticmethod - @catalog_operation(global_manager) - async def _get_mcp_tools_from_manager( - user_api_key_auth: "UserAPIKeyAuth | None", - mcp_tools_with_litellm_proxy: Iterable[Mapping[str, object]] | None, - litellm_trace_id: str | None = None, - mcp_auth_header: str | None = None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, - request_tags: list[str] | None = None, - raw_headers: dict[str, str] | None = None, - granted_toolsets: "GrantedToolsetIds | None" = None, - ) -> tuple[list[MCPTool], list[str]]: - """ - Get available tools from the MCP server manager. - - Args: - user_api_key_auth: User authentication info for access control - mcp_tools_with_litellm_proxy: ToolParam objects with server_url starting with "litellm_proxy" - mcp_auth_header: Optional deprecated auth header for MCP servers - mcp_server_auth_headers: Optional server-specific auth headers (e.g. from x-mcp-{alias}-*) - - Returns: - List of MCP tools - List names of allowed MCP servers - """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.proxy._experimental.mcp_server.server import ( - _get_allowed_mcp_servers_from_mcp_server_names, - _get_tools_from_mcp_servers, - ) - - mcp_servers: Final = [ - server_url.split("/")[-1] - for _tool in (mcp_tools_with_litellm_proxy or ()) - for server_url in (_tool.get("server_url", "") if isinstance(_tool, dict) else "",) - if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL_PREFIX) - ] - - # Resolve toolset names: collect all toolset IDs first, then apply their - # combined permissions in a single pass so multiple toolsets are unioned - # rather than the last one overwriting the others. - resolved_mcp_servers: Final[list[str]] = [] - resolved_toolset_ids: Final[list[str]] = [] - for name in mcp_servers: - if not global_mcp_server_manager.get_mcp_server_by_name(name): - try: - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is not None: - toolset = await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) - if toolset is not None: - if user_api_key_auth is not None: - from litellm.proxy._experimental.mcp_server.ui_session_utils import ( - granted_toolset_ids, - ) - from litellm.proxy.management_endpoints.common_utils import ( - user_api_key_has_admin_view, - ) - - if not user_api_key_has_admin_view(user_api_key_auth) and toolset.toolset_id not in ( - await (granted_toolsets or granted_toolset_ids)(user_api_key_auth) - ): - verbose_logger.debug("Key does not have access to toolset '%s', skipping.", name) - continue - resolved_toolset_ids.append(toolset.toolset_id) - # Don't add to resolved_mcp_servers — toolset scope - # restricts via object_permission, not server name filter. - continue - except Exception as _e: - verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, _e) - resolved_mcp_servers.append(name) - - # Apply all resolved toolsets at once (union), avoiding permission overwrite. - if resolved_toolset_ids and user_api_key_auth is not None: - user_api_key_auth = await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions( - resolved_toolset_ids=resolved_toolset_ids, - resolved_mcp_servers=resolved_mcp_servers, - user_api_key_auth=user_api_key_auth, - ) - - # When toolsets were resolved we updated object_permission.mcp_servers to the - # full union (toolset server IDs + direct server names). Passing a name-based - # filter here would exclude those toolset server IDs (which are UUIDs, not - # names), so use None and let the auth object's mcp_servers do the filtering. - effective_server_filter: Final = None if resolved_toolset_ids else (resolved_mcp_servers or None) - - listing: Final = await _get_tools_from_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=effective_server_filter, - mcp_server_auth_headers=mcp_server_auth_headers, - log_list_tools_to_spendlogs=True, - list_tools_log_source="responses", - litellm_trace_id=litellm_trace_id, - request_tags=request_tags, - raw_headers=raw_headers, - ) - tools: Final = listing.tools - - allowed_mcp_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) - allowed_mcp_servers = global_mcp_server_manager.get_mcp_servers_from_ids(allowed_mcp_server_ids) - - allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( - mcp_servers=effective_server_filter, - allowed_mcp_servers=allowed_mcp_servers, - ) - - server_names: Final = [ - server_name - for server in allowed_mcp_servers - if server is not None - for server_name in ( - getattr(server, "server_name", None) or getattr(server, "alias", None) or getattr(server, "name", None), - ) - if isinstance(server_name, str) - ] - - return tools, server_names - - @staticmethod - def _deduplicate_mcp_tools( - mcp_tools: list[MCPTool], allowed_mcp_servers: list[str] - ) -> tuple[list[MCPTool], dict[str, str]]: - """ - Deduplicate MCP tools by name, keeping the first occurrence of each tool. - - Args: - mcp_tools: List of MCP tools that may contain duplicates - - Returns: - List of deduplicated MCP tools - The returned dictionary maps each tool_name to the server_name - """ - seen_names: Final[set[str]] = set() - deduplicated_tools: Final[list[MCPTool]] = [] - tool_server_map: Final[dict[str, str]] = {} - - for tool in mcp_tools: - if isinstance(tool, dict): - tool_name = tool.get("name") - else: - tool_name = getattr(tool, "name", None) - - if tool_name and tool_name not in seen_names: - seen_names.add(tool_name) - deduplicated_tools.append(tool) - if len(allowed_mcp_servers) == 1: - tool_server_map[tool_name] = allowed_mcp_servers[0] - else: - _, tool_server_map[tool_name] = split_server_prefix_from_name(tool_name) - - return deduplicated_tools, tool_server_map - - @staticmethod - def _filter_mcp_tools_by_allowed_tools( - mcp_tools: list[MCPTool], mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]] - ) -> list[MCPTool]: - """Filter MCP tools based on allowed_tools parameter from the original tool configs.""" - # Collect all allowed tool names from all MCP tool configs - allowed_tool_names: Final[set[str]] = set() - for tool_config in mcp_tools_with_litellm_proxy: - if isinstance(tool_config, dict) and "allowed_tools" in tool_config: - allowed_tools = tool_config.get("allowed_tools", []) - if isinstance(allowed_tools, list): - allowed_tool_names.update(allowed_tools) - - # If no allowed_tools specified, return all tools - if not allowed_tool_names: - return mcp_tools - - # Filter tools based on allowed names - filtered_tools: Final = [ - mcp_tool - for mcp_tool in mcp_tools - for tool_name in (mcp_tool.get("name") if isinstance(mcp_tool, dict) else getattr(mcp_tool, "name", None),) - if tool_name - and (tool_name in allowed_tool_names or split_server_prefix_from_name(tool_name)[0] in allowed_tool_names) - ] - - return filtered_tools - - @staticmethod - async def process_mcp_tools_to_openai_format( - user_api_key_auth: "UserAPIKeyAuth | None", - mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], - litellm_trace_id: str | None = None, - request_tags: list[str] | None = None, - ) -> tuple[list[Any], dict[str, str]]: - """ - Centralized method to process MCP tools through the complete pipeline. - - Args: - user_api_key_auth: User authentication info for access control - mcp_tools_with_litellm_proxy: ToolParam objects with server_url starting with "litellm_proxy" - litellm_trace_id: Optional trace ID for linking list_mcp_tools spend logs to parent request - - Returns: - List of tools in OpenAI format ready to be sent to the LLM - The returned dictionary maps each tool_name to the server_name - """ - ( - deduplicated_mcp_tools, - tool_server_map, - ) = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_without_openai_transform( - user_api_key_auth, - mcp_tools_with_litellm_proxy, - litellm_trace_id=litellm_trace_id, - request_tags=request_tags, - ) - - openai_tools: Final = LiteLLM_Proxy_MCP_Handler.transform_mcp_tools_to_openai(deduplicated_mcp_tools) - - return openai_tools, tool_server_map - - _process_mcp_tools_to_openai_format = process_mcp_tools_to_openai_format - - @staticmethod - async def process_mcp_tools_without_openai_transform( - user_api_key_auth: "UserAPIKeyAuth | None", - mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], - litellm_trace_id: str | None = None, - mcp_auth_header: str | None = None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, - request_tags: list[str] | None = None, - raw_headers: dict[str, str] | None = None, - ) -> tuple[list[MCPTool], dict[str, str]]: - """ - Process MCP tools through filtering and deduplication pipeline without OpenAI transformation. - This is useful for cases where we need the original MCP tool objects (e.g., for events). - - Args: - user_api_key_auth: User authentication info for access control - mcp_tools_with_litellm_proxy: ToolParam objects with server_url starting with "litellm_proxy" - mcp_auth_header: Optional deprecated auth header for MCP servers - mcp_server_auth_headers: Optional server-specific auth headers (e.g. from x-mcp-{alias}-*) - - Returns: - List of filtered and deduplicated MCP tools in their original format - """ - if not mcp_tools_with_litellm_proxy: - return [], {} - - typed_user_api_key_auth: Final[UserAPIKeyAuth | None] = user_api_key_auth - - # Step 1: Fetch MCP tools from manager - ( - mcp_tools_fetched, - allowed_mcp_servers, - ) = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( - user_api_key_auth=typed_user_api_key_auth, - mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, - litellm_trace_id=litellm_trace_id, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - request_tags=request_tags, - raw_headers=raw_headers, - ) - - # Step 2: Filter tools based on allowed_tools parameter - filtered_mcp_tools: Final = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools( - mcp_tools=mcp_tools_fetched, - mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, - ) - - # Step 3: Deduplicate tools after filtering - ( - deduplicated_mcp_tools, - tool_server_map, - ) = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools(filtered_mcp_tools, allowed_mcp_servers) - - return deduplicated_mcp_tools, tool_server_map - - _process_mcp_tools_without_openai_transform = process_mcp_tools_without_openai_transform - - @overload - @staticmethod - def transform_mcp_tools_to_openai( - mcp_tools: Sequence[MCPTool], - target_format: Literal["responses"] = ..., - ) -> list[FunctionToolParam]: ... - - @overload - @staticmethod - def transform_mcp_tools_to_openai( - mcp_tools: Sequence[MCPTool], - target_format: Literal["chat"], - ) -> list[ChatCompletionToolParam]: ... - - @staticmethod - def transform_mcp_tools_to_openai( - mcp_tools: Sequence[MCPTool], - target_format: Literal["responses", "chat"] = "responses", - ) -> Sequence[FunctionToolParam | ChatCompletionToolParam]: - """Transform MCP tools to OpenAI-compatible format.""" - from litellm.experimental_mcp_client.tools import ( - transform_mcp_tool_to_openai_responses_api_tool, - transform_mcp_tool_to_openai_tool, - ) - - openai_tools: Final = [ - transform_mcp_tool_to_openai_tool(mcp_tool) - if target_format == "chat" - else transform_mcp_tool_to_openai_responses_api_tool(mcp_tool) - for mcp_tool in mcp_tools - ] - - return openai_tools - - _transform_mcp_tools_to_openai = transform_mcp_tools_to_openai - - @staticmethod - def should_auto_execute_tools( - mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], - ) -> bool: - """Check if we should auto-execute tool calls. - - Auto-execution requires EVERY MCP reference to opt in with - ``require_approval="never"``. A single reference that requires approval - ("always", "manual", the object form, or an unset value) disables - auto-execution for the whole request. This fails closed: when an - approval-required reference shares a request with a "never" one, the - model's tool calls are returned to the caller instead of being run, so - an approval-gated tool can never be invoked without approval. Returns - False for an empty list. - """ - references: Final = list(mcp_tools_with_litellm_proxy or []) - if not references: - return False - for tool in references: - approval = ( - tool.get("require_approval") if isinstance(tool, dict) else getattr(tool, "require_approval", None) - ) - if approval != "never": - return False - return True - - _should_auto_execute_tools = should_auto_execute_tools - - @staticmethod - def extract_tool_calls_from_response(response: ResponsesAPIResponse) -> list[object]: - """Extract tool calls from the response output.""" - tool_calls: Final[list[object]] = [] - for output_item in response.output: - # Check if this is a function call output item - if isinstance(output_item, dict) and output_item.get("type") == "function_call": - tool_calls.append(output_item) - elif hasattr(output_item, "type") and getattr(output_item, "type") == "function_call": - # Handle pydantic model case - tool_calls.append(output_item) - - return tool_calls - - _extract_tool_calls_from_response = extract_tool_calls_from_response - - @staticmethod - def extract_tool_calls_from_chat_response(response: ModelResponse) -> list[object]: - """Extract tool calls from a chat completion response.""" - tool_calls: Final[list[object]] = [] - - try: - for choice in response.choices: - message: Message | None = getattr(choice, "message", None) - if message is None: - continue - tool_call_entries: ( - Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None - ) = getattr(message, "tool_calls", None) - if tool_call_entries: - for tool_call in tool_call_entries: - if hasattr(tool_call, "model_dump"): - tool_calls.append(tool_call.model_dump()) - else: - tool_calls.append(tool_call) - except Exception: - verbose_logger.exception("Failed to extract tool calls from chat completion response") - - return tool_calls - - _extract_tool_calls_from_chat_response = extract_tool_calls_from_chat_response - - @staticmethod - def extract_tool_call_details( - tool_call: object, - ) -> tuple[str | None, str | None, str | None]: - """Extract tool name, arguments, and call_id from a tool call.""" - if isinstance(tool_call, dict): - tool_call_id = tool_call.get("call_id") or tool_call.get("id") - - # OpenAI chat completions wrap tool info under a `function` block - function_block: Final = tool_call.get("function") - if isinstance(function_block, dict): - tool_name = function_block.get("name") - tool_arguments = function_block.get("arguments") - else: - tool_name = tool_call.get("name") - # Anthropic tool_use blocks carry the arguments under `input` - tool_arguments = tool_call.get("arguments") - if tool_arguments is None: - tool_arguments = tool_call.get("input") - else: - tool_call_id = getattr(tool_call, "call_id", None) or getattr(tool_call, "id", None) - - function_obj: Final[object] = getattr(tool_call, "function", None) - if function_obj is not None: - tool_name = getattr(function_obj, "name", None) - tool_arguments = getattr(function_obj, "arguments", None) - else: - tool_name = getattr(tool_call, "name", None) - tool_arguments = getattr(tool_call, "arguments", None) - if tool_arguments is None: - tool_arguments = getattr(tool_call, "input", None) - - return tool_name, tool_arguments, tool_call_id - - _extract_tool_call_details = extract_tool_call_details - - @staticmethod - def _parse_tool_arguments(tool_arguments: str | None) -> dict[str, object]: - """Parse tool arguments, handling both string and dict formats.""" - import json - - if isinstance(tool_arguments, str): - try: - return json.loads(tool_arguments) - except json.JSONDecodeError: - return {} - else: - return tool_arguments or {} - - @staticmethod - def _parse_mcp_result(result: "CallToolResult") -> str: - """Parse MCP tool call result and extract meaningful content.""" - if not result or not hasattr(result, "content") or not result.content: - return "Tool executed successfully" - - # Import MCP content types for isinstance checks - try: - from mcp.types import EmbeddedResource, ImageContent, TextContent - except ImportError: - # Fallback to generic handling if MCP types not available - return "Tool executed successfully" - - text_parts: Final = [ - str(content_item.text) for content_item in result.content if isinstance(content_item, TextContent) - ] - other_content_types: Final = [ - "Image" - if isinstance(content_item, ImageContent) - else "EmbeddedResource" - if isinstance(content_item, EmbeddedResource) - else type(content_item).__name__ - for content_item in result.content - if not isinstance(content_item, TextContent) - ] - - # Combine text parts if any - result_text = " ".join(text_parts) if text_parts else "" - - # Add info about other content types - if other_content_types: - other_info: Final = f"[Generated {', '.join(other_content_types)}]" - result_text = f"{result_text} {other_info}".strip() - - return result_text or "Tool executed successfully" - - @staticmethod - @catalog_operation(global_manager) - async def execute_tool_calls( - tool_server_map: dict[str, str], - tool_calls: Sequence[object], - user_api_key_auth: "UserAPIKeyAuth | None", - mcp_auth_header: str | None = None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, - oauth2_headers: dict[str, str] | None = None, - raw_headers: dict[str, str] | None = None, - litellm_call_id: str | None = None, - litellm_trace_id: str | None = None, - request_tags: list[str] | None = None, - guardrail_context: Mapping[str, object] | None = None, - served_tools: Sequence[MCPTool] | None = None, - ) -> list[MCPToolResult]: - """Execute tool calls and return results.""" - from fastapi import HTTPException - - from litellm._uuid import uuid - from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.proxy._experimental.mcp_server.server import ( - _resolve_display_name_to_original, - ) - from litellm.proxy.proxy_server import proxy_logging_obj - - tool_results: Final[list[MCPToolResult]] = [] - tool_call_id: str | None = None - rules_obj: Final = Rules() - logging_safe_headers: Final = logging_safe_mcp_headers(raw_headers) - typed_user_api_key_auth: Final[UserAPIKeyAuth | None] = user_api_key_auth - for tool_call in tool_calls: - logging_request_data: dict[str, object] = {} - tool_name: str | None = None - try: - ( - tool_name, - tool_arguments, - tool_call_id, - ) = LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call) - - if not tool_name: - verbose_logger.warning("Tool call missing name: %s", tool_call) - continue - - parsed_arguments = LiteLLM_Proxy_MCP_Handler._parse_tool_arguments(tool_arguments) - - # Import here to avoid circular import - from litellm.proxy.proxy_server import proxy_logging_obj - - 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) - - start_time = datetime.now() - logging_input: Sequence[Mapping[str, object]] = [ - { - "role": "tool", - "content": { - "tool_name": sanitized_tool_name, - "arguments": parsed_arguments, - }, - } - ] - tool_logging_call_id = litellm_call_id or str(uuid.uuid4()) - logging_metadata: dict[str, object] = { - "tool_call_id": tool_call_id, - "tool_name": sanitized_tool_name, - "server_name": server_name, - "headers": logging_safe_headers, - } - logging_request_data = { - "model": f"MCP: {tool_name}", - "metadata": logging_metadata, - "input": logging_input, - "call_type": CallTypes.call_mcp_tool.value, - "litellm_call_id": tool_logging_call_id, - # Add proxy_server_request with arguments for callback logging - "proxy_server_request": { - "url": "/mcp/tools/call", - "method": "POST", - "headers": logging_safe_headers, - "body": { - "name": sanitized_tool_name, - "arguments": parsed_arguments, - }, - }, - } - if litellm_trace_id: - logging_request_data["litellm_trace_id"] = litellm_trace_id - if request_tags: - logging_metadata["tags"] = request_tags - if typed_user_api_key_auth is not None: - from litellm.proxy.litellm_pre_call_utils import ( - LiteLLMProxyRequestSetup, - ) - - LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( - data=logging_request_data, - user_api_key_dict=typed_user_api_key_auth, - _metadata_variable_name="metadata", - ) - user_identifier = getattr(typed_user_api_key_auth, "end_user_id", None) or getattr( - typed_user_api_key_auth, "user_id", None - ) - if user_identifier: - logging_request_data["user"] = user_identifier - - litellm_logging_obj: LiteLLMLoggingObj | None = None - try: - litellm_logging_obj, _ = function_setup( - original_function="call_mcp_tool", - rules_obj=rules_obj, - start_time=start_time, - **logging_request_data, - ) - except Exception as logging_error: - verbose_logger.debug( - "Failed to initialize logging for MCP tool call %s: %s", - tool_name, - logging_error, - ) - litellm_logging_obj = None - - logging_request_data["litellm_logging_obj"] = litellm_logging_obj - logging_request_data["arguments"] = parsed_arguments - - if litellm_logging_obj: - try: - litellm_logging_obj.pre_call( - input=logging_input, - api_key="", - ) - except Exception: - verbose_logger.exception("Failed to run pre_call for MCP tool logging") - - standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = { - "name": sanitized_tool_name, - "arguments": parsed_arguments, - "namespaced_tool_name": tool_name, - } - if mcp_server: - mcp_info = mcp_server.mcp_info or {} - standard_logging_mcp_tool_call["mcp_server_name"] = ( - mcp_info.get("server_name") or getattr(mcp_server, "server_name", None) or server_name - ) - logo_url = mcp_info.get("logo_url") - if logo_url: - standard_logging_mcp_tool_call["mcp_server_logo_url"] = logo_url - cost_info = mcp_info.get("mcp_server_cost_info") - if cost_info: - standard_logging_mcp_tool_call["mcp_server_cost_info"] = cost_info - - if litellm_logging_obj: - litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call - litellm_logging_obj.model = f"MCP: {tool_name}" - litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value - - result = await global_mcp_server_manager.call_tool( - server_name=server_name, - name=sanitized_tool_name, - arguments=parsed_arguments, - user_api_key_auth=typed_user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - proxy_logging_obj=proxy_logging_obj, - litellm_logging_obj=litellm_logging_obj, - guardrail_context=guardrail_context, - listed_tool=( - next((tool for tool in served_tools if tool.name == tool_name), None) - if served_tools is not None - else ... - ), - ) - - if proxy_logging_obj: - result = await proxy_logging_obj.post_mcp_call_hook( - response=result, - request_data=( - litellm_logging_obj.model_call_details - if litellm_logging_obj - else {"mcp_tool_name": tool_name} - ), - user_api_key_dict=typed_user_api_key_auth, - ) - - if litellm_logging_obj: - try: - litellm_logging_obj.post_call(original_response=result) - result = await litellm_logging_obj.async_post_mcp_tool_call_hook( - kwargs=litellm_logging_obj.model_call_details, - response_obj=result, - start_time=start_time, - end_time=datetime.now(), - ) - except Exception: - verbose_logger.exception("Failed to run post-call logging for MCP tool call %s", tool_name) - try: - await litellm_logging_obj.async_success_handler( - result=result, - start_time=start_time, - end_time=datetime.now(), - ) - except Exception: - verbose_logger.exception("Failed to log MCP tool call success for %s", tool_name) - - # Format result for inclusion in response - result_text = LiteLLM_Proxy_MCP_Handler._parse_mcp_result(result) - tool_results.append( - { - "tool_call_id": tool_call_id, - "result": result_text, - "name": tool_name, - } - ) - - except BlockedPiiEntityError as e: - await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( - proxy_logging_obj=proxy_logging_obj, - user_api_key_auth=typed_user_api_key_auth, - request_data=logging_request_data, - error=e, - ) - verbose_logger.error("BlockedPiiEntityError in MCP tool call: %s", e) - error_message = f"Tool call blocked: PII entity '{getattr(e, 'entity_type', 'unknown')}' detected by guardrail '{getattr(e, 'guardrail_name', 'unknown')}'. {e}" - tool_results.append( - { - "tool_call_id": tool_call_id, - "result": error_message, - "name": tool_name, - } - ) - except GuardrailRaisedException as e: - await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( - proxy_logging_obj=proxy_logging_obj, - user_api_key_auth=typed_user_api_key_auth, - request_data=logging_request_data, - error=e, - ) - verbose_logger.error("GuardrailRaisedException in MCP tool call: %s", e) - error_message = ( - f"Tool call blocked: Guardrail '{getattr(e, 'guardrail_name', 'unknown')}' violation. {e}" - ) - tool_results.append( - { - "tool_call_id": tool_call_id, - "result": error_message, - "name": tool_name, - } - ) - except HTTPException as e: - await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( - proxy_logging_obj=proxy_logging_obj, - user_api_key_auth=typed_user_api_key_auth, - request_data=logging_request_data, - error=e, - ) - verbose_logger.error("HTTPException in MCP tool call: %s", e) - error_message = f"Tool call failed: {str(e.detail) if hasattr(e, 'detail') else str(e)}" - tool_results.append( - { - "tool_call_id": tool_call_id, - "result": error_message, - "name": tool_name, - } - ) - except Exception as e: - await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( - proxy_logging_obj=proxy_logging_obj, - user_api_key_auth=typed_user_api_key_auth, - request_data=logging_request_data, - error=e, - ) - verbose_logger.exception("Error executing MCP tool call: %s", e) - tool_results.append( - { - "tool_call_id": tool_call_id, - "result": f"Error executing tool: {e}", - "name": tool_name, - } - ) - - return tool_results - - _execute_tool_calls = execute_tool_calls - - @staticmethod - def create_follow_up_messages_for_chat( - original_messages: list[object], - response: ModelResponse, - tool_results: Sequence[Mapping[str, object]], - ) -> Sequence[Mapping[str, object]]: - """Create follow-up chat messages that include tool execution results.""" - from copy import deepcopy - - from litellm.utils import convert_list_message_to_dict - - follow_up_messages: list[dict[str, object]] = convert_list_message_to_dict(deepcopy(original_messages)) - - if not follow_up_messages: - follow_up_messages = [] - - message_to_append: dict | None = None - try: - first_choice: Final = response.choices[0] - if isinstance(first_choice, Choices) and getattr(first_choice, "message", None): - message_to_append = first_choice.message.model_dump(exclude_none=True) - # Ensure tool_calls have arguments field (required by OpenAI API) - if message_to_append.get("tool_calls"): - for tool_call in message_to_append["tool_calls"]: - if isinstance(tool_call, dict) and "function" in tool_call: - if "arguments" not in tool_call["function"]: - tool_call["function"]["arguments"] = "{}" - except Exception: - verbose_logger.exception("Failed to convert assistant message for MCP flow") - - if message_to_append: - follow_up_messages.append(message_to_append) - - for tool_result in tool_results: - follow_up_messages.append( - { - "role": "tool", - "tool_call_id": tool_result.get("tool_call_id"), - "name": tool_result.get("name"), - "content": tool_result.get("result", ""), - } - ) - - return follow_up_messages - - _create_follow_up_messages_for_chat = create_follow_up_messages_for_chat - - @staticmethod - def is_persistence_disabled(call_params: Mapping[str, object]) -> bool: - """store=false means the provider kept nothing, so the follow-up call cannot chain on a response id.""" - return call_params.get("store") is False - - _is_persistence_disabled = is_persistence_disabled - - @staticmethod - def create_follow_up_input( - response: ResponsesAPIResponse, - tool_results: Sequence[Mapping[str, object]], - original_input: str | ResponseInputParam | None = None, - preserve_reasoning: bool = False, - ) -> list[object]: - """Create follow-up input with tool results in proper format.""" - follow_up_input: Final[list[object]] = [] - - # Add original user input if available to maintain conversation context - if original_input: - if isinstance(original_input, str): - follow_up_input.append({"type": "message", "role": "user", "content": original_input}) - elif isinstance(original_input, list): - follow_up_input.extend(original_input) - else: - follow_up_input.append(original_input) - - # Add the assistant message with function calls - assistant_message_content: Final[list[object]] = [] - turn_items: Final[list[Mapping[str, object]]] = [] - - for output_item in response.output: - if not isinstance(output_item, dict) and hasattr(output_item, "model_dump"): - output_item = output_item.model_dump(exclude_none=True) - - if isinstance(output_item, dict): - if output_item.get("type") == "function_call": - call_id = output_item.get("call_id") or output_item.get("id") - name = output_item.get("name") - arguments = output_item.get("arguments") - - # Only add if we have required fields - if call_id and name: - turn_items.append( - { - "type": "function_call", - "call_id": call_id, - "name": name, - "arguments": arguments, - } - ) - elif output_item.get("type") == "reasoning" and preserve_reasoning: - turn_items.append(output_item) - elif output_item.get("type") == "message": - # Extract content from message - content = output_item.get("content", []) - if isinstance(content, list): - assistant_message_content.extend(content) - else: - assistant_message_content.append(content) - - # Add assistant message only if there's actual content (not empty) - # For example, gemini requires that function call turns come immediately after user turns, - # so we should not add empty assistant messages - if assistant_message_content: - follow_up_input.append( - { - "type": "message", - "role": "assistant", - "content": assistant_message_content, - } - ) - - follow_up_input.extend(turn_items) - - # Add tool results (function call outputs) - for tool_result in tool_results: - follow_up_input.append( - { - "type": "function_call_output", - "call_id": tool_result["tool_call_id"], - "output": tool_result["result"], - } - ) - - return follow_up_input - - _create_follow_up_input = create_follow_up_input - - @staticmethod - async def make_follow_up_call( - follow_up_input: list[Any], - model: str, - all_tools: Sequence[ResponsesToolParam] | None, - response_id: str | None, - **call_params: Any, - ) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: - """Make follow-up response API call with tool results.""" - return await aresponses( - input=follow_up_input, - model=model, - tools=all_tools, # Keep tools for potential future calls - previous_response_id=response_id, # Link to previous response - **call_params, - ) - - _make_follow_up_call = make_follow_up_call - - @staticmethod - async def _log_mcp_tool_failure( - *, - proxy_logging_obj: Optional["ProxyLogging"], - user_api_key_auth: "UserAPIKeyAuth | None", - request_data: dict[str, object], - error: Exception, - ) -> None: - """Log MCP tool failures via proxy logging hooks.""" - - if proxy_logging_obj is None or user_api_key_auth is None: - return - - try: - traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) - await proxy_logging_obj.post_call_failure_hook( - request_data=request_data, - original_exception=error, - user_api_key_dict=user_api_key_auth, - route="/responses/mcp/call_tool", - traceback_str=traceback_str, - ) - except Exception: - verbose_logger.exception("Failed to log MCP tool call failure") - - @staticmethod - def _create_mcp_streaming_response( - input: str | ResponseInputParam, - model: str, - all_tools: Sequence[object] | None, - mcp_tools_with_litellm_proxy: list[Mapping[str, object]], - mcp_discovery_events: list[ResponsesAPIStreamingResponse], - call_params: Mapping[str, object], - previous_response_id: str | None, - tool_server_map: dict[str, str], - served_tools: Sequence[MCPTool] | None = None, - **kwargs, - ) -> Any: - """ - Create MCP enhanced streaming response that handles the full MCP workflow. - - This creates a streaming iterator that: - 1. Immediately emits MCP discovery events - 2. Makes the LLM call and streams the response - 3. Handles tool execution and follow-up calls - """ - from litellm.responses.mcp.mcp_streaming_iterator import ( - MCPEnhancedStreamingIterator, - ) - - # Build the complete request parameters by merging all sources - request_params: Final = LiteLLM_Proxy_MCP_Handler._build_request_params( - input=input, - model=model, - all_tools=all_tools, - call_params=call_params, - previous_response_id=previous_response_id, - **kwargs, - ) - - # Create the enhanced streaming iterator that will handle everything - return MCPEnhancedStreamingIterator( - base_iterator=None, # Will be created internally - mcp_events=mcp_discovery_events, # Pre-generated MCP discovery events - tool_server_map=tool_server_map, - served_tools=served_tools, - mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, - 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, - ) - - @staticmethod - def _build_request_params( - input: str | ResponseInputParam, - model: str, - all_tools: Sequence[object] | None, - call_params: Mapping[str, object], - previous_response_id: str | None, - **kwargs: object, - ) -> dict[str, object]: - """ - Build a clean request parameters dictionary for MCP streaming. - - Combines input, model, tools with call_params and additional kwargs - in a clean, maintainable way. - """ - # Start with the core required parameters - request_params: Final[dict[str, object]] = { - "input": input, - "model": model, - "tools": all_tools, - } - - # Add previous_response_id if provided - if previous_response_id is not None: - request_params["previous_response_id"] = previous_response_id - - # Merge in all call_params (which contains most of the API parameters) - request_params.update(call_params) - - # Merge in any additional kwargs - request_params.update(kwargs) - - return request_params - - @staticmethod - def create_tool_execution_events( - tool_calls: Sequence[object], tool_results: Sequence[MCPToolResult] - ) -> list[ResponsesAPIStreamingResponse]: - """ - Create MCP tool execution events for streaming. - - Args: - tool_calls: List of tool calls from the LLM response - tool_results: List of tool execution results - - Returns: - List of MCP tool execution events for streaming - """ - from litellm._uuid import uuid - from litellm.responses.mcp.mcp_streaming_iterator import create_mcp_call_events - - tool_execution_events: Final[list[ResponsesAPIStreamingResponse]] = [] - - # Create events for each tool execution - for tool_result in tool_results: - tool_call_id = tool_result.get("tool_call_id", "unknown") - result_text = tool_result.get("result", "") - - # Extract tool name and arguments from tool calls - tool_name = "unknown" - tool_arguments = "{}" - for tool_call in tool_calls: - ( - name, - args, - call_id, - ) = LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call) - if call_id == tool_call_id: - tool_name = name or "unknown" - tool_arguments = args or "{}" - break - - execution_events = create_mcp_call_events( - tool_name=tool_name, - tool_call_id=tool_call_id, - arguments=tool_arguments, # Use actual arguments - result=result_text, - base_item_id=f"mcp_{uuid.uuid4().hex[:8]}", # Unique ID for each tool call - sequence_start=len(tool_execution_events) + 1, - ) - tool_execution_events.extend(execution_events) - - return tool_execution_events - - _create_tool_execution_events = create_tool_execution_events - - @staticmethod - def prepare_initial_call_params(call_params: Mapping[str, object], should_auto_execute: bool) -> dict[str, Any]: - """ - Prepare call parameters for the initial LLM call. - - For auto-execute scenarios, we need to disable streaming for the initial call - so we can process the tool calls before streaming the final response. - """ - initial_params: Final = dict(call_params) - - if should_auto_execute: - # Disable streaming for initial call when auto-executing tools - initial_params["stream"] = False - - return initial_params - - _prepare_initial_call_params = prepare_initial_call_params - - @staticmethod - def prepare_follow_up_call_params( - call_params: Mapping[str, object], original_stream_setting: bool - ) -> dict[str, Any]: - """ - Prepare call parameters for the follow-up LLM call after tool execution. - - Restores the original streaming setting and removes tool_choice since - we're now providing tool results, not requesting tool calls. - """ - follow_up_params: Final = dict(call_params) - - # Restore original streaming setting for follow-up call - follow_up_params["stream"] = original_stream_setting - - # Remove tool_choice since we're providing results, not requesting tool calls - follow_up_params.pop("tool_choice", None) - - return follow_up_params - - _prepare_follow_up_call_params = prepare_follow_up_call_params - - @staticmethod - def add_mcp_output_elements_to_response( - response: ResponsesAPIResponse, - mcp_tools_fetched: Sequence[object], - tool_results: Sequence[Mapping[str, object]], - ) -> ResponsesAPIResponse: - """Add custom output elements to the final response for MCP tool execution.""" - # Import the required classes for creating output items - import json - - from litellm._uuid import uuid - from litellm.types.responses.main import GenericResponseOutputItem, OutputText - - # Create output element for initial MCP tools - mcp_tools_output: Final = GenericResponseOutputItem( - type="mcp_tools_fetched", - id=f"mcp_tools_{uuid.uuid4().hex[:8]}", - status="completed", - role="system", - content=[ - OutputText( - type="output_text", - text=json.dumps(mcp_tools_fetched, indent=2, default=str), - annotations=[], - ) - ], - ) - - # Create output element for tool execution results - tool_results_output: Final = GenericResponseOutputItem( - type="tool_execution_results", - id=f"tool_results_{uuid.uuid4().hex[:8]}", - status="completed", - role="system", - content=[ - OutputText( - type="output_text", - text=json.dumps(tool_results, indent=2, default=str), - annotations=[], - ) - ], - ) - - # Add the new output elements to the response - response.output.append(mcp_tools_output.model_dump()) - response.output.append(tool_results_output.model_dump()) - - return response - - _add_mcp_output_elements_to_response = add_mcp_output_elements_to_response +def __getattr__(name: str) -> object: + return getattr(importlib.import_module(_RELOCATED_MODULE), name) diff --git a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py similarity index 97% rename from tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py rename to tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py index f43e15cb342..eaf7570d427 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py @@ -24,8 +24,8 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ListedTool from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging from litellm.responses import main as responses_main -from litellm.responses.mcp import litellm_proxy_mcp_handler as mcp_handler_module -from litellm.responses.mcp.litellm_proxy_mcp_handler import ( +from litellm.proxy._experimental.mcp_server import litellm_proxy_mcp_handler as mcp_handler_module +from litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) from litellm.types.guardrails import GuardrailEventHooks @@ -439,7 +439,7 @@ async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_functio # NOTE: Don't patch via dotted string path here because `litellm.responses` # is a function attribute on the `litellm` package (shadowing the submodule), # which breaks monkeypatch's importpath resolution. - handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + handler_module = importlib.import_module("litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) tool_name = "deepwiki-read_wiki_structure" @@ -475,7 +475,7 @@ async def test_execute_tool_calls_threads_logging_obj_into_call_tool(monkeypatch sentinel_logging_obj.async_post_mcp_tool_call_hook = AsyncMock() sentinel_logging_obj.async_success_handler = AsyncMock() - handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + handler_module = importlib.import_module("litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler") monkeypatch.setattr( handler_module, "function_setup", @@ -524,7 +524,7 @@ async def test_execute_tool_calls_applies_post_call_hook_content(monkeypatch): return_value=CallToolResult(content=[TextContent(type="text", text="[REDACTED]")], is_error=True) ) logging_obj.async_success_handler = AsyncMock() - handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + handler_module = importlib.import_module("litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None)) tool_name = "deepwiki-read_wiki_structure" @@ -559,7 +559,7 @@ async def test_execute_tool_calls_returns_proxy_result_without_logging(monkeypat "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", fake_manager, ) - handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + handler_module = importlib.import_module("litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (None, None)) tool_name = "deepwiki-read_wiki_structure" @@ -597,7 +597,7 @@ async def test_execute_tool_calls_passes_logging_details_to_proxy_hook(monkeypat logging_obj.model_call_details = {"request_id": "request-1"} logging_obj.async_post_mcp_tool_call_hook = AsyncMock(return_value=result) logging_obj.async_success_handler = AsyncMock() - handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + handler_module = importlib.import_module("litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None)) tool_name = "deepwiki-read_wiki_structure" @@ -640,7 +640,7 @@ async def test_execute_tool_calls_continues_when_post_call_logging_fails(monkeyp logging_obj.async_success_handler = AsyncMock( side_effect=RuntimeError("success logging failed") if failure_stage == "success_handler" else None ) - handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + handler_module = importlib.import_module("litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None)) tool_name = "deepwiki-read_wiki_structure" @@ -766,7 +766,7 @@ async def test_execute_tool_calls_exposes_sanitized_client_headers_to_logging(mo captured.update(kwargs) return None, None - handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + handler_module = importlib.import_module("litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) tool_name = "deepwiki-read_wiki_structure" @@ -792,7 +792,7 @@ async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monk captured.update(kwargs) return None, None - handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + handler_module = importlib.import_module("litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) tool_name = "deepwiki-read_wiki_structure" @@ -930,7 +930,7 @@ FUNCTION_TOOL = {"type": "function", "name": "get_weather", "parameters": {}} @pytest.mark.asyncio async def test_gateway_served_names_matches_alias_server_name_name_access_group_and_toolset(): - from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names + from litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler import _gateway_served_names servers = ( _registered("id-1", "github-name", alias="github", server_name="github-server", access_groups=["prod-group"]), @@ -952,7 +952,7 @@ async def test_gateway_served_names_matches_alias_server_name_name_access_group_ @pytest.mark.asyncio async def test_gateway_served_names_matches_server_id_short_prefix_and_alias_case_like_the_gateway(): from litellm.proxy._experimental.mcp_server.utils import compute_short_server_prefix - from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names + from litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler import _gateway_served_names server_id = "0b9ae4ca-1bd2-4faa-b183-7dd812597e3b" short_prefix = compute_short_server_prefix(server_id) diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index 27477f24d92..390760872c1 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -1070,7 +1070,7 @@ async def test_execute_tool_calls_sets_proxy_server_request_arguments(monkeypatc # exported as a function on the top-level `litellm` package, which can confuse # pytest's dotted-path resolver. mcp_handler_module = importlib.import_module( - "litellm.responses.mcp.litellm_proxy_mcp_handler" + "litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler" ) monkeypatch.setattr(mcp_handler_module, "function_setup", mock_function_setup) monkeypatch.setattr( @@ -1455,7 +1455,7 @@ async def test_request_selected_mcp_guardrail_blocks_before_upstream(monkeypatch ), ] if logging_failure: - from litellm.responses.mcp import litellm_proxy_mcp_handler + from litellm.proxy._experimental.mcp_server import litellm_proxy_mcp_handler def fail_logging(*args, **kwargs): raise RuntimeError("logging initialization failed") monkeypatch.setattr(litellm_proxy_mcp_handler, "function_setup", fail_logging) From 9f88b873c0d11a5523b60fb15fd3db8caf5c3994 Mon Sep 17 00:00:00 2001 From: nate-berri Date: Sat, 10 Oct 2026 09:39:32 +0000 Subject: [PATCH 2/9] test(mcp): cover the lazy litellm_proxy_mcp_handler shim Part of LIT-9423 --- .../mcp/test_litellm_proxy_mcp_handler.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100644 tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py diff --git a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py new file mode 100644 index 00000000000..f4e2656c068 --- /dev/null +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -0,0 +1,28 @@ +from typing import Final + +import pytest + +from litellm.proxy._experimental.mcp_server import litellm_proxy_mcp_handler as relocated +from litellm.responses.mcp import litellm_proxy_mcp_handler as legacy +from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter + + +def test_legacy_import_leaves_relocated_proxy_module_unloaded() -> None: + result: Final = run_child_interpreter( + "import sys, litellm.responses.mcp.litellm_proxy_mcp_handler\n" + "print('litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler' in sys.modules)", + timeout=120, + ) + + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "False" + + +@pytest.mark.parametrize("name", legacy.__all__) +def test_legacy_path_returns_the_relocated_objects(name: str) -> None: + assert getattr(legacy, name) is getattr(relocated, name) + + +def test_legacy_path_still_raises_for_unknown_attributes() -> None: + with pytest.raises(AttributeError): + getattr(legacy, "NotAHandlerAttribute") From 4bd6bb6a0218dd2fbe3aef0b18e991af5c251b88 Mon Sep 17 00:00:00 2001 From: nate-berri Date: Sat, 10 Oct 2026 10:00:24 +0000 Subject: [PATCH 3/9] ci(mcp): run litellm_proxy_mcp_handler tests in the mcp-elicitation shard Part of LIT-9423 --- .github/workflows/test-unit.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 42b418b3ee0..98a33a89177 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -570,10 +570,12 @@ jobs: tests/unit/experimental_mcp_client/test_mcp_client.py tests/unit/proxy/_experimental/mcp_server/test_capabilities.py tests/unit/proxy/_experimental/mcp_server/test_interactions.py + tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py tests/unit/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py tests/unit/proxy/_experimental/mcp_server/test_operations.py + tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py tests/integration/mcp/test_interactions.py python-version: "3.12" workers: 2 From 36bc2d483bf4499058c7f05fdcb6899fa44c9d29 Mon Sep 17 00:00:00 2001 From: nate-berri Date: Sat, 10 Oct 2026 10:34:49 +0000 Subject: [PATCH 4/9] test(mcp): cover more of the relocated litellm_proxy_mcp_handler Part of LIT-9423 --- .../test_litellm_proxy_mcp_handler.py | 381 +++++++++++++++++- 1 file changed, 380 insertions(+), 1 deletion(-) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py b/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py index eaf7570d427..2e3cdd12997 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py @@ -33,7 +33,7 @@ from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.responses.main import OutputFunctionToolCall -from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse +from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message, ModelResponse class _DummyMCPResult: @@ -1574,3 +1574,382 @@ async def test_apply_toolset_permissions_pins_the_auth_to_explicit_grants_only(m assert scoped.object_permission is not None assert scoped.object_permission.mcp_servers == ["srv-1"] assert scoped.object_permission.mcp_tool_permissions == {"srv-1": ["add"]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "lookup,expected", + [ + (AsyncMock(return_value=types.SimpleNamespace(toolset_id="ts-1")), True), + (AsyncMock(return_value=None), False), + (AsyncMock(side_effect=RuntimeError("db down")), False), + ], +) +async def test_toolset_exists_reflects_the_toolset_lookup(monkeypatch, lookup: AsyncMock, expected: bool): + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + types.SimpleNamespace(get_toolset_by_name_cached=lookup), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + + assert await mcp_handler_module._toolset_exists("team-toolset") is expected + + +@pytest.mark.asyncio +async def test_apply_toolset_permissions_merges_into_existing_object_permission(monkeypatch): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + types.SimpleNamespace( + resolve_toolset_tool_permissions=AsyncMock(return_value={"srv-1": ["add"], "srv-2": ["sub"]}) + ), + ) + existing: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", + mcp_servers=["old"], + mcp_tool_permissions={"srv-1": ["mul"]}, + mcp_toolsets=["ts-1"], + ) + + scoped: Final = await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions( + resolved_toolset_ids=["ts-1"], + resolved_mcp_servers=["direct"], + user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", object_permission=existing), + ) + + op: Final = scoped.object_permission + assert op is not None + assert op.object_permission_id == "op-key" + assert sorted(op.mcp_servers or []) == ["direct", "srv-1", "srv-2"] + assert {server: sorted(tools) for server, tools in (op.mcp_tool_permissions or {}).items()} == { + "srv-1": ["add", "mul"], + "srv-2": ["sub"], + } + assert op.mcp_toolsets == [] + assert scoped.mcp_explicit_grants_only is True + + +@pytest.mark.asyncio +async def test_apply_toolset_permissions_returns_auth_unchanged_when_resolution_fails(monkeypatch): + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + types.SimpleNamespace(resolve_toolset_tool_permissions=AsyncMock(side_effect=RuntimeError("db down"))), + ) + auth: Final = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + + scoped: Final = await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions( + resolved_toolset_ids=["ts-1"], + resolved_mcp_servers=[], + user_api_key_auth=auth, + ) + + assert scoped is auth + assert scoped.object_permission is None + assert scoped.mcp_explicit_grants_only is not True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prisma_client", [None, MagicMock()]) +async def test_get_mcp_tools_from_manager_treats_unresolvable_toolsets_as_server_names(monkeypatch, prisma_client): + mock_get_tools = AsyncMock(return_value=AggregateToolListing(tools=[], outcomes={})) + monkeypatch.setattr("litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", mock_get_tools) + manager = _toolset_gateway_manager("unused", "unused") + manager.get_toolset_by_name_cached = AsyncMock(side_effect=[None, RuntimeError("db down")]) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + manager, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + auth: Final = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + + tools, server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( + user_api_key_auth=auth, + mcp_tools_with_litellm_proxy=[ + {"type": "mcp", "server_url": "litellm_proxy/mcp/missing"}, + {"type": "mcp", "server_url": "litellm_proxy/mcp/broken"}, + ], + ) + + assert tools == [] + assert server_names == [] + assert mock_get_tools.await_args is not None + assert mock_get_tools.await_args.kwargs["mcp_servers"] == ["missing", "broken"] + assert mock_get_tools.await_args.kwargs["user_api_key_auth"] is auth + + +@pytest.mark.asyncio +async def test_process_mcp_tools_to_openai_format_filters_dedupes_and_maps_to_the_single_server(monkeypatch): + server: Final = _registered("srv-1", "deepwiki", server_name="deepwiki") + listed: Final = [ + MCPTool(name="deepwiki-read", inputSchema={"type": "object"}), + MCPTool(name="deepwiki-read", inputSchema={"type": "object"}), + MCPTool(name="deepwiki-write", inputSchema={"type": "object"}), + ] + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", + AsyncMock(return_value=AggregateToolListing(tools=listed, outcomes={})), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_mcp_server_by_name=MagicMock(return_value=server), + get_allowed_mcp_servers=AsyncMock(return_value=["srv-1"]), + get_mcp_servers_from_ids=MagicMock(return_value=[server]), + ), + ) + + openai_tools, tool_server_map = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_to_openai_format( + user_api_key_auth=UserAPIKeyAuth(api_key="sk-test"), + mcp_tools_with_litellm_proxy=[ + {"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki", "allowed_tools": ["deepwiki-read"]} + ], + ) + + assert [(tool["type"], tool["name"]) for tool in openai_tools] == [("function", "deepwiki-read")] + assert tool_server_map == {"deepwiki-read": "deepwiki"} + + +@pytest.mark.asyncio +async def test_process_mcp_tools_to_openai_format_returns_nothing_without_proxy_tools(): + assert await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_to_openai_format(None, []) == ([], {}) + + +def test_extract_tool_calls_from_response_keeps_dict_and_model_function_calls_only(): + model_call: Final = OutputFunctionToolCall( + id="id-2", type="function_call", call_id="call-2", name="bar", arguments="{}", status="completed" + ) + dict_call: Final = {"type": "function_call", "call_id": "call-1", "name": "foo"} + response: Final = types.SimpleNamespace( + output=[dict_call, {"type": "message", "content": []}, model_call, types.SimpleNamespace(type="reasoning")] + ) + + assert LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(cast(Any, response)) == [dict_call, model_call] + + +def test_extract_tool_calls_from_chat_response_skips_choices_without_message_and_keeps_raw_entries(): + raw_call: Final = {"id": "call-raw", "function": {"name": "foo", "arguments": "{}"}} + response: Final = types.SimpleNamespace( + choices=[ + types.SimpleNamespace(message=None), + types.SimpleNamespace(message=types.SimpleNamespace(tool_calls=[raw_call])), + ] + ) + + assert LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(cast(Any, response)) == [raw_call] + + +def test_extract_tool_calls_from_chat_response_returns_empty_on_malformed_response(): + response: Final = types.SimpleNamespace(choices=None) + + assert LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(cast(Any, response)) == [] + + +@pytest.mark.parametrize( + "raw,expected", + [ + ('{"q": 1}', {"q": 1}), + ("not json", {}), + ({"q": 2}, {"q": 2}), + (None, {}), + ], +) +def test_parse_tool_arguments(raw: Any, expected: dict[str, object]): + assert LiteLLM_Proxy_MCP_Handler._parse_tool_arguments(raw) == expected + + +def test_parse_mcp_result_summarizes_non_text_content(): + from mcp.types import EmbeddedResource, ImageContent, TextResourceContents + + result: Final = CallToolResult( + content=[ + TextContent(type="text", text="hello"), + ImageContent(type="image", data="aGk=", mimeType="image/png"), + EmbeddedResource( + type="resource", + resource=TextResourceContents(uri="file:///a.txt", text="a", mimeType="text/plain"), + ), + ], + isError=False, + ) + + assert LiteLLM_Proxy_MCP_Handler._parse_mcp_result(result) == "hello [Generated Image, EmbeddedResource]" + + +def test_parse_mcp_result_falls_back_when_empty(): + assert ( + LiteLLM_Proxy_MCP_Handler._parse_mcp_result(CallToolResult(content=[], isError=False)) + == "Tool executed successfully" + ) + + +@pytest.mark.asyncio +async def test_execute_tool_calls_skips_tool_calls_without_a_name(monkeypatch): + call_tool_mock: Final = _setup_mcp_call_environment(monkeypatch) + + results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={}, + tool_calls=[{"id": "call-1", "function": {"arguments": "{}"}}], + user_api_key_auth=None, + ) + + assert results == [] + call_tool_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_execute_tool_calls_survives_pre_call_failure_and_logs_server_branding(monkeypatch): + result: Final = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_mcp_call_hook = AsyncMock(side_effect=lambda response, **_: response) + monkeypatch.setitem( + sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj) + ) + server: Final = _registered("srv-1", "deepwiki", server_name="deepwiki").model_copy( + update={"mcp_info": {"logo_url": "https://example.com/logo.png", "mcp_server_cost_info": {"default": 0.5}}} + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=result), + get_mcp_server_by_name=MagicMock(return_value=server), + get_mcp_server_from_tool_name=MagicMock(return_value=None), + ), + ) + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.pre_call.side_effect = RuntimeError("pre_call broke") + logging_obj.async_post_mcp_tool_call_hook = AsyncMock(side_effect=lambda response_obj, **_: response_obj) + logging_obj.async_success_handler = AsyncMock() + monkeypatch.setattr(mcp_handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None)) + + tool_name: Final = "deepwiki-read_wiki_structure" + results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}], + user_api_key_auth=None, + ) + + assert results == [{"tool_call_id": "call-1", "result": "ok", "name": tool_name}] + metadata: Final = logging_obj.model_call_details["mcp_tool_call_metadata"] + assert metadata["mcp_server_name"] == "deepwiki" + assert metadata["mcp_server_logo_url"] == "https://example.com/logo.png" + assert metadata["mcp_server_cost_info"] == {"default": 0.5} + + +@pytest.mark.asyncio +async def test_execute_tool_calls_reports_blocked_pii_as_tool_result(monkeypatch): + from litellm.exceptions import BlockedPiiEntityError + + call_tool_mock: Final = _setup_mcp_call_environment(monkeypatch) + call_tool_mock.side_effect = BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="presidio") + post_call_failure_hook: Final = _setup_proxy_logging(monkeypatch) + + tool_name: Final = "deepwiki-read_wiki_structure" + results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=[{"id": "call-pii", "function": {"name": tool_name, "arguments": "{}"}}], + user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", user_id="u1"), + ) + + assert len(results) == 1 + assert results[0]["tool_call_id"] == "call-pii" + assert results[0]["name"] == tool_name + assert results[0]["result"].startswith( + "Tool call blocked: PII entity 'EMAIL_ADDRESS' detected by guardrail 'presidio'." + ) + assert post_call_failure_hook.await_args is not None + assert isinstance(post_call_failure_hook.await_args.kwargs["original_exception"], BlockedPiiEntityError) + + +@pytest.mark.asyncio +async def test_execute_tool_calls_reports_http_errors_even_when_failure_logging_breaks(monkeypatch): + call_tool_mock: Final = _setup_mcp_call_environment(monkeypatch) + call_tool_mock.side_effect = HTTPException(status_code=403, detail="forbidden tool") + post_call_failure_hook: Final = _setup_proxy_logging(monkeypatch) + post_call_failure_hook.side_effect = RuntimeError("logging down") + + tool_name: Final = "deepwiki-read_wiki_structure" + results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=[{"id": "call-http", "function": {"name": tool_name, "arguments": "{}"}}], + user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", user_id="u1"), + ) + + assert results == [{"tool_call_id": "call-http", "result": "Tool call failed: forbidden tool", "name": tool_name}] + post_call_failure_hook.assert_awaited_once() + + +def test_create_follow_up_messages_for_chat_starts_from_assistant_turn_without_prior_messages(): + response: Final = types.SimpleNamespace( + choices=[Choices(index=0, message=Message(role="assistant", content="calling foo"))] + ) + + follow_up: Final = LiteLLM_Proxy_MCP_Handler.create_follow_up_messages_for_chat( + [], cast(Any, response), [{"tool_call_id": "call-1", "name": "foo", "result": "done"}] + ) + + assert len(follow_up) == 2 + assert follow_up[0]["role"] == "assistant" + assert follow_up[0]["content"] == "calling foo" + assert "tool_calls" not in follow_up[0] + assert follow_up[1] == {"role": "tool", "tool_call_id": "call-1", "name": "foo", "content": "done"} + + +def test_create_follow_up_messages_for_chat_keeps_tool_results_when_response_has_no_choices(): + response: Final = types.SimpleNamespace(choices=[]) + + follow_up: Final = LiteLLM_Proxy_MCP_Handler.create_follow_up_messages_for_chat( + [{"role": "user", "content": "hi"}], + cast(Any, response), + [{"tool_call_id": "call-1", "name": "foo", "result": "done"}], + ) + + assert [message["role"] for message in follow_up] == ["user", "tool"] + + +def test_create_follow_up_input_extends_list_input_and_collects_assistant_message_content(): + original_input: Final = [{"type": "message", "role": "user", "content": "hi"}] + response: Final = types.SimpleNamespace( + output=[ + {"type": "message", "content": [{"type": "output_text", "text": "a"}]}, + {"type": "message", "content": "b"}, + {"type": "function_call", "call_id": "call-1", "arguments": "{}"}, + {"type": "function_call", "id": "call-2", "name": "foo", "arguments": "{}"}, + {"type": "reasoning", "summary": []}, + object(), + ] + ) + + follow_up: Final = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( + response=cast(Any, response), + tool_results=[{"tool_call_id": "call-2", "result": "done"}], + original_input=cast(Any, original_input), + ) + + assert follow_up == [ + {"type": "message", "role": "user", "content": "hi"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "a"}, "b"]}, + {"type": "function_call", "call_id": "call-2", "name": "foo", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call-2", "output": "done"}, + ] + + +@pytest.mark.parametrize( + "original_input,expected_first", + [ + ("hi", {"type": "message", "role": "user", "content": "hi"}), + ({"type": "message", "role": "user", "content": "raw"}, {"type": "message", "role": "user", "content": "raw"}), + ], +) +def test_create_follow_up_input_prepends_non_list_original_input(original_input: Any, expected_first: object): + follow_up: Final = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( + response=cast(Any, types.SimpleNamespace(output=[])), + tool_results=[], + original_input=original_input, + ) + + assert follow_up == [expected_first] From 27f11f0544b410178f245f1b8be58489088dc6a8 Mon Sep 17 00:00:00 2001 From: nate-berri Date: Sat, 10 Oct 2026 19:08:50 +0000 Subject: [PATCH 5/9] test(mcp): keep handler tests in-process and off doubles of our own code Part of LIT-9423 --- .../test_litellm_proxy_mcp_handler.py | 115 +++++------------- .../mcp/test_litellm_proxy_mcp_handler.py | 19 +-- 2 files changed, 42 insertions(+), 92 deletions(-) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py b/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py index 2e3cdd12997..341d6f014d4 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py @@ -22,6 +22,7 @@ from litellm.proxy._experimental.mcp_server import operations as mcp_operations from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing from litellm.proxy._experimental.mcp_server.mcp_server_manager import ListedToolsCaller, MCPServerManager from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import ProxyLogging from litellm.responses import main as responses_main from litellm.proxy._experimental.mcp_server import litellm_proxy_mcp_handler as mcp_handler_module @@ -32,6 +33,7 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.types.mcp_server.mcp_toolset import MCPToolset from litellm.types.responses.main import OutputFunctionToolCall from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message, ModelResponse @@ -1576,21 +1578,25 @@ async def test_apply_toolset_permissions_pins_the_auth_to_explicit_grants_only(m assert scoped.object_permission.mcp_tool_permissions == {"srv-1": ["add"]} +def _toolset_prisma_client(find_first: AsyncMock | None = None, find_many: AsyncMock | None = None) -> MagicMock: + client: Final = MagicMock() + client.db.litellm_mcptoolsettable.find_first = find_first or AsyncMock(return_value=None) + client.db.litellm_mcptoolsettable.find_many = find_many or AsyncMock(return_value=[]) + return client + + @pytest.mark.asyncio @pytest.mark.parametrize( - "lookup,expected", + "find_first,expected", [ - (AsyncMock(return_value=types.SimpleNamespace(toolset_id="ts-1")), True), + (AsyncMock(return_value=MCPToolset(toolset_id="ts-1", toolset_name="team-toolset")), True), (AsyncMock(return_value=None), False), (AsyncMock(side_effect=RuntimeError("db down")), False), ], ) -async def test_toolset_exists_reflects_the_toolset_lookup(monkeypatch, lookup: AsyncMock, expected: bool): - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - types.SimpleNamespace(get_toolset_by_name_cached=lookup), - ) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) +async def test_toolset_exists_reflects_the_toolset_lookup(monkeypatch, find_first: AsyncMock, expected: bool): + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _toolset_prisma_client(find_first=find_first)) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) assert await mcp_handler_module._toolset_exists("team-toolset") is expected @@ -1599,12 +1605,16 @@ async def test_toolset_exists_reflects_the_toolset_lookup(monkeypatch, lookup: A async def test_apply_toolset_permissions_merges_into_existing_object_permission(monkeypatch): from litellm.proxy._types import LiteLLM_ObjectPermissionTable - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - types.SimpleNamespace( - resolve_toolset_tool_permissions=AsyncMock(return_value={"srv-1": ["add"], "srv-2": ["sub"]}) - ), + toolset: Final = MCPToolset( + toolset_id="ts-1", + toolset_name="team-toolset", + tools=[{"server_id": "srv-1", "tool_name": "add"}, {"server_id": "srv-2", "tool_name": "sub"}], ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + _toolset_prisma_client(find_many=AsyncMock(return_value=[toolset])), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) existing: Final = LiteLLM_ObjectPermissionTable( object_permission_id="op-key", mcp_servers=["old"], @@ -1631,84 +1641,23 @@ async def test_apply_toolset_permissions_merges_into_existing_object_permission( @pytest.mark.asyncio -async def test_apply_toolset_permissions_returns_auth_unchanged_when_resolution_fails(monkeypatch): +async def test_apply_toolset_permissions_grants_no_toolset_servers_when_the_toolset_table_is_down(monkeypatch): monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - types.SimpleNamespace(resolve_toolset_tool_permissions=AsyncMock(side_effect=RuntimeError("db down"))), + "litellm.proxy.proxy_server.prisma_client", + _toolset_prisma_client(find_many=AsyncMock(side_effect=RuntimeError("db down"))), ) - auth: Final = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) scoped: Final = await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions( resolved_toolset_ids=["ts-1"], resolved_mcp_servers=[], - user_api_key_auth=auth, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", user_id="u1"), ) - assert scoped is auth - assert scoped.object_permission is None - assert scoped.mcp_explicit_grants_only is not True - - -@pytest.mark.asyncio -@pytest.mark.parametrize("prisma_client", [None, MagicMock()]) -async def test_get_mcp_tools_from_manager_treats_unresolvable_toolsets_as_server_names(monkeypatch, prisma_client): - mock_get_tools = AsyncMock(return_value=AggregateToolListing(tools=[], outcomes={})) - monkeypatch.setattr("litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", mock_get_tools) - manager = _toolset_gateway_manager("unused", "unused") - manager.get_toolset_by_name_cached = AsyncMock(side_effect=[None, RuntimeError("db down")]) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - manager, - ) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) - auth: Final = UserAPIKeyAuth(api_key="sk-test", user_id="u1") - - tools, server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( - user_api_key_auth=auth, - mcp_tools_with_litellm_proxy=[ - {"type": "mcp", "server_url": "litellm_proxy/mcp/missing"}, - {"type": "mcp", "server_url": "litellm_proxy/mcp/broken"}, - ], - ) - - assert tools == [] - assert server_names == [] - assert mock_get_tools.await_args is not None - assert mock_get_tools.await_args.kwargs["mcp_servers"] == ["missing", "broken"] - assert mock_get_tools.await_args.kwargs["user_api_key_auth"] is auth - - -@pytest.mark.asyncio -async def test_process_mcp_tools_to_openai_format_filters_dedupes_and_maps_to_the_single_server(monkeypatch): - server: Final = _registered("srv-1", "deepwiki", server_name="deepwiki") - listed: Final = [ - MCPTool(name="deepwiki-read", inputSchema={"type": "object"}), - MCPTool(name="deepwiki-read", inputSchema={"type": "object"}), - MCPTool(name="deepwiki-write", inputSchema={"type": "object"}), - ] - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", - AsyncMock(return_value=AggregateToolListing(tools=listed, outcomes={})), - ) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_mcp_server_by_name=MagicMock(return_value=server), - get_allowed_mcp_servers=AsyncMock(return_value=["srv-1"]), - get_mcp_servers_from_ids=MagicMock(return_value=[server]), - ), - ) - - openai_tools, tool_server_map = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_to_openai_format( - user_api_key_auth=UserAPIKeyAuth(api_key="sk-test"), - mcp_tools_with_litellm_proxy=[ - {"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki", "allowed_tools": ["deepwiki-read"]} - ], - ) - - assert [(tool["type"], tool["name"]) for tool in openai_tools] == [("function", "deepwiki-read")] - assert tool_server_map == {"deepwiki-read": "deepwiki"} + assert scoped.object_permission is not None + assert scoped.object_permission.mcp_servers == [] + assert scoped.object_permission.mcp_tool_permissions == {} + assert scoped.mcp_explicit_grants_only is True @pytest.mark.asyncio 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 f4e2656c068..b80ff31333b 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -1,21 +1,22 @@ +import importlib +import sys from typing import Final import pytest from litellm.proxy._experimental.mcp_server import litellm_proxy_mcp_handler as relocated from litellm.responses.mcp import litellm_proxy_mcp_handler as legacy -from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter -def test_legacy_import_leaves_relocated_proxy_module_unloaded() -> None: - result: Final = run_child_interpreter( - "import sys, litellm.responses.mcp.litellm_proxy_mcp_handler\n" - "print('litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler' in sys.modules)", - timeout=120, - ) +def test_legacy_import_leaves_relocated_proxy_module_unloaded(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delitem(sys.modules, legacy.__name__) + monkeypatch.delitem(sys.modules, relocated.__name__) + monkeypatch.setattr(sys.modules["litellm.responses.mcp"], "litellm_proxy_mcp_handler", legacy) - assert result.returncode == 0, result.stderr - assert result.stdout.strip() == "False" + reimported: Final = importlib.import_module(legacy.__name__) + + assert reimported is not legacy + assert relocated.__name__ not in sys.modules @pytest.mark.parametrize("name", legacy.__all__) From 3456e88c2bdc69dbe3c01c488854b1e3c390deda Mon Sep 17 00:00:00 2001 From: nate-berri Date: Sat, 10 Oct 2026 19:20:44 +0000 Subject: [PATCH 6/9] fix(docs-test): import litellm.proxy._types explicitly in the API docs check import litellm no longer loads proxy modules after the layer cleanup, so litellm.proxy was unbound and the docs job failed on main --- tests/documentation_tests/test_api_docs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/documentation_tests/test_api_docs.py b/tests/documentation_tests/test_api_docs.py index d8536f13b9c..6c93f100c5d 100644 --- a/tests/documentation_tests/test_api_docs.py +++ b/tests/documentation_tests/test_api_docs.py @@ -5,7 +5,7 @@ from dataclasses import dataclass import argparse import re -import litellm +import litellm.proxy._types @dataclass From ca478b73824867703cb55373ca9cecd8cffd3a6a Mon Sep 17 00:00:00 2001 From: nate-berri Date: Sat, 10 Oct 2026 19:23:35 +0000 Subject: [PATCH 7/9] fix(sdk): resolve litellm.proxy lazily on attribute access The layer cleanup stopped import litellm from loading any proxy module, so litellm.proxy was no longer bound and litellm.proxy._types.X raised AttributeError. The module __getattr__ now imports litellm.proxy on first access, which keeps import litellm free of proxy._types and fixes the docs and code-quality checks on main --- litellm/__init__.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litellm/__init__.py b/litellm/__init__.py index 4196912488f..a6bbd132cbf 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -2386,6 +2386,11 @@ def __getattr__(name: str) -> Any: handler_func: Final = registry[name] return handler_func(name) + if name == "proxy": + import importlib + + return importlib.import_module("litellm.proxy") + if name == "harness" or name in _AGENT_EXPORTS: import importlib From d2e5868d903ff1e463b384aa08d07eb9f9de1523 Mon Sep 17 00:00:00 2001 From: nate-berri Date: Sat, 10 Oct 2026 19:23:59 +0000 Subject: [PATCH 8/9] Revert "fix(docs-test): import litellm.proxy._types explicitly in the API docs check" This reverts commit 3456e88c2bdc69dbe3c01c488854b1e3c390deda. --- tests/documentation_tests/test_api_docs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/documentation_tests/test_api_docs.py b/tests/documentation_tests/test_api_docs.py index 6c93f100c5d..d8536f13b9c 100644 --- a/tests/documentation_tests/test_api_docs.py +++ b/tests/documentation_tests/test_api_docs.py @@ -5,7 +5,7 @@ from dataclasses import dataclass import argparse import re -import litellm.proxy._types +import litellm @dataclass From 1237353e401ebb7f1317146e128377a5f0e462d9 Mon Sep 17 00:00:00 2001 From: nate-berri Date: Sat, 10 Oct 2026 21:04:24 +0000 Subject: [PATCH 9/9] fix(types): drop unused TypeIs import from proxy/_types.py --- litellm/proxy/_types.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7b2a63974d8..b6c81f17761 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -25,7 +25,6 @@ from typing_extensions import ( ReadOnly, Required, TypedDict, - TypeIs, # noqa: TID251 # narrows the before-validator input without copying it ) from litellm._uuid import uuid