diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 1dc9db0a0c7..f53c11ab948 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -573,10 +573,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 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..4dd0394005d --- /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.utils import ProxyLogging + from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +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 4dd0394005d..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.utils import ProxyLogging - from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth -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/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py b/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py new file mode 100644 index 00000000000..341d6f014d4 --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/test_litellm_proxy_mcp_handler.py @@ -0,0 +1,1904 @@ +import asyncio +import importlib +import subprocess +import sys +import textwrap +import types +from contextlib import nullcontext +from typing import Any, Final, Literal, cast +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException +from mcp.types import CallToolResult, TextContent +from mcp.types import Tool as MCPTool +from openai.types.responses.tool_param import Mcp + +import litellm +from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging +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 +from litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, +) +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 + + +class _DummyMCPResult: + def __init__(self): + self.content = [] + + +def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + """Patch MCP globals so _execute_tool_calls can run in tests.""" + proxy_module = types.SimpleNamespace(proxy_logging_obj=object(), prisma_client=None) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) + + fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=_DummyMCPResult()), + # Newer logging path calls this to enrich spend logs metadata + get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + return fake_manager.call_tool + + +def _setup_proxy_logging(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + """Patch proxy_logging_obj so failure hook can be asserted.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock() + proxy_module = types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj, prisma_client=None) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) + return proxy_logging_obj.post_call_failure_hook + + +def test_deduplicate_mcp_tools_single_allowed_server(): + tools = [{"name": "search"}, {"name": "search"}] # duplicate on purpose + + deduped, server_map = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools( + tools, + ["everything"], + ) + + assert len(deduped) == 1 + assert server_map == {"search": "everything"} + + +@pytest.mark.parametrize( + "tool_name,expected_server", + [ + ("alpha-tool", "alpha"), + ("beta-another_tool", "beta"), + ], +) +def test_deduplicate_mcp_tools_prefixed_names(tool_name, expected_server): + tools = [{"name": tool_name}] + + _, server_map = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools( + tools, + ["alpha", "beta"], + ) + + assert server_map[tool_name] == expected_server + + +def test_extract_tool_calls_from_chat_response_handles_tool_calls(): + response = ModelResponse( + id="resp-1", + choices=[ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call-123", + "type": "function", + "function": {"name": "foo", "arguments": "{}"}, + } + ], + }, + "finish_reason": "tool_calls", + } + ], + model="gpt", + created=0, + object="chat.completion", + ) + + tool_calls = LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response) + + assert len(tool_calls) == 1 + assert tool_calls[0]["function"]["name"] == "foo" + + +def test_create_follow_up_messages_for_chat_appends_tool_results(): + original_messages = [{"role": "user", "content": "hi"}] + response = ModelResponse( + id="resp-2", + choices=[ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call-abc", + "type": "function", + "function": {"name": "foo", "arguments": "{}"}, + } + ], + }, + } + ], + model="gpt", + created=0, + object="chat.completion", + ) + tool_results = [ + { + "tool_call_id": "call-abc", + "name": "foo", + "result": "done", + } + ] + + follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_messages_for_chat( + original_messages, + response, + tool_results, + ) + + assert follow_up[0]["role"] == "user" + assert follow_up[-1]["role"] == "tool" + assert follow_up[-1]["name"] == "foo" + assert follow_up[-1]["content"] == "done" + + +def test_transform_mcp_tools_to_openai_uses_chat_format(monkeypatch): + captured = {} + + def fake_transform_chat(tool): + captured.setdefault("chat", []).append(tool) + return {"chat": True} + + def fake_transform_responses(tool): + captured.setdefault("responses", []).append(tool) + return {"responses": True} + + monkeypatch.setattr( + "litellm.experimental_mcp_client.tools.transform_mcp_tool_to_openai_tool", + fake_transform_chat, + ) + monkeypatch.setattr( + "litellm.experimental_mcp_client.tools.transform_mcp_tool_to_openai_responses_api_tool", + fake_transform_responses, + ) + + chat_tools = LiteLLM_Proxy_MCP_Handler.transform_mcp_tools_to_openai(["tool"], target_format="chat") + resp_tools = LiteLLM_Proxy_MCP_Handler.transform_mcp_tools_to_openai(["tool"]) + + assert chat_tools == [{"chat": True}] + assert resp_tools == [{"responses": True}] + assert captured["chat"] == ["tool"] + assert captured["responses"] == ["tool"] + + +def test_create_follow_up_input_handles_response_function_tool_call(): + response = types.SimpleNamespace( + output=[ + OutputFunctionToolCall( + id="id", + type="function_call", + call_id="call-1", + name="foo", + arguments="{}", + status="completed", + ) + ] + ) + + follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( + response=cast(Any, response), + tool_results=[], + original_input=None, + ) + + assert follow_up == [ + { + "type": "function_call", + "call_id": "call-1", + "name": "foo", + "arguments": "{}", + } + ] + + +@pytest.mark.asyncio +async def test_execute_tool_calls_strips_server_prefix(monkeypatch): + call_tool_mock = _setup_mcp_call_environment(monkeypatch) + tool_name = "deepwiki-read_wiki_structure" + tool_calls = [ + { + "id": "call-1", + "function": {"name": tool_name, "arguments": "{}"}, + } + ] + + await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=tool_calls, + user_api_key_auth=None, + ) + + assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_args is not None + assert call_tool_mock.await_args.kwargs["name"] == "read_wiki_structure" + + +@pytest.mark.asyncio +async def test_execute_tool_calls_keeps_tool_name_without_prefix(monkeypatch): + call_tool_mock = _setup_mcp_call_environment(monkeypatch) + tool_name = "read_wiki_structure" + tool_calls = [ + { + "id": "call-2", + "function": {"name": tool_name, "arguments": "{}"}, + } + ] + + await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=tool_calls, + user_api_key_auth=None, + ) + + assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_args is not None + assert call_tool_mock.await_args.kwargs["name"] == tool_name + + +@pytest.mark.asyncio +async def test_execute_tool_calls_keeps_tool_name_when_equal_to_server(monkeypatch): + call_tool_mock = _setup_mcp_call_environment(monkeypatch) + tool_name = "echo" + tool_calls = [ + { + "id": "call-3", + "function": {"name": tool_name, "arguments": "{}"}, + } + ] + + await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={tool_name: "echo"}, + tool_calls=tool_calls, + user_api_key_auth=None, + ) + + assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_args is not None + assert call_tool_mock.await_args.kwargs["name"] == tool_name + + +@pytest.mark.asyncio +async def test_execute_tool_calls_strips_prefix_when_alias_differs_from_server_name( + monkeypatch, +): + call_tool_mock = _setup_mcp_call_environment(monkeypatch) + fake_server = types.SimpleNamespace( + alias="my_deepwiki", + server_name="deepwiki_test", + server_id="test-server-id", + short_prefix=None, + mcp_info=None, + tool_name_to_display_name=None, + ) + from litellm.proxy._experimental.mcp_server import mcp_server_manager as _msm + + _msm.global_mcp_server_manager.get_mcp_server_from_tool_name = MagicMock(return_value=fake_server) + + tool_name = "my_deepwiki-read_wiki_structure" + tool_calls = [ + { + "id": "call-4", + "function": {"name": tool_name, "arguments": "{}"}, + } + ] + + await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={tool_name: "deepwiki_test"}, + tool_calls=tool_calls, + user_api_key_auth=None, + ) + + assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_args is not None + assert call_tool_mock.await_args.kwargs["name"] == "read_wiki_structure" + + +@pytest.mark.asyncio +async def test_execute_tool_calls_reverse_maps_display_name(monkeypatch): + call_tool_mock = _setup_mcp_call_environment(monkeypatch) + colliding_server = types.SimpleNamespace( + alias=None, + server_name="other_mcp", + server_id="other-server-id", + short_prefix=None, + mcp_info=None, + tool_name_to_display_name={"search": "search_docs"}, + ) + fake_server = types.SimpleNamespace( + alias=None, + server_name="deepwiki_mcp", + server_id="test-server-id", + short_prefix=None, + mcp_info=None, + tool_name_to_display_name={"read_wiki_structure": "browse_repo_docs"}, + ) + from litellm.proxy._experimental.mcp_server import mcp_server_manager as _msm + + _msm.global_mcp_server_manager.get_mcp_server_from_tool_name = MagicMock(return_value=colliding_server) + _msm.global_mcp_server_manager.get_mcp_server_by_name = MagicMock(return_value=fake_server) + + tool_name = "browse_repo_docs" + tool_calls = [ + { + "id": "call-5", + "function": {"name": tool_name, "arguments": "{}"}, + } + ] + + await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={tool_name: "deepwiki_mcp"}, + tool_calls=tool_calls, + user_api_key_auth=None, + ) + + assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_args is not None + assert call_tool_mock.await_args.kwargs["name"] == "read_wiki_structure" + + +@pytest.mark.asyncio +async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkeypatch): + """ + Regression test for ae4d92ad...: + Ensure responses-side MCP tool execution logs failures via proxy_logging_obj.post_call_failure_hook. + """ + post_call_failure_hook = _setup_proxy_logging(monkeypatch) + + fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(side_effect=HTTPException(status_code=500, detail="boom")), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + tool_name = "deepwiki-read_wiki_structure" + tool_calls = [{"id": "call-err", "function": {"name": tool_name, "arguments": "{}"}}] + + user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") + + results = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=tool_calls, + user_api_key_auth=user_auth, + litellm_call_id="cid", + litellm_trace_id="tid", + ) + + assert len(results) == 1 + assert results[0]["tool_call_id"] == "call-err" + assert results[0]["name"] == tool_name + + post_call_failure_hook.assert_awaited_once() + assert post_call_failure_hook.await_args is not None + assert post_call_failure_hook.await_args.kwargs.get("route") == "/responses/mcp/call_tool" + + +@pytest.mark.asyncio +async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_function_setup( + monkeypatch, +): + """ + Regression test for ae4d92ad...: + Ensure litellm_call_id / litellm_trace_id are forwarded into function_setup kwargs. + """ + _setup_proxy_logging(monkeypatch) + call_tool_mock = _setup_mcp_call_environment(monkeypatch) + + captured = {} + + def fake_function_setup(*_args, **kwargs): + captured.update(kwargs) + return None, None + + # 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.proxy._experimental.mcp_server.litellm_proxy_mcp_handler") + monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) + + tool_name = "deepwiki-read_wiki_structure" + tool_calls = [{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}] + + await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=tool_calls, + user_api_key_auth=None, + litellm_call_id="cid", + litellm_trace_id="tid", + ) + + # Ensure the tool call was attempted (sanity) + assert call_tool_mock.await_count == 1 + + assert captured.get("litellm_call_id") == "cid" + assert captured.get("litellm_trace_id") == "tid" + + +@pytest.mark.asyncio +async def test_execute_tool_calls_threads_logging_obj_into_call_tool(monkeypatch): + """The Responses-API MCP path must hand the request's litellm_logging_obj to + global_mcp_server_manager.call_tool, otherwise pre_call_tool_check / + _create_during_hook_task get None and no guardrail evaluation is bridged onto + the request logger, so MCP tool calls made through the Responses API report zero + guardrail evaluations in the monitor. Drop the litellm_logging_obj kwarg on the + call_tool invocation and this fails.""" + _setup_proxy_logging(monkeypatch) + call_tool_mock = _setup_mcp_call_environment(monkeypatch) + + sentinel_logging_obj = MagicMock() + sentinel_logging_obj.async_post_mcp_tool_call_hook = AsyncMock() + sentinel_logging_obj.async_success_handler = AsyncMock() + + handler_module = importlib.import_module("litellm.proxy._experimental.mcp_server.litellm_proxy_mcp_handler") + monkeypatch.setattr( + handler_module, + "function_setup", + lambda *_args, **_kwargs: (sentinel_logging_obj, None), + ) + + tool_name = "deepwiki-read_wiki_structure" + tool_calls = [{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}] + + await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=tool_calls, + user_api_key_auth=None, + ) + + assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_args is not None + assert call_tool_mock.await_args.kwargs["litellm_logging_obj"] is sentinel_logging_obj + + +@pytest.mark.asyncio +async def test_execute_tool_calls_applies_post_call_hook_content(monkeypatch): + proxy_module = types.SimpleNamespace(proxy_logging_obj=None) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) + + result = CallToolResult( + content=[TextContent(type="text", text="SECRET-1234")], + structuredContent={"result": "SECRET-1234"}, + isError=False, + ) + fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=result), + get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.async_post_mcp_tool_call_hook = AsyncMock( + return_value=CallToolResult(content=[TextContent(type="text", text="[REDACTED]")], is_error=True) + ) + logging_obj.async_success_handler = AsyncMock() + 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" + results = 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": "[REDACTED]", "name": tool_name}] + assert logging_obj.async_success_handler.await_args.kwargs["result"].content[0].text == "[REDACTED]" + assert logging_obj.async_success_handler.await_args.kwargs["result"].structured_content is None + + +@pytest.mark.asyncio +async def test_execute_tool_calls_returns_proxy_result_without_logging(monkeypatch): + result = 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) + ) + + fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=result), + get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + 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" + results = 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}] + proxy_logging_obj.post_mcp_call_hook.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_execute_tool_calls_passes_logging_details_to_proxy_hook(monkeypatch): + result = 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) + ) + + fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=result), + get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + logging_obj = MagicMock() + 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.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" + results = 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}] + assert proxy_logging_obj.post_mcp_call_hook.await_args.kwargs["request_data"] == logging_obj.model_call_details + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_stage", ["post_call_hook", "success_handler"]) +async def test_execute_tool_calls_continues_when_post_call_logging_fails(monkeypatch, failure_stage: str): + proxy_module = types.SimpleNamespace(proxy_logging_obj=None) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) + + result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=AsyncMock(return_value=result), + get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.post_call = MagicMock() + logging_obj.async_post_mcp_tool_call_hook = AsyncMock( + side_effect=RuntimeError("hook failed") if failure_stage == "post_call_hook" else None, + return_value=result, + ) + 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.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" + results = 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}] + + +@pytest.mark.asyncio +async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch): + """ + Regression test for 872e5b98...: + Ensure responses-side tool discovery enables list-tools SpendLogs logging flags. + """ + served_tools: Final = [ + MCPTool(name="safe", description="Safe lookup", inputSchema={"type": "object"}), + MCPTool( + name="masked", + description="Contact [MASKED]", + inputSchema={"type": "object", "properties": {"query": {"type": "string", "description": "For [MASKED]"}}}, + ), + ] + mock_get_tools = AsyncMock(return_value=AggregateToolListing(tools=served_tools, outcomes={})) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", + mock_get_tools, + ) + + # Patch manager methods used by _get_mcp_tools_from_manager to avoid needing full UserAPIKeyAuth fields. + fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + get_allowed_mcp_servers=AsyncMock(return_value=[]), + get_mcp_servers_from_ids=MagicMock(return_value=[]), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") + tools, _server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( + user_api_key_auth=user_auth, + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}], + ) + + forwarded: Final = LiteLLM_Proxy_MCP_Handler.transform_mcp_tools_to_openai(tools) + assert [tool["name"] for tool in forwarded] == ["safe", "masked"] + assert forwarded[0]["description"] == "Safe lookup" + assert forwarded[1]["description"] == "Contact [MASKED]" + assert forwarded[1]["parameters"] == { + "type": "object", + "properties": {"query": {"type": "string", "description": "For [MASKED]"}}, + "additionalProperties": False, + } + assert mock_get_tools.await_count == 1 + assert mock_get_tools.await_args is not None + assert mock_get_tools.await_args.kwargs["log_list_tools_to_spendlogs"] is True + assert mock_get_tools.await_args.kwargs["list_tools_log_source"] == "responses" + + +def test_get_parent_request_tags_from_metadata(): + tags = LiteLLM_Proxy_MCP_Handler.get_parent_request_tags({"metadata": {"tags": ["team-a", "prod"]}}) + assert tags == ["team-a", "prod"] + + +def test_get_parent_request_tags_from_nested_litellm_params(): + tags = LiteLLM_Proxy_MCP_Handler.get_parent_request_tags( + { + "metadata": {"tags": ["top-level"]}, + "litellm_params": { + "metadata": {"tags": ["nested"]}, + "proxy_server_request": {"headers": {"user-agent": "client/1.0"}}, + }, + } + ) + assert tags == ["nested", "User-Agent: client", "User-Agent: client/1.0"] + + +@pytest.mark.asyncio +async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch): + 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, + ) + fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + get_allowed_mcp_servers=AsyncMock(return_value=[]), + get_mcp_servers_from_ids=MagicMock(return_value=[]), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( + user_api_key_auth=types.SimpleNamespace(api_key="k", user_id="u"), + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}], + request_tags=["team-a"], + ) + + assert mock_get_tools.await_args.kwargs["request_tags"] == ["team-a"] + + +@pytest.mark.asyncio +async def test_execute_tool_calls_exposes_sanitized_client_headers_to_logging(monkeypatch): + """The Responses API MCP bridge used to log an empty header dict, hiding the caller's + headers from logging callbacks and hooks.""" + _setup_proxy_logging(monkeypatch) + _setup_mcp_call_environment(monkeypatch) + + captured = {} + + def fake_function_setup(*_args, **kwargs): + captured.update(kwargs) + return None, None + + 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" + 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, + raw_headers={"x-nuid": "nuid-1", "x-litellm-api-key": "sk-proxy", "cookie": "s=1"}, + ) + + expected = {"x-nuid": "nuid-1", "cookie": "***REDACTED***"} + assert captured["metadata"]["headers"] == expected + assert captured["proxy_server_request"]["headers"] == expected + + +@pytest.mark.asyncio +async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monkeypatch): + _setup_proxy_logging(monkeypatch) + _setup_mcp_call_environment(monkeypatch) + captured = {} + + def fake_function_setup(*_args, **kwargs): + captured.update(kwargs) + return None, None + + 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" + 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, + request_tags=["team-a", "prod"], + ) + + assert captured["metadata"]["tags"] == ["team-a", "prod"] + + +def test_completion_with_function_tools_works_without_fastapi_installed(): + script = textwrap.dedent( + """ + import sys + + class _FastapiBlocker: + def find_spec(self, fullname, path=None, target=None): + if fullname == "fastapi" or fullname.startswith("fastapi."): + raise ModuleNotFoundError("No module named 'fastapi'") + return None + + sys.meta_path.insert(0, _FastapiBlocker()) + + import litellm + + response = litellm.completion( + model="openai/gpt-5.5", + messages=[{"role": "user", "content": "What is the weather in SF?"}], + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ], + mock_response="sunny", + ) + assert response.choices[0].message.content == "sunny" + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + + +def test_extract_tool_call_details_reads_anthropic_tool_use_input(): + """ + Regression test (LIT-4517): an Anthropic tool_use block carries its arguments + under `input`, not `arguments`. + + Given: A tool_use content block as /v1/messages returns it + When: The shared extractor reads it + Then: The arguments come back, so the MCP tool is called with them + + Reading only `arguments` fails silently rather than loudly: _parse_tool_arguments + turns the resulting None into {}, so the tool still executes, just with every + argument dropped. + """ + tool_use_block = { + "type": "tool_use", + "id": "toolu_01ABC", + "name": "read_wiki_structure", + "input": {"repoName": "BerriAI/litellm"}, + } + + name, arguments, call_id = LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_use_block) + + assert name == "read_wiki_structure" + assert call_id == "toolu_01ABC" + assert arguments == {"repoName": "BerriAI/litellm"} + assert LiteLLM_Proxy_MCP_Handler._parse_tool_arguments(arguments) == {"repoName": "BerriAI/litellm"} + + +def test_extract_tool_call_details_still_prefers_openai_arguments(): + """The OpenAI chat shape must keep winning; `input` is only the fallback.""" + openai_tool_call = { + "id": "call_123", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + + name, arguments, call_id = LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(openai_tool_call) + + assert name == "get_weather" + assert call_id == "call_123" + assert arguments == '{"city": "Paris"}' + + +def _registered( + server_id: str, + name: str, + alias: str | None = None, + server_name: str | None = None, + access_groups: list[str] | None = None, +): + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + return MCPServer( + server_id=server_id, + name=name, + alias=alias, + server_name=server_name, + transport=MCPTransport.http, + access_groups=access_groups, + ) + + +async def _no_toolset(_: str) -> bool: + return False + + +ZAPIER_TOOL: Mcp = { + "type": "mcp", + "server_label": "zapier", + "server_url": "https://mcp.zapier.com/api/mcp/mcp", + "require_approval": "never", +} +EXPLICIT_GATEWAY_TOOL = {"type": "mcp", "server_label": "github", "server_url": "litellm_proxy/mcp/github"} +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.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"]), + _registered("id-2", "deepwiki"), + ) + + async def toolset_exists(name: str) -> bool: + return name == "my-toolset" + + served = await _gateway_served_names( + {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset", "mcp", "nope"}, + servers=lambda: servers, + toolset_exists=toolset_exists, + ) + + assert served == {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset"} + + +@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.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) + servers = (_registered(server_id, "github-name", alias="github", access_groups=["prod-group"]),) + + served = await _gateway_served_names( + {server_id, short_prefix, "GitHub", "PROD-GROUP", "nope"}, servers=lambda: servers, toolset_exists=_no_toolset + ) + + assert served == {server_id, short_prefix, "GitHub"} + + +@pytest.mark.asyncio +async def test_routes_through_gateway_flags_explicit_and_served_tools_only(): + served_tool = {"type": "mcp", "server_label": "github", "server_url": "http://localhost:4000/mcp/github"} + + async def served_names(names): + assert names == {"github", "mcp"} + return frozenset({"github"}) + + flags = await LiteLLM_Proxy_MCP_Handler.routes_through_gateway( + [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, served_tool, FUNCTION_TOOL], served_names=served_names + ) + + assert flags == (False, True, True, False) + + +@pytest.mark.asyncio +async def test_split_mcp_tools_leaves_external_mcp_path_urls_for_the_provider(): + + async def served_names(names): + assert names == {"mcp"} + return frozenset() + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler.split_mcp_tools( + [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names + ) + + assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] + assert other_tools == [ZAPIER_TOOL, FUNCTION_TOOL] + + +@pytest.mark.asyncio +async def test_split_mcp_tools_repoints_served_proxy_urls_at_the_gateway(): + served_tool = { + "type": "mcp", + "server_label": "toolset", + "server_url": "http://localhost:4000/mcp/my-toolset", + "require_approval": "never", + "allowed_tools": ["get_me"], + } + unserved_tool = {"type": "mcp", "server_label": "typo", "server_url": "http://localhost:4000/mcp/githb"} + + async def served_names(names): + return frozenset({"my-toolset"}) + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler.split_mcp_tools( + [served_tool, unserved_tool], served_names=served_names + ) + + assert gateway_tools == [{**served_tool, "server_url": "litellm_proxy/mcp/my-toolset"}] + assert other_tools == [unserved_tool] + + +@pytest.mark.asyncio +async def test_split_mcp_tools_skips_resolution_when_nothing_points_at_the_proxy(): + async def served_names(names): + raise AssertionError("no lookup expected") + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler.split_mcp_tools( + [EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names + ) + + assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] + assert other_tools == [FUNCTION_TOOL] + + +def test_should_use_gateway_still_triggers_on_http_mcp_path(): + assert LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway([ZAPIER_TOOL]) is True + assert LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway([EXPLICIT_GATEWAY_TOOL]) is True + assert LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway([FUNCTION_TOOL]) is False + assert LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(None) is False + + +@pytest.mark.asyncio +async def test_aresponses_api_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.responses import main as responses_main + from litellm.types.llms.openai import ResponsesAPIResponse + + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None)) + monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {}) + provider_tools: list[object] = [] + + def fake_provider(**kwargs: object) -> object: + request_params = cast(dict[str, object], kwargs["response_api_optional_request_params"]) + provider_tools.append(request_params.get("tools")) + + async def respond() -> ResponsesAPIResponse: + return ResponsesAPIResponse(id="resp_zapier", created_at=0, output=[]) + + return respond() + + monkeypatch.setattr(responses_main.base_llm_http_handler, "response_api_handler", fake_provider) + + response = await responses_main.aresponses_api_with_mcp( + input="Reply with the single word ok.", model="openai/gpt-4.1", tools=[ZAPIER_TOOL] + ) + + assert isinstance(response, ResponsesAPIResponse) + assert provider_tools == [[ZAPIER_TOOL]] + + +def _response_with_reasoning_and_tool_call() -> Any: + """A first-turn response as a reasoning model returns it: reasoning item, then a function call.""" + return ResponsesAPIResponse( + id="resp_first", + created_at=1234567890, + model="gpt-5", + object="response", + status="completed", + output=[ + { + "type": "reasoning", + "id": "rs_1", + "summary": [], + "encrypted_content": "gAAAAA-opaque-blob", + }, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call-1", + "name": "foo", + "arguments": "{}", + "status": "completed", + }, + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + + +def test_create_follow_up_input_preserves_reasoning_when_stateless(): + """ + Regression test (LIT-5427): a store=false follow-up has to replay the reasoning + item, including reasoning.encrypted_content, since the provider kept no state. + """ + follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( + response=_response_with_reasoning_and_tool_call(), + tool_results=[{"tool_call_id": "call-1", "name": "foo", "result": "done"}], + original_input="hi", + preserve_reasoning=True, + ) + + assert follow_up[1] == { + "type": "reasoning", + "id": "rs_1", + "summary": [], + "encrypted_content": "gAAAAA-opaque-blob", + } + assert follow_up[2] == { + "type": "function_call", + "call_id": "call-1", + "name": "foo", + "arguments": "{}", + } + assert follow_up[3] == { + "type": "function_call_output", + "call_id": "call-1", + "output": "done", + } + + +def _response_with_interleaved_reasoning_and_tool_calls() -> Any: + """A first-turn response that reasons before each of two function calls.""" + return ResponsesAPIResponse( + id="resp_first", + created_at=1234567890, + model="gpt-5", + object="response", + status="completed", + output=[ + {"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "blob-1"}, + {"type": "function_call", "id": "fc_1", "call_id": "call-1", "name": "foo", "arguments": "{}"}, + {"type": "reasoning", "id": "rs_2", "summary": [], "encrypted_content": "blob-2"}, + {"type": "function_call", "id": "fc_2", "call_id": "call-2", "name": "bar", "arguments": "{}"}, + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + + +def test_create_follow_up_input_keeps_each_reasoning_item_before_its_function_call(): + """ + Regression test (LIT-5427): the provider pairs a replayed reasoning item with the + item that follows it, so the replay has to keep the response's output order instead + of grouping every reasoning item ahead of every function call. + """ + follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( + response=_response_with_interleaved_reasoning_and_tool_calls(), + tool_results=[ + {"tool_call_id": "call-1", "name": "foo", "result": "one"}, + {"tool_call_id": "call-2", "name": "bar", "result": "two"}, + ], + original_input="hi", + preserve_reasoning=True, + ) + + assert [cast(dict[str, Any], item)["type"] for item in follow_up] == [ + "message", + "reasoning", + "function_call", + "reasoning", + "function_call", + "function_call_output", + "function_call_output", + ] + assert [ + cast(dict[str, Any], item).get("id") or cast(dict[str, Any], item).get("call_id") for item in follow_up[1:5] + ] == [ + "rs_1", + "call-1", + "rs_2", + "call-2", + ] + + +def test_create_follow_up_input_omits_reasoning_when_stateful(): + """With store=true the provider still holds the reasoning item, so don't resend it.""" + follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( + response=_response_with_reasoning_and_tool_call(), + tool_results=[{"tool_call_id": "call-1", "name": "foo", "result": "done"}], + original_input="hi", + ) + + assert not [item for item in follow_up if isinstance(item, dict) and item.get("type") == "reasoning"] + + +@pytest.mark.parametrize( + "call_params, expected", + [ + ({"store": False}, True), + ({"store": True}, False), + ({"store": None}, False), + ({}, False), + ], +) +def test_is_persistence_disabled(call_params: dict[str, Any], expected: bool): + assert LiteLLM_Proxy_MCP_Handler.is_persistence_disabled(call_params) is expected + + +@pytest.mark.parametrize( + "store, caller_previous_response_id, expected_previous_response_id", + [ + (False, None, None), + (False, "resp_caller", "resp_caller"), + (True, None, "resp_first"), + (True, "resp_caller", "resp_first"), + ], +) +@pytest.mark.asyncio +async def test_mcp_follow_up_call_is_stateless_when_store_is_false( + monkeypatch: pytest.MonkeyPatch, + store: bool, + caller_previous_response_id: str | None, + expected_previous_response_id: str | None, +): + """ + Regression test (LIT-5427): linking the MCP follow-up call to the first response's id + fails for zero data retention callers, because store=false means it was never persisted. + The caller's own previous_response_id was valid for the first call, so it stays. + """ + captured_calls: list[dict[str, Any]] = [] + first_response = _response_with_reasoning_and_tool_call() + + async def fake_aresponses(**kwargs: Any) -> ResponsesAPIResponse: + captured_calls.append(kwargs) + return ( + first_response + if len(captured_calls) == 1 + else ResponsesAPIResponse( + id="resp_follow_up", + created_at=1234567891, + model="gpt-5", + object="response", + status="completed", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + ) + + async def fake_process(**kwargs: Any) -> tuple[list[Any], dict[str, str]]: + assert kwargs["raw_headers"] == {"x-app-id": "follow-up-caller"} + return ([], {"foo": "litellm_proxy"}) + + async def fake_execute(**kwargs: Any) -> list[dict[str, Any]]: + assert kwargs["guardrail_context"]["metadata"]["guardrails"] == ("block-all",) + assert kwargs["guardrail_context"]["model"] == "gpt-5" + return [{"tool_call_id": "call-1", "name": "foo", "result": "done"}] + + monkeypatch.setattr(responses_main, "aresponses", fake_aresponses) + monkeypatch.setattr(mcp_handler_module, "aresponses", fake_aresponses) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, "process_mcp_tools_without_openai_transform", staticmethod(fake_process) + ) + monkeypatch.setattr(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", staticmethod(fake_execute)) + + await responses_main.aresponses_api_with_mcp( + input="hi", + model="gpt-5", + tools=[{"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"}], + litellm_metadata={"guardrails": ["block-all"]}, + secret_fields={"raw_headers": {"x-app-id": "follow-up-caller"}}, + store=store, + previous_response_id=caller_previous_response_id, + ) + + assert len(captured_calls) == 2 + follow_up_call = captured_calls[1] + assert follow_up_call["previous_response_id"] == expected_previous_response_id + + reasoning_items = [ + item for item in follow_up_call["input"] if isinstance(item, dict) and item.get("type") == "reasoning" + ] + assert bool(reasoning_items) is (store is False) + + +@pytest.mark.asyncio +async def test_responses_discovery_logs_sanitized_caller_headers(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy._experimental.mcp_server import mcp_server_manager, operations + + headers: Final = { + "x-app-id": "app-a", + "x-nuid": "user-a", + "x-user-id": "identity-a", + "x-mcp-deepwiki-authorization": "upstream-sentinel", + "authorization": "proxy-sentinel", + } + manager: Final = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + get_allowed_mcp_servers=AsyncMock(return_value=[]), + get_mcp_servers_from_ids=MagicMock(return_value=[]), + ) + logger: Final = MagicMock(model_call_details={}) + logger.async_success_handler = AsyncMock() + setup: Final = MagicMock(return_value=(logger, None)) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + monkeypatch.setattr(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[])) + monkeypatch.setattr(operations, "function_setup", setup) + response: Final = ResponsesAPIResponse( + id="resp_test", + created_at=1234567891, + model="test-model", + object="response", + status="completed", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + monkeypatch.setattr(responses_main, "aresponses", AsyncMock(return_value=response)) + result: Final = await responses_main.aresponses_api_with_mcp( + input="hi", + model="test-model", + tools=[{"type": "mcp", "server_url": "litellm_proxy"}], + secret_fields={"raw_headers": headers}, + ) + assert result is response + logger.async_success_handler.assert_awaited_once() + logged: Final = setup.call_args.kwargs["metadata"]["headers"] + assert logged == {"x-app-id": "app-a", "x-nuid": "user-a", "x-user-id": "identity-a"} + assert headers["x-mcp-deepwiki-authorization"] == "upstream-sentinel" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("real_listing", [False, True]) +@pytest.mark.parametrize( + ("allowed_tools", "expected_names"), + [ + ([], ["responses_slot-echo", "responses_slot-status"]), + (["echo"], ["responses_slot-echo"]), + (["responses_slot-echo"], ["responses_slot-echo"]), + (["absent"], []), + ], +) +async def test_bridge_listing_leaves_the_callers_catalog_unchanged( + monkeypatch: pytest.MonkeyPatch, allowed_tools: list[str], expected_names: list[str], real_listing: bool +) -> None: + manager: Final = mcp_operations.global_mcp_server_manager + server: Final = MCPServer( + server_id="responses-slot", name="responses_slot", alias="responses_slot", transport=MCPTransport.http + ) + user: Final = UserAPIKeyAuth(api_key="sk-responses-slot", user_id="responder") + upstream: Final = [ + MCPTool(name="echo", description="Echo text back", inputSchema={"type": "object"}), + MCPTool(name="status", description="Report status", inputSchema={"type": "object"}), + MCPTool(name="echo", description="Duplicate echo", inputSchema={"type": "object", "properties": {}}), + ] + fake_manager: Final = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + get_allowed_mcp_servers=AsyncMock(return_value=[]), + get_mcp_servers_from_ids=MagicMock(return_value=[]), + get_mcp_server_by_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + with ( + patch.dict(manager.tool_name_to_mcp_server_name_mapping), + patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])), + patch.object(manager, "create_mcp_client", AsyncMock(return_value=object())), + patch.object(manager, "_fetch_tools_with_timeout", AsyncMock(return_value=upstream)), + ): + try: + if real_listing: + await manager._get_tools_from_server(server, user_api_key_auth=user, record_listing=True) + caller: Final = ListedToolsCaller(user_api_key_auth=user) + before: Final = { + tool.name: (listed.description, listed.input_schema) + for tool in upstream + if (listed := manager.get_listed_tool(server, tool.name, caller)) is not None + } + assert bool(before) is real_listing + tools, _server_names = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_without_openai_transform( + user_api_key_auth=user, + mcp_tools_with_litellm_proxy=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp/responses-slot", + "allowed_tools": allowed_tools, + } + ], + ) + recorded: Final = { + tool.name: (listed.description, listed.input_schema) + for tool in upstream + if (listed := manager.get_listed_tool(server, tool.name, caller)) is not None + } + assert recorded == before + assert ( + manager.get_listed_tool( + server, "echo", ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(api_key="sk-other-caller")) + ) + is None + ) + finally: + manager._drop_listed_tools(server.server_id) + + assert [tool.name for tool in tools] == expected_names + + +class _BridgeMetadataGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="bridge-metadata", event_hook=GuardrailEventHooks.pre_mcp_call, default_on=True) + self.calls: tuple[tuple[object, object], ...] = () + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Logging | None = None, + ) -> GenericGuardrailAPIInputs: + if request_data.get("mcp_arguments") == {"probe": "bridge"}: + self.calls += ((request_data.get("mcp_tool_description"), request_data.get("mcp_input_schema")),) + return inputs + + +@pytest.mark.asyncio +async def test_concurrent_bridge_calls_use_their_own_served_metadata(monkeypatch: pytest.MonkeyPatch) -> None: + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="bridge", name="bridge", transport=MCPTransport.http, url="http://upstream") + manager.registry = {server.server_id: server} + user: Final = UserAPIKeyAuth(api_key="sk-bridge", user_id="bridge-user") + upstream: Final = [ + MCPTool( + name="echo", + description="Echo text", + inputSchema={"type": "object", "properties": {"text": {"type": "string"}}}, + ), + MCPTool(name="status", description="Read status", inputSchema={"type": "object"}), + ] + client: Final = AsyncMock() + client.call_tool.return_value = CallToolResult(content=[TextContent(type="text", text="ok")]) + manager.create_mcp_client = AsyncMock(return_value=client) + manager._fetch_tools_with_timeout = AsyncMock(return_value=upstream) + guardrail: Final = _BridgeMetadataGuardrail() + logger: Final = ProxyLogging(user_api_key_cache=DualCache()) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", logger) + monkeypatch.setattr(mcp_operations, "global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])) + monkeypatch.setattr("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", manager) + first_listed: Final = asyncio.Event() + second_listed: Final = asyncio.Event() + + async def bridge(name: str, first: bool) -> None: + if not first: + await first_listed.wait() + tools, server_map = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_without_openai_transform( + user_api_key_auth=user, + mcp_tools_with_litellm_proxy=[ + {"type": "mcp", "server_url": "litellm_proxy/mcp/bridge", "allowed_tools": [name]} + ], + ) + (first_listed if first else second_listed).set() + await second_listed.wait() + result: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( + tool_server_map=server_map, + tool_calls=[ + {"type": "function_call", "name": f"bridge-{name}", "arguments": '{"probe":"bridge"}', "call_id": name} + ], + user_api_key_auth=user, + served_tools=tools, + ) + assert [entry["result"] for entry in result] == ["ok"] + + try: + await asyncio.gather(bridge("echo", True), bridge("status", False)) + assert sorted(guardrail.calls, key=str) == sorted( + ((tool.description, tool.input_schema) for tool in upstream), key=str + ) + await manager.call_tool("bridge", "echo", {"probe": "bridge"}, user_api_key_auth=user, proxy_logging_obj=logger) + assert guardrail.calls[-1] == (None, None) + await manager._get_tools_from_server( + server, user_api_key_auth=user, proxy_logging_obj=logger, record_listing=True + ) + await manager.call_tool("bridge", "echo", {"probe": "bridge"}, user_api_key_auth=user, proxy_logging_obj=logger) + assert guardrail.calls[-1] == (upstream[0].description, upstream[0].input_schema) + finally: + manager._drop_listed_tools(server.server_id) + ProxyLogging._callback_capabilities_cache.clear() + + +def _toolset_gateway_manager(toolset_id: str, server_id: str) -> types.SimpleNamespace: + return types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + get_allowed_mcp_servers=AsyncMock(return_value=[]), + get_mcp_servers_from_ids=MagicMock(return_value=[]), + get_mcp_server_by_name=MagicMock(return_value=None), + get_toolset_by_name_cached=AsyncMock(return_value=types.SimpleNamespace(toolset_id=toolset_id)), + resolve_toolset_tool_permissions=AsyncMock(return_value={server_id: ["add"]}), + ) + + +async def _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id: str) -> dict[str, object]: + from litellm.proxy._experimental.mcp_server.ui_session_utils import granted_toolset_ids + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth + + 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) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + _toolset_gateway_manager("ts-granted", "srv-1"), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + + async def team_permission(context: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable: + return LiteLLM_ObjectPermissionTable(object_permission_id="op-team", mcp_toolsets=[team_toolset_id]) + + async def granted_through_team(context: UserAPIKeyAuth) -> frozenset[str]: + return await granted_toolset_ids(context, team_object_permission=team_permission, require_key_access=False) + + team_key: Final = UserAPIKeyAuth(api_key="sk-team", team_id="team-1", user_role=LitellmUserRoles.INTERNAL_USER) + await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( + user_api_key_auth=team_key, + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/team-toolset"}], + granted_toolsets=granted_through_team, + ) + assert mock_get_tools.await_args is not None + return mock_get_tools.await_args.kwargs + + +@pytest.mark.asyncio +async def test_toolset_gateway_url_scopes_a_team_granted_toolset_for_a_key_without_its_own_grant(monkeypatch): + kwargs: Final = await _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id="ts-granted") + scoped = kwargs["user_api_key_auth"].object_permission + assert scoped is not None + assert scoped.mcp_servers == ["srv-1"] + assert scoped.mcp_tool_permissions == {"srv-1": ["add"]} + assert kwargs["mcp_servers"] is None + + +@pytest.mark.asyncio +async def test_toolset_gateway_url_skips_a_toolset_the_team_does_not_grant(monkeypatch): + kwargs: Final = await _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id="ts-other") + assert kwargs["user_api_key_auth"].object_permission is None + assert kwargs["mcp_servers"] is None + + +@pytest.mark.asyncio +async def test_apply_toolset_permissions_pins_the_auth_to_explicit_grants_only(monkeypatch: pytest.MonkeyPatch): + """A toolset gateway URL must not widen to operator-open (allow_all_keys) servers.""" + from litellm.proxy._types import UserAPIKeyAuth + + fake_manager = types.SimpleNamespace( + resolve_toolset_tool_permissions=AsyncMock(return_value={"srv-1": ["add"]}), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + scoped = await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions( + resolved_toolset_ids=["ts-1"], + resolved_mcp_servers=[], + user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", user_id="u1"), + ) + + assert scoped.mcp_explicit_grants_only is True + 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"]} + + +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( + "find_first,expected", + [ + (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, 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 + + +@pytest.mark.asyncio +async def test_apply_toolset_permissions_merges_into_existing_object_permission(monkeypatch): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + 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"], + 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_grants_no_toolset_servers_when_the_toolset_table_is_down(monkeypatch): + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + _toolset_prisma_client(find_many=AsyncMock(side_effect=RuntimeError("db down"))), + ) + 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=UserAPIKeyAuth(api_key="sk-test", user_id="u1"), + ) + + 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 +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] 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) 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 f43e15cb342..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,1576 +1,29 @@ -import asyncio import importlib -import subprocess import sys -import textwrap -import types -from contextlib import nullcontext -from typing import Any, Final, Literal, cast -from unittest.mock import AsyncMock, MagicMock, patch +from typing import Final import pytest -from fastapi import HTTPException -from mcp.types import CallToolResult, TextContent -from mcp.types import Tool as MCPTool -from openai.types.responses.tool_param import Mcp -import litellm -from litellm.caching.caching import DualCache -from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.litellm_core_utils.litellm_logging import Logging -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.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 ( - LiteLLM_Proxy_MCP_Handler, -) -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.responses.main import OutputFunctionToolCall -from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse +from litellm.proxy._experimental.mcp_server import litellm_proxy_mcp_handler as relocated +from litellm.responses.mcp import litellm_proxy_mcp_handler as legacy -class _DummyMCPResult: - def __init__(self): - self.content = [] +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) + reimported: Final = importlib.import_module(legacy.__name__) -def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: - """Patch MCP globals so _execute_tool_calls can run in tests.""" - proxy_module = types.SimpleNamespace(proxy_logging_obj=object(), prisma_client=None) - monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) + assert reimported is not legacy + assert relocated.__name__ not in sys.modules - fake_manager = types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - call_tool=AsyncMock(return_value=_DummyMCPResult()), - # Newer logging path calls this to enrich spend logs metadata - get_mcp_server_from_tool_name=MagicMock(return_value=None), - get_mcp_server_by_name=MagicMock(return_value=None), - ) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - fake_manager, - ) - return fake_manager.call_tool +@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 _setup_proxy_logging(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: - """Patch proxy_logging_obj so failure hook can be asserted.""" - proxy_logging_obj = MagicMock() - proxy_logging_obj.post_call_failure_hook = AsyncMock() - proxy_module = types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj, prisma_client=None) - monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) - return proxy_logging_obj.post_call_failure_hook - -def test_deduplicate_mcp_tools_single_allowed_server(): - tools = [{"name": "search"}, {"name": "search"}] # duplicate on purpose - - deduped, server_map = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools( - tools, - ["everything"], - ) - - assert len(deduped) == 1 - assert server_map == {"search": "everything"} - - -@pytest.mark.parametrize( - "tool_name,expected_server", - [ - ("alpha-tool", "alpha"), - ("beta-another_tool", "beta"), - ], -) -def test_deduplicate_mcp_tools_prefixed_names(tool_name, expected_server): - tools = [{"name": tool_name}] - - _, server_map = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools( - tools, - ["alpha", "beta"], - ) - - assert server_map[tool_name] == expected_server - - -def test_extract_tool_calls_from_chat_response_handles_tool_calls(): - response = ModelResponse( - id="resp-1", - choices=[ - { - "index": 0, - "message": { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call-123", - "type": "function", - "function": {"name": "foo", "arguments": "{}"}, - } - ], - }, - "finish_reason": "tool_calls", - } - ], - model="gpt", - created=0, - object="chat.completion", - ) - - tool_calls = LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response) - - assert len(tool_calls) == 1 - assert tool_calls[0]["function"]["name"] == "foo" - - -def test_create_follow_up_messages_for_chat_appends_tool_results(): - original_messages = [{"role": "user", "content": "hi"}] - response = ModelResponse( - id="resp-2", - choices=[ - { - "index": 0, - "message": { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call-abc", - "type": "function", - "function": {"name": "foo", "arguments": "{}"}, - } - ], - }, - } - ], - model="gpt", - created=0, - object="chat.completion", - ) - tool_results = [ - { - "tool_call_id": "call-abc", - "name": "foo", - "result": "done", - } - ] - - follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_messages_for_chat( - original_messages, - response, - tool_results, - ) - - assert follow_up[0]["role"] == "user" - assert follow_up[-1]["role"] == "tool" - assert follow_up[-1]["name"] == "foo" - assert follow_up[-1]["content"] == "done" - - -def test_transform_mcp_tools_to_openai_uses_chat_format(monkeypatch): - captured = {} - - def fake_transform_chat(tool): - captured.setdefault("chat", []).append(tool) - return {"chat": True} - - def fake_transform_responses(tool): - captured.setdefault("responses", []).append(tool) - return {"responses": True} - - monkeypatch.setattr( - "litellm.experimental_mcp_client.tools.transform_mcp_tool_to_openai_tool", - fake_transform_chat, - ) - monkeypatch.setattr( - "litellm.experimental_mcp_client.tools.transform_mcp_tool_to_openai_responses_api_tool", - fake_transform_responses, - ) - - chat_tools = LiteLLM_Proxy_MCP_Handler.transform_mcp_tools_to_openai(["tool"], target_format="chat") - resp_tools = LiteLLM_Proxy_MCP_Handler.transform_mcp_tools_to_openai(["tool"]) - - assert chat_tools == [{"chat": True}] - assert resp_tools == [{"responses": True}] - assert captured["chat"] == ["tool"] - assert captured["responses"] == ["tool"] - - -def test_create_follow_up_input_handles_response_function_tool_call(): - response = types.SimpleNamespace( - output=[ - OutputFunctionToolCall( - id="id", - type="function_call", - call_id="call-1", - name="foo", - arguments="{}", - status="completed", - ) - ] - ) - - follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( - response=cast(Any, response), - tool_results=[], - original_input=None, - ) - - assert follow_up == [ - { - "type": "function_call", - "call_id": "call-1", - "name": "foo", - "arguments": "{}", - } - ] - - -@pytest.mark.asyncio -async def test_execute_tool_calls_strips_server_prefix(monkeypatch): - call_tool_mock = _setup_mcp_call_environment(monkeypatch) - tool_name = "deepwiki-read_wiki_structure" - tool_calls = [ - { - "id": "call-1", - "function": {"name": tool_name, "arguments": "{}"}, - } - ] - - await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( - tool_server_map={tool_name: "deepwiki"}, - tool_calls=tool_calls, - user_api_key_auth=None, - ) - - assert call_tool_mock.await_count == 1 - assert call_tool_mock.await_args is not None - assert call_tool_mock.await_args.kwargs["name"] == "read_wiki_structure" - - -@pytest.mark.asyncio -async def test_execute_tool_calls_keeps_tool_name_without_prefix(monkeypatch): - call_tool_mock = _setup_mcp_call_environment(monkeypatch) - tool_name = "read_wiki_structure" - tool_calls = [ - { - "id": "call-2", - "function": {"name": tool_name, "arguments": "{}"}, - } - ] - - await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( - tool_server_map={tool_name: "deepwiki"}, - tool_calls=tool_calls, - user_api_key_auth=None, - ) - - assert call_tool_mock.await_count == 1 - assert call_tool_mock.await_args is not None - assert call_tool_mock.await_args.kwargs["name"] == tool_name - - -@pytest.mark.asyncio -async def test_execute_tool_calls_keeps_tool_name_when_equal_to_server(monkeypatch): - call_tool_mock = _setup_mcp_call_environment(monkeypatch) - tool_name = "echo" - tool_calls = [ - { - "id": "call-3", - "function": {"name": tool_name, "arguments": "{}"}, - } - ] - - await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( - tool_server_map={tool_name: "echo"}, - tool_calls=tool_calls, - user_api_key_auth=None, - ) - - assert call_tool_mock.await_count == 1 - assert call_tool_mock.await_args is not None - assert call_tool_mock.await_args.kwargs["name"] == tool_name - - -@pytest.mark.asyncio -async def test_execute_tool_calls_strips_prefix_when_alias_differs_from_server_name( - monkeypatch, -): - call_tool_mock = _setup_mcp_call_environment(monkeypatch) - fake_server = types.SimpleNamespace( - alias="my_deepwiki", - server_name="deepwiki_test", - server_id="test-server-id", - short_prefix=None, - mcp_info=None, - tool_name_to_display_name=None, - ) - from litellm.proxy._experimental.mcp_server import mcp_server_manager as _msm - - _msm.global_mcp_server_manager.get_mcp_server_from_tool_name = MagicMock(return_value=fake_server) - - tool_name = "my_deepwiki-read_wiki_structure" - tool_calls = [ - { - "id": "call-4", - "function": {"name": tool_name, "arguments": "{}"}, - } - ] - - await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( - tool_server_map={tool_name: "deepwiki_test"}, - tool_calls=tool_calls, - user_api_key_auth=None, - ) - - assert call_tool_mock.await_count == 1 - assert call_tool_mock.await_args is not None - assert call_tool_mock.await_args.kwargs["name"] == "read_wiki_structure" - - -@pytest.mark.asyncio -async def test_execute_tool_calls_reverse_maps_display_name(monkeypatch): - call_tool_mock = _setup_mcp_call_environment(monkeypatch) - colliding_server = types.SimpleNamespace( - alias=None, - server_name="other_mcp", - server_id="other-server-id", - short_prefix=None, - mcp_info=None, - tool_name_to_display_name={"search": "search_docs"}, - ) - fake_server = types.SimpleNamespace( - alias=None, - server_name="deepwiki_mcp", - server_id="test-server-id", - short_prefix=None, - mcp_info=None, - tool_name_to_display_name={"read_wiki_structure": "browse_repo_docs"}, - ) - from litellm.proxy._experimental.mcp_server import mcp_server_manager as _msm - - _msm.global_mcp_server_manager.get_mcp_server_from_tool_name = MagicMock(return_value=colliding_server) - _msm.global_mcp_server_manager.get_mcp_server_by_name = MagicMock(return_value=fake_server) - - tool_name = "browse_repo_docs" - tool_calls = [ - { - "id": "call-5", - "function": {"name": tool_name, "arguments": "{}"}, - } - ] - - await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( - tool_server_map={tool_name: "deepwiki_mcp"}, - tool_calls=tool_calls, - user_api_key_auth=None, - ) - - assert call_tool_mock.await_count == 1 - assert call_tool_mock.await_args is not None - assert call_tool_mock.await_args.kwargs["name"] == "read_wiki_structure" - - -@pytest.mark.asyncio -async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkeypatch): - """ - Regression test for ae4d92ad...: - Ensure responses-side MCP tool execution logs failures via proxy_logging_obj.post_call_failure_hook. - """ - post_call_failure_hook = _setup_proxy_logging(monkeypatch) - - fake_manager = types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - call_tool=AsyncMock(side_effect=HTTPException(status_code=500, detail="boom")), - ) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - fake_manager, - ) - - tool_name = "deepwiki-read_wiki_structure" - tool_calls = [{"id": "call-err", "function": {"name": tool_name, "arguments": "{}"}}] - - user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") - - results = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( - tool_server_map={tool_name: "deepwiki"}, - tool_calls=tool_calls, - user_api_key_auth=user_auth, - litellm_call_id="cid", - litellm_trace_id="tid", - ) - - assert len(results) == 1 - assert results[0]["tool_call_id"] == "call-err" - assert results[0]["name"] == tool_name - - post_call_failure_hook.assert_awaited_once() - assert post_call_failure_hook.await_args is not None - assert post_call_failure_hook.await_args.kwargs.get("route") == "/responses/mcp/call_tool" - - -@pytest.mark.asyncio -async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_function_setup( - monkeypatch, -): - """ - Regression test for ae4d92ad...: - Ensure litellm_call_id / litellm_trace_id are forwarded into function_setup kwargs. - """ - _setup_proxy_logging(monkeypatch) - call_tool_mock = _setup_mcp_call_environment(monkeypatch) - - captured = {} - - def fake_function_setup(*_args, **kwargs): - captured.update(kwargs) - return None, None - - # 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") - monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) - - tool_name = "deepwiki-read_wiki_structure" - tool_calls = [{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}] - - await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( - tool_server_map={tool_name: "deepwiki"}, - tool_calls=tool_calls, - user_api_key_auth=None, - litellm_call_id="cid", - litellm_trace_id="tid", - ) - - # Ensure the tool call was attempted (sanity) - assert call_tool_mock.await_count == 1 - - assert captured.get("litellm_call_id") == "cid" - assert captured.get("litellm_trace_id") == "tid" - - -@pytest.mark.asyncio -async def test_execute_tool_calls_threads_logging_obj_into_call_tool(monkeypatch): - """The Responses-API MCP path must hand the request's litellm_logging_obj to - global_mcp_server_manager.call_tool, otherwise pre_call_tool_check / - _create_during_hook_task get None and no guardrail evaluation is bridged onto - the request logger, so MCP tool calls made through the Responses API report zero - guardrail evaluations in the monitor. Drop the litellm_logging_obj kwarg on the - call_tool invocation and this fails.""" - _setup_proxy_logging(monkeypatch) - call_tool_mock = _setup_mcp_call_environment(monkeypatch) - - sentinel_logging_obj = MagicMock() - 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") - monkeypatch.setattr( - handler_module, - "function_setup", - lambda *_args, **_kwargs: (sentinel_logging_obj, None), - ) - - tool_name = "deepwiki-read_wiki_structure" - tool_calls = [{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}] - - await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( - tool_server_map={tool_name: "deepwiki"}, - tool_calls=tool_calls, - user_api_key_auth=None, - ) - - assert call_tool_mock.await_count == 1 - assert call_tool_mock.await_args is not None - assert call_tool_mock.await_args.kwargs["litellm_logging_obj"] is sentinel_logging_obj - - -@pytest.mark.asyncio -async def test_execute_tool_calls_applies_post_call_hook_content(monkeypatch): - proxy_module = types.SimpleNamespace(proxy_logging_obj=None) - monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) - - result = CallToolResult( - content=[TextContent(type="text", text="SECRET-1234")], - structuredContent={"result": "SECRET-1234"}, - isError=False, - ) - fake_manager = types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - call_tool=AsyncMock(return_value=result), - get_mcp_server_from_tool_name=MagicMock(return_value=None), - get_mcp_server_by_name=MagicMock(return_value=None), - ) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - fake_manager, - ) - - logging_obj = MagicMock() - logging_obj.model_call_details = {} - logging_obj.async_post_mcp_tool_call_hook = AsyncMock( - 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") - monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None)) - - tool_name = "deepwiki-read_wiki_structure" - results = 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": "[REDACTED]", "name": tool_name}] - assert logging_obj.async_success_handler.await_args.kwargs["result"].content[0].text == "[REDACTED]" - assert logging_obj.async_success_handler.await_args.kwargs["result"].structured_content is None - - -@pytest.mark.asyncio -async def test_execute_tool_calls_returns_proxy_result_without_logging(monkeypatch): - result = 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) - ) - - fake_manager = types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - call_tool=AsyncMock(return_value=result), - get_mcp_server_from_tool_name=MagicMock(return_value=None), - get_mcp_server_by_name=MagicMock(return_value=None), - ) - monkeypatch.setattr( - "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") - monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (None, None)) - - tool_name = "deepwiki-read_wiki_structure" - results = 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}] - proxy_logging_obj.post_mcp_call_hook.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_execute_tool_calls_passes_logging_details_to_proxy_hook(monkeypatch): - result = 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) - ) - - fake_manager = types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - call_tool=AsyncMock(return_value=result), - get_mcp_server_from_tool_name=MagicMock(return_value=None), - get_mcp_server_by_name=MagicMock(return_value=None), - ) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - fake_manager, - ) - logging_obj = MagicMock() - 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") - monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None)) - - tool_name = "deepwiki-read_wiki_structure" - results = 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}] - assert proxy_logging_obj.post_mcp_call_hook.await_args.kwargs["request_data"] == logging_obj.model_call_details - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure_stage", ["post_call_hook", "success_handler"]) -async def test_execute_tool_calls_continues_when_post_call_logging_fails(monkeypatch, failure_stage: str): - proxy_module = types.SimpleNamespace(proxy_logging_obj=None) - monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) - - result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) - fake_manager = types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - call_tool=AsyncMock(return_value=result), - get_mcp_server_from_tool_name=MagicMock(return_value=None), - get_mcp_server_by_name=MagicMock(return_value=None), - ) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - fake_manager, - ) - - logging_obj = MagicMock() - logging_obj.model_call_details = {} - logging_obj.post_call = MagicMock() - logging_obj.async_post_mcp_tool_call_hook = AsyncMock( - side_effect=RuntimeError("hook failed") if failure_stage == "post_call_hook" else None, - return_value=result, - ) - 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") - monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None)) - - tool_name = "deepwiki-read_wiki_structure" - results = 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}] - - -@pytest.mark.asyncio -async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch): - """ - Regression test for 872e5b98...: - Ensure responses-side tool discovery enables list-tools SpendLogs logging flags. - """ - served_tools: Final = [ - MCPTool(name="safe", description="Safe lookup", inputSchema={"type": "object"}), - MCPTool( - name="masked", - description="Contact [MASKED]", - inputSchema={"type": "object", "properties": {"query": {"type": "string", "description": "For [MASKED]"}}}, - ), - ] - mock_get_tools = AsyncMock(return_value=AggregateToolListing(tools=served_tools, outcomes={})) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", - mock_get_tools, - ) - - # Patch manager methods used by _get_mcp_tools_from_manager to avoid needing full UserAPIKeyAuth fields. - fake_manager = types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - get_allowed_mcp_servers=AsyncMock(return_value=[]), - get_mcp_servers_from_ids=MagicMock(return_value=[]), - get_mcp_server_by_name=MagicMock(return_value=None), - ) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - fake_manager, - ) - - user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") - tools, _server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( - user_api_key_auth=user_auth, - mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}], - ) - - forwarded: Final = LiteLLM_Proxy_MCP_Handler.transform_mcp_tools_to_openai(tools) - assert [tool["name"] for tool in forwarded] == ["safe", "masked"] - assert forwarded[0]["description"] == "Safe lookup" - assert forwarded[1]["description"] == "Contact [MASKED]" - assert forwarded[1]["parameters"] == { - "type": "object", - "properties": {"query": {"type": "string", "description": "For [MASKED]"}}, - "additionalProperties": False, - } - assert mock_get_tools.await_count == 1 - assert mock_get_tools.await_args is not None - assert mock_get_tools.await_args.kwargs["log_list_tools_to_spendlogs"] is True - assert mock_get_tools.await_args.kwargs["list_tools_log_source"] == "responses" - - -def test_get_parent_request_tags_from_metadata(): - tags = LiteLLM_Proxy_MCP_Handler.get_parent_request_tags({"metadata": {"tags": ["team-a", "prod"]}}) - assert tags == ["team-a", "prod"] - - -def test_get_parent_request_tags_from_nested_litellm_params(): - tags = LiteLLM_Proxy_MCP_Handler.get_parent_request_tags( - { - "metadata": {"tags": ["top-level"]}, - "litellm_params": { - "metadata": {"tags": ["nested"]}, - "proxy_server_request": {"headers": {"user-agent": "client/1.0"}}, - }, - } - ) - assert tags == ["nested", "User-Agent: client", "User-Agent: client/1.0"] - - -@pytest.mark.asyncio -async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch): - 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, - ) - fake_manager = types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - get_allowed_mcp_servers=AsyncMock(return_value=[]), - get_mcp_servers_from_ids=MagicMock(return_value=[]), - get_mcp_server_by_name=MagicMock(return_value=None), - ) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - fake_manager, - ) - - await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( - user_api_key_auth=types.SimpleNamespace(api_key="k", user_id="u"), - mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}], - request_tags=["team-a"], - ) - - assert mock_get_tools.await_args.kwargs["request_tags"] == ["team-a"] - - -@pytest.mark.asyncio -async def test_execute_tool_calls_exposes_sanitized_client_headers_to_logging(monkeypatch): - """The Responses API MCP bridge used to log an empty header dict, hiding the caller's - headers from logging callbacks and hooks.""" - _setup_proxy_logging(monkeypatch) - _setup_mcp_call_environment(monkeypatch) - - captured = {} - - def fake_function_setup(*_args, **kwargs): - captured.update(kwargs) - return None, None - - handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") - monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) - - tool_name = "deepwiki-read_wiki_structure" - 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, - raw_headers={"x-nuid": "nuid-1", "x-litellm-api-key": "sk-proxy", "cookie": "s=1"}, - ) - - expected = {"x-nuid": "nuid-1", "cookie": "***REDACTED***"} - assert captured["metadata"]["headers"] == expected - assert captured["proxy_server_request"]["headers"] == expected - - -@pytest.mark.asyncio -async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monkeypatch): - _setup_proxy_logging(monkeypatch) - _setup_mcp_call_environment(monkeypatch) - captured = {} - - def fake_function_setup(*_args, **kwargs): - captured.update(kwargs) - return None, None - - handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") - monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) - - tool_name = "deepwiki-read_wiki_structure" - 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, - request_tags=["team-a", "prod"], - ) - - assert captured["metadata"]["tags"] == ["team-a", "prod"] - - -def test_completion_with_function_tools_works_without_fastapi_installed(): - script = textwrap.dedent( - """ - import sys - - class _FastapiBlocker: - def find_spec(self, fullname, path=None, target=None): - if fullname == "fastapi" or fullname.startswith("fastapi."): - raise ModuleNotFoundError("No module named 'fastapi'") - return None - - sys.meta_path.insert(0, _FastapiBlocker()) - - import litellm - - response = litellm.completion( - model="openai/gpt-5.5", - messages=[{"role": "user", "content": "What is the weather in SF?"}], - tools=[ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the current weather for a location", - "parameters": { - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"], - }, - }, - } - ], - mock_response="sunny", - ) - assert response.choices[0].message.content == "sunny" - """ - ) - result = subprocess.run( - [sys.executable, "-c", script], - capture_output=True, - text=True, - timeout=120, - ) - assert result.returncode == 0, result.stderr - - -def test_extract_tool_call_details_reads_anthropic_tool_use_input(): - """ - Regression test (LIT-4517): an Anthropic tool_use block carries its arguments - under `input`, not `arguments`. - - Given: A tool_use content block as /v1/messages returns it - When: The shared extractor reads it - Then: The arguments come back, so the MCP tool is called with them - - Reading only `arguments` fails silently rather than loudly: _parse_tool_arguments - turns the resulting None into {}, so the tool still executes, just with every - argument dropped. - """ - tool_use_block = { - "type": "tool_use", - "id": "toolu_01ABC", - "name": "read_wiki_structure", - "input": {"repoName": "BerriAI/litellm"}, - } - - name, arguments, call_id = LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_use_block) - - assert name == "read_wiki_structure" - assert call_id == "toolu_01ABC" - assert arguments == {"repoName": "BerriAI/litellm"} - assert LiteLLM_Proxy_MCP_Handler._parse_tool_arguments(arguments) == {"repoName": "BerriAI/litellm"} - - -def test_extract_tool_call_details_still_prefers_openai_arguments(): - """The OpenAI chat shape must keep winning; `input` is only the fallback.""" - openai_tool_call = { - "id": "call_123", - "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, - } - - name, arguments, call_id = LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(openai_tool_call) - - assert name == "get_weather" - assert call_id == "call_123" - assert arguments == '{"city": "Paris"}' - - -def _registered( - server_id: str, - name: str, - alias: str | None = None, - server_name: str | None = None, - access_groups: list[str] | None = None, -): - from litellm.types.mcp import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer - - return MCPServer( - server_id=server_id, - name=name, - alias=alias, - server_name=server_name, - transport=MCPTransport.http, - access_groups=access_groups, - ) - - -async def _no_toolset(_: str) -> bool: - return False - - -ZAPIER_TOOL: Mcp = { - "type": "mcp", - "server_label": "zapier", - "server_url": "https://mcp.zapier.com/api/mcp/mcp", - "require_approval": "never", -} -EXPLICIT_GATEWAY_TOOL = {"type": "mcp", "server_label": "github", "server_url": "litellm_proxy/mcp/github"} -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 - - servers = ( - _registered("id-1", "github-name", alias="github", server_name="github-server", access_groups=["prod-group"]), - _registered("id-2", "deepwiki"), - ) - - async def toolset_exists(name: str) -> bool: - return name == "my-toolset" - - served = await _gateway_served_names( - {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset", "mcp", "nope"}, - servers=lambda: servers, - toolset_exists=toolset_exists, - ) - - assert served == {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset"} - - -@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 - - server_id = "0b9ae4ca-1bd2-4faa-b183-7dd812597e3b" - short_prefix = compute_short_server_prefix(server_id) - servers = (_registered(server_id, "github-name", alias="github", access_groups=["prod-group"]),) - - served = await _gateway_served_names( - {server_id, short_prefix, "GitHub", "PROD-GROUP", "nope"}, servers=lambda: servers, toolset_exists=_no_toolset - ) - - assert served == {server_id, short_prefix, "GitHub"} - - -@pytest.mark.asyncio -async def test_routes_through_gateway_flags_explicit_and_served_tools_only(): - served_tool = {"type": "mcp", "server_label": "github", "server_url": "http://localhost:4000/mcp/github"} - - async def served_names(names): - assert names == {"github", "mcp"} - return frozenset({"github"}) - - flags = await LiteLLM_Proxy_MCP_Handler.routes_through_gateway( - [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, served_tool, FUNCTION_TOOL], served_names=served_names - ) - - assert flags == (False, True, True, False) - - -@pytest.mark.asyncio -async def test_split_mcp_tools_leaves_external_mcp_path_urls_for_the_provider(): - - async def served_names(names): - assert names == {"mcp"} - return frozenset() - - gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler.split_mcp_tools( - [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names - ) - - assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] - assert other_tools == [ZAPIER_TOOL, FUNCTION_TOOL] - - -@pytest.mark.asyncio -async def test_split_mcp_tools_repoints_served_proxy_urls_at_the_gateway(): - served_tool = { - "type": "mcp", - "server_label": "toolset", - "server_url": "http://localhost:4000/mcp/my-toolset", - "require_approval": "never", - "allowed_tools": ["get_me"], - } - unserved_tool = {"type": "mcp", "server_label": "typo", "server_url": "http://localhost:4000/mcp/githb"} - - async def served_names(names): - return frozenset({"my-toolset"}) - - gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler.split_mcp_tools( - [served_tool, unserved_tool], served_names=served_names - ) - - assert gateway_tools == [{**served_tool, "server_url": "litellm_proxy/mcp/my-toolset"}] - assert other_tools == [unserved_tool] - - -@pytest.mark.asyncio -async def test_split_mcp_tools_skips_resolution_when_nothing_points_at_the_proxy(): - async def served_names(names): - raise AssertionError("no lookup expected") - - gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler.split_mcp_tools( - [EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names - ) - - assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] - assert other_tools == [FUNCTION_TOOL] - - -def test_should_use_gateway_still_triggers_on_http_mcp_path(): - assert LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway([ZAPIER_TOOL]) is True - assert LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway([EXPLICIT_GATEWAY_TOOL]) is True - assert LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway([FUNCTION_TOOL]) is False - assert LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(None) is False - - -@pytest.mark.asyncio -async def test_aresponses_api_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch): - from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager - from litellm.responses import main as responses_main - from litellm.types.llms.openai import ResponsesAPIResponse - - monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None)) - monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {}) - provider_tools: list[object] = [] - - def fake_provider(**kwargs: object) -> object: - request_params = cast(dict[str, object], kwargs["response_api_optional_request_params"]) - provider_tools.append(request_params.get("tools")) - - async def respond() -> ResponsesAPIResponse: - return ResponsesAPIResponse(id="resp_zapier", created_at=0, output=[]) - - return respond() - - monkeypatch.setattr(responses_main.base_llm_http_handler, "response_api_handler", fake_provider) - - response = await responses_main.aresponses_api_with_mcp( - input="Reply with the single word ok.", model="openai/gpt-4.1", tools=[ZAPIER_TOOL] - ) - - assert isinstance(response, ResponsesAPIResponse) - assert provider_tools == [[ZAPIER_TOOL]] - - -def _response_with_reasoning_and_tool_call() -> Any: - """A first-turn response as a reasoning model returns it: reasoning item, then a function call.""" - return ResponsesAPIResponse( - id="resp_first", - created_at=1234567890, - model="gpt-5", - object="response", - status="completed", - output=[ - { - "type": "reasoning", - "id": "rs_1", - "summary": [], - "encrypted_content": "gAAAAA-opaque-blob", - }, - { - "type": "function_call", - "id": "fc_1", - "call_id": "call-1", - "name": "foo", - "arguments": "{}", - "status": "completed", - }, - ], - parallel_tool_calls=False, - tool_choice="auto", - tools=[], - ) - - -def test_create_follow_up_input_preserves_reasoning_when_stateless(): - """ - Regression test (LIT-5427): a store=false follow-up has to replay the reasoning - item, including reasoning.encrypted_content, since the provider kept no state. - """ - follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( - response=_response_with_reasoning_and_tool_call(), - tool_results=[{"tool_call_id": "call-1", "name": "foo", "result": "done"}], - original_input="hi", - preserve_reasoning=True, - ) - - assert follow_up[1] == { - "type": "reasoning", - "id": "rs_1", - "summary": [], - "encrypted_content": "gAAAAA-opaque-blob", - } - assert follow_up[2] == { - "type": "function_call", - "call_id": "call-1", - "name": "foo", - "arguments": "{}", - } - assert follow_up[3] == { - "type": "function_call_output", - "call_id": "call-1", - "output": "done", - } - - -def _response_with_interleaved_reasoning_and_tool_calls() -> Any: - """A first-turn response that reasons before each of two function calls.""" - return ResponsesAPIResponse( - id="resp_first", - created_at=1234567890, - model="gpt-5", - object="response", - status="completed", - output=[ - {"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "blob-1"}, - {"type": "function_call", "id": "fc_1", "call_id": "call-1", "name": "foo", "arguments": "{}"}, - {"type": "reasoning", "id": "rs_2", "summary": [], "encrypted_content": "blob-2"}, - {"type": "function_call", "id": "fc_2", "call_id": "call-2", "name": "bar", "arguments": "{}"}, - ], - parallel_tool_calls=False, - tool_choice="auto", - tools=[], - ) - - -def test_create_follow_up_input_keeps_each_reasoning_item_before_its_function_call(): - """ - Regression test (LIT-5427): the provider pairs a replayed reasoning item with the - item that follows it, so the replay has to keep the response's output order instead - of grouping every reasoning item ahead of every function call. - """ - follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( - response=_response_with_interleaved_reasoning_and_tool_calls(), - tool_results=[ - {"tool_call_id": "call-1", "name": "foo", "result": "one"}, - {"tool_call_id": "call-2", "name": "bar", "result": "two"}, - ], - original_input="hi", - preserve_reasoning=True, - ) - - assert [cast(dict[str, Any], item)["type"] for item in follow_up] == [ - "message", - "reasoning", - "function_call", - "reasoning", - "function_call", - "function_call_output", - "function_call_output", - ] - assert [ - cast(dict[str, Any], item).get("id") or cast(dict[str, Any], item).get("call_id") for item in follow_up[1:5] - ] == [ - "rs_1", - "call-1", - "rs_2", - "call-2", - ] - - -def test_create_follow_up_input_omits_reasoning_when_stateful(): - """With store=true the provider still holds the reasoning item, so don't resend it.""" - follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_input( - response=_response_with_reasoning_and_tool_call(), - tool_results=[{"tool_call_id": "call-1", "name": "foo", "result": "done"}], - original_input="hi", - ) - - assert not [item for item in follow_up if isinstance(item, dict) and item.get("type") == "reasoning"] - - -@pytest.mark.parametrize( - "call_params, expected", - [ - ({"store": False}, True), - ({"store": True}, False), - ({"store": None}, False), - ({}, False), - ], -) -def test_is_persistence_disabled(call_params: dict[str, Any], expected: bool): - assert LiteLLM_Proxy_MCP_Handler.is_persistence_disabled(call_params) is expected - - -@pytest.mark.parametrize( - "store, caller_previous_response_id, expected_previous_response_id", - [ - (False, None, None), - (False, "resp_caller", "resp_caller"), - (True, None, "resp_first"), - (True, "resp_caller", "resp_first"), - ], -) -@pytest.mark.asyncio -async def test_mcp_follow_up_call_is_stateless_when_store_is_false( - monkeypatch: pytest.MonkeyPatch, - store: bool, - caller_previous_response_id: str | None, - expected_previous_response_id: str | None, -): - """ - Regression test (LIT-5427): linking the MCP follow-up call to the first response's id - fails for zero data retention callers, because store=false means it was never persisted. - The caller's own previous_response_id was valid for the first call, so it stays. - """ - captured_calls: list[dict[str, Any]] = [] - first_response = _response_with_reasoning_and_tool_call() - - async def fake_aresponses(**kwargs: Any) -> ResponsesAPIResponse: - captured_calls.append(kwargs) - return ( - first_response - if len(captured_calls) == 1 - else ResponsesAPIResponse( - id="resp_follow_up", - created_at=1234567891, - model="gpt-5", - object="response", - status="completed", - output=[], - parallel_tool_calls=False, - tool_choice="auto", - tools=[], - ) - ) - - async def fake_process(**kwargs: Any) -> tuple[list[Any], dict[str, str]]: - assert kwargs["raw_headers"] == {"x-app-id": "follow-up-caller"} - return ([], {"foo": "litellm_proxy"}) - - async def fake_execute(**kwargs: Any) -> list[dict[str, Any]]: - assert kwargs["guardrail_context"]["metadata"]["guardrails"] == ("block-all",) - assert kwargs["guardrail_context"]["model"] == "gpt-5" - return [{"tool_call_id": "call-1", "name": "foo", "result": "done"}] - - monkeypatch.setattr(responses_main, "aresponses", fake_aresponses) - monkeypatch.setattr(mcp_handler_module, "aresponses", fake_aresponses) - monkeypatch.setattr( - LiteLLM_Proxy_MCP_Handler, "process_mcp_tools_without_openai_transform", staticmethod(fake_process) - ) - monkeypatch.setattr(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", staticmethod(fake_execute)) - - await responses_main.aresponses_api_with_mcp( - input="hi", - model="gpt-5", - tools=[{"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"}], - litellm_metadata={"guardrails": ["block-all"]}, - secret_fields={"raw_headers": {"x-app-id": "follow-up-caller"}}, - store=store, - previous_response_id=caller_previous_response_id, - ) - - assert len(captured_calls) == 2 - follow_up_call = captured_calls[1] - assert follow_up_call["previous_response_id"] == expected_previous_response_id - - reasoning_items = [ - item for item in follow_up_call["input"] if isinstance(item, dict) and item.get("type") == "reasoning" - ] - assert bool(reasoning_items) is (store is False) - - -@pytest.mark.asyncio -async def test_responses_discovery_logs_sanitized_caller_headers(monkeypatch: pytest.MonkeyPatch): - from litellm.proxy._experimental.mcp_server import mcp_server_manager, operations - - headers: Final = { - "x-app-id": "app-a", - "x-nuid": "user-a", - "x-user-id": "identity-a", - "x-mcp-deepwiki-authorization": "upstream-sentinel", - "authorization": "proxy-sentinel", - } - manager: Final = types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - get_allowed_mcp_servers=AsyncMock(return_value=[]), - get_mcp_servers_from_ids=MagicMock(return_value=[]), - ) - logger: Final = MagicMock(model_call_details={}) - logger.async_success_handler = AsyncMock() - setup: Final = MagicMock(return_value=(logger, None)) - monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) - monkeypatch.setattr(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[])) - monkeypatch.setattr(operations, "function_setup", setup) - response: Final = ResponsesAPIResponse( - id="resp_test", - created_at=1234567891, - model="test-model", - object="response", - status="completed", - output=[], - parallel_tool_calls=False, - tool_choice="auto", - tools=[], - ) - monkeypatch.setattr(responses_main, "aresponses", AsyncMock(return_value=response)) - result: Final = await responses_main.aresponses_api_with_mcp( - input="hi", - model="test-model", - tools=[{"type": "mcp", "server_url": "litellm_proxy"}], - secret_fields={"raw_headers": headers}, - ) - assert result is response - logger.async_success_handler.assert_awaited_once() - logged: Final = setup.call_args.kwargs["metadata"]["headers"] - assert logged == {"x-app-id": "app-a", "x-nuid": "user-a", "x-user-id": "identity-a"} - assert headers["x-mcp-deepwiki-authorization"] == "upstream-sentinel" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("real_listing", [False, True]) -@pytest.mark.parametrize( - ("allowed_tools", "expected_names"), - [ - ([], ["responses_slot-echo", "responses_slot-status"]), - (["echo"], ["responses_slot-echo"]), - (["responses_slot-echo"], ["responses_slot-echo"]), - (["absent"], []), - ], -) -async def test_bridge_listing_leaves_the_callers_catalog_unchanged( - monkeypatch: pytest.MonkeyPatch, allowed_tools: list[str], expected_names: list[str], real_listing: bool -) -> None: - manager: Final = mcp_operations.global_mcp_server_manager - server: Final = MCPServer( - server_id="responses-slot", name="responses_slot", alias="responses_slot", transport=MCPTransport.http - ) - user: Final = UserAPIKeyAuth(api_key="sk-responses-slot", user_id="responder") - upstream: Final = [ - MCPTool(name="echo", description="Echo text back", inputSchema={"type": "object"}), - MCPTool(name="status", description="Report status", inputSchema={"type": "object"}), - MCPTool(name="echo", description="Duplicate echo", inputSchema={"type": "object", "properties": {}}), - ] - fake_manager: Final = types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - get_allowed_mcp_servers=AsyncMock(return_value=[]), - get_mcp_servers_from_ids=MagicMock(return_value=[]), - get_mcp_server_by_name=MagicMock(return_value=None), - ) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - fake_manager, - ) - with ( - patch.dict(manager.tool_name_to_mcp_server_name_mapping), - patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])), - patch.object(manager, "create_mcp_client", AsyncMock(return_value=object())), - patch.object(manager, "_fetch_tools_with_timeout", AsyncMock(return_value=upstream)), - ): - try: - if real_listing: - await manager._get_tools_from_server(server, user_api_key_auth=user, record_listing=True) - caller: Final = ListedToolsCaller(user_api_key_auth=user) - before: Final = { - tool.name: (listed.description, listed.input_schema) - for tool in upstream - if (listed := manager.get_listed_tool(server, tool.name, caller)) is not None - } - assert bool(before) is real_listing - tools, _server_names = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_without_openai_transform( - user_api_key_auth=user, - mcp_tools_with_litellm_proxy=[ - { - "type": "mcp", - "server_url": "litellm_proxy/mcp/responses-slot", - "allowed_tools": allowed_tools, - } - ], - ) - recorded: Final = { - tool.name: (listed.description, listed.input_schema) - for tool in upstream - if (listed := manager.get_listed_tool(server, tool.name, caller)) is not None - } - assert recorded == before - assert ( - manager.get_listed_tool( - server, "echo", ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(api_key="sk-other-caller")) - ) - is None - ) - finally: - manager._drop_listed_tools(server.server_id) - - assert [tool.name for tool in tools] == expected_names - - -class _BridgeMetadataGuardrail(CustomGuardrail): - def __init__(self) -> None: - super().__init__(guardrail_name="bridge-metadata", event_hook=GuardrailEventHooks.pre_mcp_call, default_on=True) - self.calls: tuple[tuple[object, object], ...] = () - - async def apply_guardrail( - self, - inputs: GenericGuardrailAPIInputs, - request_data: dict[str, object], - input_type: Literal["request", "response"], - logging_obj: Logging | None = None, - ) -> GenericGuardrailAPIInputs: - if request_data.get("mcp_arguments") == {"probe": "bridge"}: - self.calls += ((request_data.get("mcp_tool_description"), request_data.get("mcp_input_schema")),) - return inputs - - -@pytest.mark.asyncio -async def test_concurrent_bridge_calls_use_their_own_served_metadata(monkeypatch: pytest.MonkeyPatch) -> None: - manager: Final = MCPServerManager() - server: Final = MCPServer(server_id="bridge", name="bridge", transport=MCPTransport.http, url="http://upstream") - manager.registry = {server.server_id: server} - user: Final = UserAPIKeyAuth(api_key="sk-bridge", user_id="bridge-user") - upstream: Final = [ - MCPTool( - name="echo", - description="Echo text", - inputSchema={"type": "object", "properties": {"text": {"type": "string"}}}, - ), - MCPTool(name="status", description="Read status", inputSchema={"type": "object"}), - ] - client: Final = AsyncMock() - client.call_tool.return_value = CallToolResult(content=[TextContent(type="text", text="ok")]) - manager.create_mcp_client = AsyncMock(return_value=client) - manager._fetch_tools_with_timeout = AsyncMock(return_value=upstream) - guardrail: Final = _BridgeMetadataGuardrail() - logger: Final = ProxyLogging(user_api_key_cache=DualCache()) - monkeypatch.setattr(litellm, "callbacks", [guardrail]) - monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", logger) - monkeypatch.setattr(mcp_operations, "global_mcp_server_manager", manager) - monkeypatch.setattr(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])) - monkeypatch.setattr("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", manager) - first_listed: Final = asyncio.Event() - second_listed: Final = asyncio.Event() - - async def bridge(name: str, first: bool) -> None: - if not first: - await first_listed.wait() - tools, server_map = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_without_openai_transform( - user_api_key_auth=user, - mcp_tools_with_litellm_proxy=[ - {"type": "mcp", "server_url": "litellm_proxy/mcp/bridge", "allowed_tools": [name]} - ], - ) - (first_listed if first else second_listed).set() - await second_listed.wait() - result: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( - tool_server_map=server_map, - tool_calls=[ - {"type": "function_call", "name": f"bridge-{name}", "arguments": '{"probe":"bridge"}', "call_id": name} - ], - user_api_key_auth=user, - served_tools=tools, - ) - assert [entry["result"] for entry in result] == ["ok"] - - try: - await asyncio.gather(bridge("echo", True), bridge("status", False)) - assert sorted(guardrail.calls, key=str) == sorted( - ((tool.description, tool.input_schema) for tool in upstream), key=str - ) - await manager.call_tool("bridge", "echo", {"probe": "bridge"}, user_api_key_auth=user, proxy_logging_obj=logger) - assert guardrail.calls[-1] == (None, None) - await manager._get_tools_from_server( - server, user_api_key_auth=user, proxy_logging_obj=logger, record_listing=True - ) - await manager.call_tool("bridge", "echo", {"probe": "bridge"}, user_api_key_auth=user, proxy_logging_obj=logger) - assert guardrail.calls[-1] == (upstream[0].description, upstream[0].input_schema) - finally: - manager._drop_listed_tools(server.server_id) - ProxyLogging._callback_capabilities_cache.clear() - - -def _toolset_gateway_manager(toolset_id: str, server_id: str) -> types.SimpleNamespace: - return types.SimpleNamespace( - catalog=types.SimpleNamespace(operation=nullcontext), - get_registry=MagicMock(return_value={}), - get_allowed_mcp_servers=AsyncMock(return_value=[]), - get_mcp_servers_from_ids=MagicMock(return_value=[]), - get_mcp_server_by_name=MagicMock(return_value=None), - get_toolset_by_name_cached=AsyncMock(return_value=types.SimpleNamespace(toolset_id=toolset_id)), - resolve_toolset_tool_permissions=AsyncMock(return_value={server_id: ["add"]}), - ) - - -async def _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id: str) -> dict[str, object]: - from litellm.proxy._experimental.mcp_server.ui_session_utils import granted_toolset_ids - from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth - - 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) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - _toolset_gateway_manager("ts-granted", "srv-1"), - ) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) - - async def team_permission(context: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable: - return LiteLLM_ObjectPermissionTable(object_permission_id="op-team", mcp_toolsets=[team_toolset_id]) - - async def granted_through_team(context: UserAPIKeyAuth) -> frozenset[str]: - return await granted_toolset_ids(context, team_object_permission=team_permission, require_key_access=False) - - team_key: Final = UserAPIKeyAuth(api_key="sk-team", team_id="team-1", user_role=LitellmUserRoles.INTERNAL_USER) - await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( - user_api_key_auth=team_key, - mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/team-toolset"}], - granted_toolsets=granted_through_team, - ) - assert mock_get_tools.await_args is not None - return mock_get_tools.await_args.kwargs - - -@pytest.mark.asyncio -async def test_toolset_gateway_url_scopes_a_team_granted_toolset_for_a_key_without_its_own_grant(monkeypatch): - kwargs: Final = await _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id="ts-granted") - scoped = kwargs["user_api_key_auth"].object_permission - assert scoped is not None - assert scoped.mcp_servers == ["srv-1"] - assert scoped.mcp_tool_permissions == {"srv-1": ["add"]} - assert kwargs["mcp_servers"] is None - - -@pytest.mark.asyncio -async def test_toolset_gateway_url_skips_a_toolset_the_team_does_not_grant(monkeypatch): - kwargs: Final = await _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id="ts-other") - assert kwargs["user_api_key_auth"].object_permission is None - assert kwargs["mcp_servers"] is None - - -@pytest.mark.asyncio -async def test_apply_toolset_permissions_pins_the_auth_to_explicit_grants_only(monkeypatch: pytest.MonkeyPatch): - """A toolset gateway URL must not widen to operator-open (allow_all_keys) servers.""" - from litellm.proxy._types import UserAPIKeyAuth - - fake_manager = types.SimpleNamespace( - resolve_toolset_tool_permissions=AsyncMock(return_value={"srv-1": ["add"]}), - ) - monkeypatch.setattr( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", - fake_manager, - ) - - scoped = await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions( - resolved_toolset_ids=["ts-1"], - resolved_mcp_servers=[], - user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", user_id="u1"), - ) - - assert scoped.mcp_explicit_grants_only is True - 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"]} +def test_legacy_path_still_raises_for_unknown_attributes() -> None: + with pytest.raises(AttributeError): + getattr(legacy, "NotAHandlerAttribute")