diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 06d369eabcd..592d8edf6b8 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -130,6 +130,10 @@ jobs: echo "File content around line 43:" head -50 litellm/litellm_core_utils/custom_logger_registry.py | tail -10 + - name: Check MCP operation boundary + if: steps.changes.outputs.decision != 'skip' + run: uv run --no-sync python scripts/check_mcp_operation_boundary.py + - name: Run Ruff linting if: steps.changes.outputs.decision != 'skip' run: | diff --git a/Makefile b/Makefile index 0e9d2bbf82c..ab7fab6aa99 100644 --- a/Makefile +++ b/Makefile @@ -164,6 +164,7 @@ lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) # Linting targets lint-ruff: $(LINT_DEP_INSTALL) + $(UV_RUN) python scripts/check_mcp_operation_boundary.py cd litellm && $(UV_RUN) ruff check . && cd .. $(UV_RUN) ruff check --config ruff-tests.toml tests diff --git a/litellm/proxy/_experimental/mcp_server/contracts.py b/litellm/proxy/_experimental/mcp_server/contracts.py new file mode 100644 index 00000000000..c3129d171ad --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/contracts.py @@ -0,0 +1,95 @@ +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import dataclass, field +from datetime import datetime +from types import MappingProxyType +from typing import Final, Protocol + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +def copy_caller(auth: UserAPIKeyAuth | None) -> UserAPIKeyAuth | None: + if auth is None: + return None + span: Final = auth.parent_otel_span + return deepcopy(auth, {id(span): span} if span is not None else None) # mutable-ok: deepcopy mutates its memo + + +@dataclass(frozen=True, slots=True) +class OperationContext: + _caller: UserAPIKeyAuth | None = field(repr=False) + mcp_auth_header: str | None = field(default=None, repr=False) + mcp_servers: tuple[str, ...] | None = None + mcp_server_auth_headers: Mapping[str, Mapping[str, str]] | None = field(default=None, repr=False) + oauth2_headers: Mapping[str, str] | None = field(default=None, repr=False) + raw_headers: Mapping[str, str] | None = field(default=None, repr=False) + client_ip: str | None = None + mcp_proxy_mode: bool = False + + def __post_init__(self) -> None: + object.__setattr__(self, "_caller", copy_caller(self._caller)) + object.__setattr__(self, "mcp_servers", tuple(self.mcp_servers) if self.mcp_servers is not None else None) + object.__setattr__( + self, + "oauth2_headers", + MappingProxyType(dict(self.oauth2_headers)) if self.oauth2_headers is not None else None, + ) + object.__setattr__( + self, "raw_headers", MappingProxyType(dict(self.raw_headers)) if self.raw_headers is not None else None + ) + object.__setattr__( + self, + "mcp_server_auth_headers", + MappingProxyType( + {key: MappingProxyType(dict(value)) for key, value in self.mcp_server_auth_headers.items()} + ) + if self.mcp_server_auth_headers is not None + else None, + ) + + @property + def user_api_key_auth(self) -> UserAPIKeyAuth | None: + return copy_caller(self._caller) + + def legacy_auth( + self, + ) -> tuple[ + UserAPIKeyAuth | None, + str | None, + list[str] | None, # mutable-ok: detached legacy server-list payload + dict[str, dict[str, str]] | None, # mutable-ok: legacy auth dispatch requires concrete dict headers + dict[str, str] | None, # mutable-ok: detached legacy header payload + dict[str, str] | None, # mutable-ok: detached legacy header payload + str | None, + ]: + return ( + self.user_api_key_auth, + self.mcp_auth_header, + list(self.mcp_servers) if self.mcp_servers is not None else None, # mutable-ok: legacy policy list input + { + key: dict(value) for key, value in self.mcp_server_auth_headers.items() + } # mutable-ok: legacy auth dispatch checks concrete dict headers + if self.mcp_server_auth_headers is not None + else None, + dict(self.oauth2_headers) + if self.oauth2_headers is not None + else None, # mutable-ok: legacy OAuth header input + dict(self.raw_headers) if self.raw_headers is not None else None, # mutable-ok: legacy request header input + self.client_ip, + ) + + +class ProgressCallback(Protocol): + async def __call__(self, progress: float, total: float | None, /) -> None: ... + + +@dataclass(frozen=True, slots=True) +class AuthorizedToolCall: + name: str + arguments: Mapping[str, object] + allowed_mcp_servers: tuple[MCPServer, ...] + start_time: datetime + host_progress_callback: ProgressCallback | None + guardrail_context: Mapping[str, object] | None + logging_data: Mapping[str, object] diff --git a/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py b/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py new file mode 100644 index 00000000000..9e321062643 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py @@ -0,0 +1,83 @@ +from collections.abc import Mapping +from typing import Final, Protocol + +from mcp.client.session import ClientRequestContext +from mcp.types import ( + CreateMessageRequestParams, + CreateMessageResult, + CreateMessageResultWithTools, + ElicitRequestParams, + ElicitResult, + ErrorData, +) + +from litellm.proxy._experimental.mcp_server.contracts import OperationContext +from litellm.proxy._types import UserAPIKeyAuth + + +class SamplingCallback(Protocol): + async def __call__( + self, context: ClientRequestContext, params: CreateMessageRequestParams, / + ) -> CreateMessageResult | CreateMessageResultWithTools | ErrorData: ... + + +class ElicitationCallback(Protocol): + async def __call__(self, context: object, params: ElicitRequestParams, /) -> ElicitResult | ErrorData: ... + + +def create_sampling_callback( + user_api_key_auth: UserAPIKeyAuth | None = None, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, + operation_context: OperationContext | None = None, +) -> SamplingCallback: + from litellm.proxy._experimental.mcp_server.server import get_active_auth_context + + auth: Final = get_active_auth_context() if operation_context is None and user_api_key_auth is None else None + captured: Final = ( + operation_context + if operation_context is not None + else OperationContext( + _caller=user_api_key_auth if user_api_key_auth is not None else (auth.user_api_key_auth if auth else None), + raw_headers=raw_headers if raw_headers is not None else (auth.raw_headers if auth else None), + client_ip=client_ip if client_ip is not None else (auth.client_ip if auth else None), + ) + ) + + async def callback( + context: ClientRequestContext, params: CreateMessageRequestParams + ) -> CreateMessageResult | CreateMessageResultWithTools | ErrorData: + import litellm + from litellm.proxy._experimental.mcp_server.sampling_handler import handle_sampling_create_message + + return await handle_sampling_create_message( + context=context, + params=params, + default_model=getattr(litellm, "default_mcp_sampling_model", None), + user_api_key_auth=captured.user_api_key_auth, + raw_headers=dict(captured.raw_headers) + if captured.raw_headers is not None + else None, # mutable-ok: handler consumes an owned request header dict + client_ip=captured.client_ip, + ) + + return callback + + +def create_elicitation_callback() -> ElicitationCallback: + from litellm.proxy._experimental.mcp_server.server import get_active_mcp_session + + downstream_session: Final = get_active_mcp_session() + downstream_capabilities: Final = getattr(downstream_session, "capabilities", None) + + async def callback(context: object, params: ElicitRequestParams) -> ElicitResult | ErrorData: + from litellm.proxy._experimental.mcp_server.elicitation_handler import handle_elicitation_request + + return await handle_elicitation_request( + context=context, + params=params, + downstream_session=downstream_session, + downstream_capabilities=downstream_capabilities, + ) + + return callback diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b49bac8a4cd..4a2713cb19c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -73,6 +73,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPServerAccess, _is_mcp_admitted_user_subject, ) +from litellm.proxy._experimental.mcp_server.contracts import OperationContext from litellm.proxy._experimental.mcp_server.elicitation_handler import ( MCP_ELICITATION_AVAILABLE, ) @@ -195,9 +196,6 @@ from litellm.types.mcp_server.mcp_server_manager import ( from litellm.types.utils import CallTypes if TYPE_CHECKING: - from mcp.client.session import ClientRequestContext - from mcp.types import CreateMessageRequestParams - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.mcp_server.mcp_toolset import MCPToolset @@ -1218,7 +1216,7 @@ async def _resolve_byok_mcp_auth_header( if not mcp_server.is_byok: return mcp_auth_header - from litellm.proxy._experimental.mcp_server.server import ( + from litellm.proxy._experimental.mcp_server.operations import ( _check_byok_credential, _get_byok_credential, ) @@ -1577,77 +1575,25 @@ def _normalize_mcp_server_cost_info(mcp_info: MCPInfo) -> None: mcp_info["mcp_server_cost_info"] = normalized -def _create_sampling_callback(user_api_key_auth: UserAPIKeyAuth | None = None): - """ - Create a sampling callback for MCP ClientSession. - Returns a callable that handles sampling/createMessage requests from - upstream MCP servers by routing them through litellm.acompletion(). - """ +def _create_sampling_callback( + user_api_key_auth: UserAPIKeyAuth | None = None, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, + operation_context: OperationContext | None = None, +): if not MCP_SAMPLING_AVAILABLE: return None + from litellm.proxy._experimental.mcp_server.legacy_callbacks import create_sampling_callback - async def _sampling_callback( - context: "ClientRequestContext", - params: "CreateMessageRequestParams", - ): - import litellm - from litellm.proxy._experimental.mcp_server.sampling_handler import ( - handle_sampling_create_message, - ) - from litellm.proxy._experimental.mcp_server.server import ( - get_active_auth_context, - ) - - auth_context: Final = get_active_auth_context() - resolved_auth: Final = user_api_key_auth or (auth_context.user_api_key_auth if auth_context else None) - # Forward original HTTP headers and client IP so that - # header-dependent guardrails, tag-based routing, trace - # correlation, and forward_llm_provider_auth_headers work - # correctly for sampling sub-calls. - _raw_headers: Final = getattr(auth_context, "raw_headers", None) - _client_ip: Final = getattr(auth_context, "client_ip", None) - - return await handle_sampling_create_message( - context=context, - params=params, - default_model=getattr(litellm, "default_mcp_sampling_model", None), - user_api_key_auth=resolved_auth, - raw_headers=_raw_headers, - client_ip=_client_ip, - ) - - return _sampling_callback + return create_sampling_callback(user_api_key_auth, raw_headers, client_ip, operation_context) def _create_elicitation_callback(): - """ - Create an elicitation callback for MCP ClientSession. - Returns a callable that handles elicitation/create requests from - upstream MCP servers. In gateway mode, this relays to the downstream - client; in tool bridge mode, it returns a decline response. - """ if not MCP_ELICITATION_AVAILABLE: return None + from litellm.proxy._experimental.mcp_server.legacy_callbacks import create_elicitation_callback - async def _elicitation_callback(context, params): - from litellm.proxy._experimental.mcp_server.elicitation_handler import ( - handle_elicitation_request, - ) - from litellm.proxy._experimental.mcp_server.server import get_active_mcp_session - - # In Gateway mode, we relay the elicitation request to the downstream client - # that triggered the current operation. - downstream_session: Final = get_active_mcp_session() - downstream_capabilities = getattr(downstream_session, "capabilities", None) if downstream_session else None - - return await handle_elicitation_request( - context=context, - params=params, - downstream_session=downstream_session, - downstream_capabilities=downstream_capabilities, - ) - - return _elicitation_callback + return create_elicitation_callback() def _record_mcp_guardrail_evaluations( @@ -3386,17 +3332,13 @@ class MCPServerManager: listable but uninvokable. Empty inside a toolset scope: toolset_mcp_route / dynamic_mcp_route set - ``_mcp_active_toolset_id`` before calling the handler, pinning the request to the toolset's + the caller's server-only ``mcp_toolset_id`` before calling the handler, pinning the request to the toolset's own servers (checking op.mcp_toolsets==[] instead would false-positive on DB-default rows where Postgres initialises the column to ARRAY[]::TEXT[]). ``allow_all_server_ids`` / ``submitted_server_ids`` are injectable so the server union, which precomputes both for its fallback path, does not compute them twice.""" - from litellm.proxy._experimental.mcp_server.mcp_context import ( # noqa: PLC0415 - _mcp_active_toolset_id, - ) - - if _mcp_active_toolset_id.get() is not None: + if user_api_key_auth is not None and user_api_key_auth.mcp_toolset_id is not None: return set() if allow_all_server_ids is None: allow_all_server_ids = self.get_allow_all_keys_server_ids() @@ -4164,6 +4106,8 @@ class MCPServerManager: subject_token: str | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, cred_provider: UpstreamCredentialProvider | None = None, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, ) -> MCPClient: """ Create an MCPClient instance for the given server. @@ -4212,7 +4156,13 @@ class MCPServerManager: # Create sampling and elicitation callbacks for this client sampling_cb = ( - _create_sampling_callback(user_api_key_auth=user_api_key_auth) if resolved_server.allow_sampling else None + _create_sampling_callback( + operation_context=OperationContext( + _caller=user_api_key_auth, raw_headers=raw_headers, client_ip=client_ip + ) + ) + if resolved_server.allow_sampling + else None ) elicitation_cb: Final = _create_elicitation_callback() if resolved_server.allow_elicitation else None @@ -4357,6 +4307,7 @@ class MCPServerManager: raw_headers: dict[str, str] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, oauth2_headers: dict[str, str] | None = None, + client_ip: str | None = None, ) -> list[MCPTool]: """ Helper method to get tools from a single MCP server with prefixed names. @@ -4446,6 +4397,8 @@ class MCPServerManager: stdio_env=stdio_env, subject_token=subject_token, user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, ) ## HANDLE OPENAPI TOOLS @@ -4556,6 +4509,7 @@ class MCPServerManager: extra_headers: dict[str, str] | None = None, add_prefix: bool = True, raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, ) -> list[Prompt]: try: headers: Final = ( @@ -4576,6 +4530,8 @@ class MCPServerManager: stdio_env=stdio_env, subject_token=subject_token, user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, ) credential_fingerprint: Final = await client.discovery_auth_fingerprint() key: Final = self._discovery_key( @@ -4599,6 +4555,7 @@ class MCPServerManager: extra_headers: dict[str, str] | None = None, add_prefix: bool = True, raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, ) -> list[Resource]: try: headers: Final = ( @@ -4619,6 +4576,8 @@ class MCPServerManager: stdio_env=stdio_env, subject_token=subject_token, user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, ) credential_fingerprint: Final = await client.discovery_auth_fingerprint() key: Final = self._discovery_key( @@ -4642,6 +4601,7 @@ class MCPServerManager: extra_headers: dict[str, str] | None = None, add_prefix: bool = True, raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, ) -> list[ResourceTemplate]: try: headers: Final = ( @@ -4662,6 +4622,8 @@ class MCPServerManager: stdio_env=stdio_env, subject_token=subject_token, user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, ) credential_fingerprint: Final = await client.discovery_auth_fingerprint() key: Final = self._discovery_key( @@ -4685,6 +4647,7 @@ class MCPServerManager: mcp_auth_header: str | dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, ) -> ReadResourceResult: """Read resource contents from a specific MCP server.""" @@ -4705,6 +4668,9 @@ class MCPServerManager: extra_headers=extra_headers, stdio_env=stdio_env, subject_token=subject_token, + raw_headers=raw_headers, + client_ip=client_ip, + user_api_key_auth=user_api_key_auth, ) return await client.read_resource(url) @@ -4718,6 +4684,7 @@ class MCPServerManager: mcp_auth_header: str | dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, ) -> GetPromptResult: """Fetch a specific prompt definition from a single MCP server.""" @@ -4738,6 +4705,9 @@ class MCPServerManager: extra_headers=extra_headers, stdio_env=stdio_env, subject_token=subject_token, + raw_headers=raw_headers, + client_ip=client_ip, + user_api_key_auth=user_api_key_auth, ) get_prompt_request_params: Final = GetPromptRequestParams( @@ -5818,6 +5788,8 @@ class MCPServerManager: stdio_env: dict[str, str] | None, subject_token: str | None, user_api_key_auth: UserAPIKeyAuth | None, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, ) -> CallToolResult: """Call a token_exchange (OBO) tool; on an upstream 401/403 re-mint the token once and retry. @@ -5843,6 +5815,8 @@ class MCPServerManager: stdio_env=stdio_env, subject_token=subject_token, user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, ) return await retry_client.call_tool(call_tool_params, host_progress_callback=host_progress_callback) @@ -5860,6 +5834,7 @@ class MCPServerManager: host_progress_callback: Callable | None = None, hook_extra_headers: dict[str, str] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, + client_ip: str | None = None, ) -> CallToolResult: """ Call a regular MCP tool using the MCP client. @@ -6004,6 +5979,8 @@ class MCPServerManager: stdio_env=stdio_env, subject_token=subject_token, user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, ) call_tool_params: Final = MCPCallToolRequestParams( @@ -6027,6 +6004,8 @@ class MCPServerManager: stdio_env=stdio_env, subject_token=subject_token, user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, ) tool_call_coro = _obo_call_tool_limited() @@ -6202,7 +6181,7 @@ class MCPServerManager: return oauth2_headers try: - from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415 + from litellm.proxy._experimental.mcp_server.operations import ( # noqa: PLC0415 _get_user_oauth_extra_headers_from_db, ) @@ -6308,6 +6287,7 @@ class MCPServerManager: host_progress_callback: Callable | None = None, litellm_logging_obj: "LiteLLMLoggingObj | None" = None, guardrail_context: Mapping[str, object] | None = None, + client_ip: str | None = None, ) -> CallToolResult: """ Call a tool with the given name and arguments @@ -6434,6 +6414,7 @@ class MCPServerManager: mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, + client_ip=client_ip, proxy_logging_obj=proxy_logging_obj, host_progress_callback=host_progress_callback, hook_extra_headers=hook_result.get("extra_headers"), diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py new file mode 100644 index 00000000000..fcee3483e15 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -0,0 +1,3102 @@ +"""Shared MCP operation policy and dispatch.""" + +import asyncio +import traceback +import types +import uuid +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import Any, Final, NoReturn, TypeAlias, overload + +from fastapi import HTTPException +from mcp import ReadResourceResult, Resource +from mcp.types import ( + CallToolRequest, + CallToolRequestParams, + CallToolResult, + GetPromptRequest, + GetPromptRequestParams, + GetPromptResult, + ListPromptsRequest, + ListPromptsResult, + ListResourcesRequest, + ListResourcesResult, + ListResourceTemplatesRequest, + ListResourceTemplatesResult, + ListToolsRequest, + ListToolsResult, + PaginatedRequestParams, + Prompt, + ReadResourceRequest, + ReadResourceRequestParams, + ResourceTemplate, + TextContent, +) +from mcp.types import Tool as MCPTool +from pydantic import AnyUrl, ConfigDict, Field, TypeAdapter +from typing_extensions import ReadOnly, TypedDict, assert_never + +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.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, +) +from litellm.proxy._experimental.mcp_server.byok_credential_cache import ( + byok_credential_cache, + byok_credential_cache_key, + cache_byok_credential, + get_cached_byok_credential, +) +from litellm.proxy._experimental.mcp_server.contracts import ( + AuthorizedToolCall, + OperationContext, + ProgressCallback, +) +from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload +from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPToolResultError, + MCPUpstreamAuthError, +) +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + SERVER_OUTCOMES_META_KEY, + AggregateToolListing, + ServerListOk, + ServerOutcome, + classify_list_exception, + outcome_wire_value, +) +from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + _caller_authorization_fans_out, + _client_forwarded_authorization_headers, + _resolve_openapi_tool_auth, + _should_strip_caller_authorization, + global_mcp_server_manager, +) +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + _redact_mcp_resource_url, + get_byok_www_authenticate, +) +from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _request_auth_header, + _request_extra_headers, + _request_resolved_auth_headers, +) +from litellm.proxy._experimental.mcp_server.tool_registry import ( + global_mcp_tool_registry, +) +from litellm.proxy._experimental.mcp_server.utils import ( + MCP_TOOL_PREFIX_SEPARATOR, + MCPMissingUserEnvVarsError, + add_server_prefix_to_name, + build_synthetic_mcp_request, + extract_mcp_tool_result_error_message, + get_server_prefix, + is_tool_name_prefixed, + iter_known_server_prefixes, + logging_safe_mcp_headers, + match_known_tool_name, + normalize_server_name, + split_server_prefix_from_name, + strip_known_server_prefix, +) +from litellm.proxy._types import ( + UserAPIKeyAuth, +) +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + publish_auth_cache_invalidation, +) +from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + get_chain_id_from_headers, +) +from litellm.types.mcp import ( + DEFAULT_CREDENTIAL_HEADER, + MCPAuth, + without_header, +) +from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer +from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall +from litellm.utils import Rules, client, function_setup + +__all__ = ( + "_MCP_CREDENTIAL_REQUEST_FIELDS", + "ListMCPToolsRestAPIResponseObject", + "MCPInfo", + "MCPServer", + "_McpDeniedDetail", + "_aggregate_server_key", + "_build_virtual_call_logging_obj", + "_check_byok_credential", + "_client_has_passthrough_authorization", + "_client_has_per_server_auth_header", + "_dispatch_virtual_mcp_tool", + "_fire_mcp_tool_call_logging", + "_get_allowed_mcp_servers", + "_get_allowed_mcp_servers_from_mcp_server_names", + "_get_byok_credential", + "_get_prompts_from_mcp_servers", + "_get_resource_templates_from_mcp_servers", + "_get_resources_from_mcp_servers", + "_get_standard_logging_mcp_tool_call", + "_get_tools_from_mcp_servers", + "_get_user_oauth_extra_headers_from_db", + "_handle_local_mcp_tool", + "_handle_managed_mcp_tool", + "_http_detail_message", + "_invalidate_byok_cred_cache", + "_list_mcp_prompts", + "_list_mcp_resource_templates", + "_list_mcp_resources", + "_list_mcp_tools", + "_list_tools_before_first_call", + "_mcp_session_id_from_headers", + "_merge_gateway_initialize_instructions", + "_prefetch_oauth_creds_for_user", + "_prepare_mcp_server_headers", + "_raise_if_initialize_grants_no_mcp_servers", + "_resolve_display_name_to_original", + "_run_post_mcp_call_guardrails", + "_server_answers_to", + "_tool_name_matches", + "apply_tool_overrides", + "call_mcp_tool", + "execute_mcp_tool", + "filter_tools_by_allowed_tools", + "filter_tools_by_key_team_permissions", + "fire_mcp_tool_call_failure_logging", + "mcp_get_prompt", + "mcp_read_resource", + "raise_denied_scoped_mcp_access", +) + + +async def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: + """Drop a stored-or-deleted BYOK credential from this worker's cache and from every peer worker's.""" + cache_key: Final = byok_credential_cache_key(user_id, server_id) + byok_credential_cache.delete_cache(cache_key) + await publish_auth_cache_invalidation(cache_key=cache_key) + + +def _mcp_session_id_from_headers( + raw_headers: dict[str, str] | None, +) -> str | None: + """The ``mcp-session-id`` of a stateful MCP session, read case-insensitively + from the request headers. ``None`` for stateless calls (no such header).""" + if not raw_headers: + return None + for key, value in raw_headers.items(): + if isinstance(key, str) and key.lower() == "mcp-session-id": + return value or None + return None + + +class ListMCPToolsRestAPIResponseObject(MCPTool): + """ + Object returned by the /tools/list REST API route. + """ + + mcp_info: MCPInfo | None = Field(default=None, alias="mcp_info") + model_config = ConfigDict(arbitrary_types_allowed=True) + + +async def _build_virtual_call_logging_obj( + name: str, + arguments: dict[str, object], + user_api_key_auth: UserAPIKeyAuth, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, +) -> LiteLLMLoggingObj | None: + """Run the pre-call pipeline (guardrails + logging setup) for a virtual + mcp_tool_call so the SSE path spend-logs like the REST path.""" + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + from litellm.proxy.proxy_server import ( + general_settings, + proxy_config, + proxy_logging_obj, + ) + + request: Final = build_synthetic_mcp_request( + path="/mcp/tools/call", + raw_headers=raw_headers, + client_ip=client_ip, + ) + _, virtual_logging_obj = await ProxyBaseLLMRequestProcessing( + data={"name": name, "arguments": arguments} + ).common_processing_pre_call_logic( + request=request, + user_api_key_dict=user_api_key_auth, + proxy_config=proxy_config, + route_type=CallTypes.call_mcp_tool.value, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + ) + return virtual_logging_obj + + +async def _dispatch_virtual_mcp_tool( + name: str, + arguments: dict[str, object] | None, + user_api_key_auth: UserAPIKeyAuth | None, + client_ip: str | None, + mcp_servers: list[str] | None = 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, + mcp_proxy_mode: bool = False, +) -> CallToolResult | None: + """Handle the mcp_tool_search / mcp_tool_call virtual tools. + + Returns a CallToolResult when ``name`` is a virtual tool, else ``None`` so + the caller falls through to normal tool routing. + """ + from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K + from litellm.proxy._experimental.mcp_server.tool_search import ( + AGENT_SEARCH_TOOL_NAME, + DEFAULT_AGENT_SEARCH_TOP_K, + MCP_PROXY_CALL_TOOL_NAME, + MCP_PROXY_TOOL_NAMES, + MCP_TOOL_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, + VIRTUAL_TOOL_NAMES, + coerce_top_k, + handle_agent_search, + handle_mcp_proxy_tool, + handle_mcp_tool_call, + handle_mcp_tool_search, + handle_skill_search, + ) + + if mcp_proxy_mode and name not in MCP_PROXY_TOOL_NAMES: + return CallToolResult( + content=[ # mutable-ok: MCP result content + TextContent(type="text", text=f"Tool {name} is unavailable on /mcp/proxy") + ], + is_error=True, + ) + + if mcp_proxy_mode and name in MCP_PROXY_TOOL_NAMES: + assert user_api_key_auth is not None + proxy_call_start: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes + proxy_logging_obj: Final = ( + await _build_virtual_call_logging_obj( + name=name, + arguments=arguments or {}, # mutable-ok: logging pipeline payload + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, + ) + if name == MCP_PROXY_CALL_TOOL_NAME + else None + ) + try: + proxy_result: Final = await handle_mcp_proxy_tool( + name=name, + arguments=arguments or {}, # mutable-ok: proxy handler payload + user_api_key_dict=user_api_key_auth, + client_ip=client_ip, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + litellm_logging_obj=proxy_logging_obj, + ) + except Exception as exc: + if proxy_logging_obj is not None: + from litellm.proxy.proxy_server import proxy_logging_obj as request_logging_obj + + failure_end: Final = datetime.now() # noqa: DTZ005 # matches the logging pipeline start time + failure_traceback: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + try: + proxy_logging_obj.failure_handler(exc, failure_traceback, proxy_call_start, failure_end) + await proxy_logging_obj.async_failure_handler(exc, failure_traceback, proxy_call_start, failure_end) + if not isinstance(exc, MCPUpstreamAuthError): + await request_logging_obj.post_call_failure_hook( + request_data={ # mutable-ok: failure hook mutates its request payload + "name": name, + "arguments": arguments, + "litellm_logging_obj": proxy_logging_obj, + }, + original_exception=exc, + user_api_key_dict=user_api_key_auth, + route="/mcp/call_tool", + traceback_str=failure_traceback, + ) + except Exception: # noqa: BLE001 # a failing failure hook must not mask the tool call's own error + verbose_logger.exception("Error logging failed MCP proxy tool call") + raise + if proxy_logging_obj is not None: + return await _fire_mcp_tool_call_logging( + logging_obj=proxy_logging_obj, + result=proxy_result, + start_time=proxy_call_start, + end_time=datetime.now(), # noqa: DTZ005 # matches the logging pipeline start time + user_api_key_auth=user_api_key_auth, + request_data=types.MappingProxyType({"name": name, "arguments": arguments}), + ) + return proxy_result + + if name not in VIRTUAL_TOOL_NAMES: + return None + + if not getattr( + getattr(user_api_key_auth, "object_permission", None), + "mcp_tool_search_enabled", + False, + ): + return CallToolResult( + content=[ + TextContent( + type="text", + text=f"Tool {name} requires mcp_tool_search_enabled on the key", + ) + ], + is_error=True, + ) + + args: Final = arguments or {} + if name == MCP_TOOL_SEARCH_TOOL_NAME: + return await handle_mcp_tool_search( + query=TypeAdapter(str).validate_python(args.get("query", "")), + top_k=coerce_top_k(args.get("top_k", 5)), + user_api_key_dict=user_api_key_auth, + client_ip=client_ip, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + ) + + assert user_api_key_auth is not None # guaranteed by the flag check above + if name == AGENT_SEARCH_TOOL_NAME: + return await handle_agent_search( + query=str(args.get("query", "")), + top_k=coerce_top_k(args.get("top_k", DEFAULT_AGENT_SEARCH_TOP_K), default=DEFAULT_AGENT_SEARCH_TOP_K), + user_api_key_dict=user_api_key_auth, + ) + if name == SKILL_SEARCH_TOOL_NAME: + return await handle_skill_search( + query=str(args.get("query", "")), + top_k=coerce_top_k(args.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K), + user_api_key_dict=user_api_key_auth, + ) + virtual_logging_obj: Final = await _build_virtual_call_logging_obj( + name=name, + arguments=args, + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, + ) + tool_request: Final = CallToolRequestParams.model_validate( + types.MappingProxyType({"name": args.get("tool_name", ""), "arguments": args.get("arguments") or {}}) + ) + return await handle_mcp_tool_call( + tool_name=tool_request.name, + arguments=tool_request.arguments or {}, + user_api_key_dict=user_api_key_auth, + client_ip=client_ip, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + litellm_logging_obj=virtual_logging_obj, + ) + + +async def _get_allowed_mcp_servers_from_mcp_server_names( + mcp_servers: Sequence[str] | None, + allowed_mcp_servers: list[MCPServer], +) -> list[MCPServer]: + """ + Get the filtered MCP servers from the MCP server names. + + Fails closed when ``mcp_servers`` is explicitly provided (path- or + header-derived) but none of the names resolve to a server alias or + access group the caller can access. The previous behavior returned + the full ``allowed_mcp_servers`` set, which silently widened scope + when a client targeted ``/mcp//`` and made URL/header + namespacing appear to work when it did not. + """ + + filtered_server: Final[dict[str, MCPServer]] = {} + # Filter servers based on mcp_servers parameter if provided + if mcp_servers is not None: + for server_or_group in mcp_servers: + server_name_matched = False + + for server in allowed_mcp_servers: + if server and _server_answers_to(server, server_or_group): + filtered_server[server.server_id] = server + server_name_matched = True + break + + if not server_name_matched: + try: + access_group_server_ids = await MCPRequestHandler._get_mcp_servers_from_access_groups( + [server_or_group] + ) + # Only include servers that the user has access to + for server_id in access_group_server_ids: + for server in allowed_mcp_servers: + if server_id == server.server_id: + filtered_server[server.server_id] = server + except Exception as e: + verbose_logger.debug("Could not resolve '%s' as access group: %s", server_or_group, e) + + if filtered_server: + return list(filtered_server.values()) + + if mcp_servers is not None: + # Caller asked for a specific scope but nothing resolved. Fail + # closed so URL/header namespacing cannot silently fall back to + # the caller's full allowed-server set. + verbose_logger.debug( + "MCP scope filter resolved to no servers for requested names %s; returning empty list (fail-closed).", + mcp_servers, + ) + return [] + + return allowed_mcp_servers + + +def _http_detail_message(detail: object) -> str: + return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail) + + +def _server_answers_to(server: MCPServer, name: str) -> bool: + requested: Final = name.lower() + return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known) + + +async def raise_denied_scoped_mcp_access( + requested_names: Sequence[str], + user_api_key_auth: UserAPIKeyAuth | None, + client_ip: str | None = None, +) -> None: + """A scoped request (``/mcp/`` path or ``x-mcp-servers`` header) resolved to zero + allowed servers, so the denial must be loud: a silent 200 with no tools reads as a healthy + server with no tools. Unknown, unauthorized, and access-group names all share one generic + error so scoping cannot probe which servers exist; the agent variant fires only when the + same request resolves once the agent binding is stripped, proving the binding caused the veto.""" + agent_id: Final = user_api_key_auth.agent_id if user_api_key_auth else None + if user_api_key_auth is not None and agent_id: + resolved_without_agent: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None})), + mcp_servers=requested_names, + client_ip=client_ip, + ) + + def _resolved_to_server(name: str) -> bool: + return any(_server_answers_to(server, name) for server in resolved_without_agent) + + vetoed_server: Final = next((name for name in requested_names if _resolved_to_server(name)), None) + if vetoed_server is not None: + agent_denial: Final[_McpDeniedDetail] = { + "error": ( + f"MCP server '{vetoed_server}' is not available to this key: the key is bound to " + f"agent '{agent_id}', whose MCP grants do not include this server. Add the server " + f"to the agent's object_permission.mcp_servers (edit the agent in the Admin UI or " + f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent." + ) + } + raise HTTPException(status_code=403, detail=agent_denial) + vetoed_group: Final = next( + ( + name + for name in requested_names + if not _resolved_to_server(name) + and any(name in (server.access_groups or ()) for server in resolved_without_agent) + ), + None, + ) + if vetoed_group is not None: + group_denial: Final[_McpDeniedDetail] = { + "error": ( + f"MCP access group '{vetoed_group}' is not available to this key: the key is bound to " + f"agent '{agent_id}', whose MCP grants do not include it. Add the group to the " + f"agent's object_permission.mcp_access_groups (edit the agent in the Admin UI or " + f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent." + ) + } + raise HTTPException(status_code=403, detail=group_denial) + generic_denial: Final[_McpDeniedDetail] = { + "error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}" + } + raise HTTPException(status_code=403, detail=generic_denial) + + +def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool: + """ + Check if a tool name matches any name in the filter list. + + Reads the same owner the server-level permission checks use, so discovery hides + exactly what dispatch refuses. ``mcp_server`` is required: guessing the boundary + at the first separator mismatches every tool on a server whose prefix contains + the separator. + """ + bare_name: Final = strip_known_server_prefix(tool_name, mcp_server) + return match_known_tool_name(bare_name, mcp_server, filter_list) is not None + + +def filter_tools_by_allowed_tools( + tools: list[MCPTool], + mcp_server: MCPServer, +) -> list[MCPTool]: + """ + Filter tools by allowed/disallowed tools configuration. + + If allowed_tools is set, only tools in that list are returned. + If disallowed_tools is set, tools in that list are excluded. + Tool names are matched with and without server prefixes for flexibility. + + Args: + tools: List of tools to filter + mcp_server: Server configuration with allowed_tools/disallowed_tools + + Returns: + Filtered list of tools + """ + from litellm.proxy._experimental.mcp_server.utils import ( + server_applies_tool_allowlist, + ) + + tools_to_return = tools + + # Filter by allowed_tools (whitelist) + if server_applies_tool_allowlist(mcp_server): + if not mcp_server.allowed_tools: + return [] + tools_to_return = [ + tool for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools, mcp_server) + ] + + # Filter by disallowed_tools (blacklist) + if mcp_server.disallowed_tools: + tools_to_return = [ + tool + for tool in tools_to_return + if not _tool_name_matches(tool.name, mcp_server.disallowed_tools, mcp_server) + ] + + return tools_to_return + + +def apply_tool_overrides( + tools: list[MCPTool], + mcp_server: MCPServer, +) -> list[MCPTool]: + """Apply admin-configured display name/description overrides to tools. + + Overrides are keyed by the unprefixed tool name, same convention as + allowed_tools configuration. + """ + display_name_map: Final = mcp_server.tool_name_to_display_name or {} + description_map: Final = mcp_server.tool_name_to_description or {} + if not display_name_map and not description_map: + return tools + + for tool in tools: + unprefixed = strip_known_server_prefix(tool.name, mcp_server) + lookup_key = unprefixed or tool.name + if lookup_key in display_name_map: + tool.name = display_name_map[lookup_key] + if lookup_key in description_map: + tool.description = description_map[lookup_key] + return tools + + +async def _get_allowed_mcp_servers( + user_api_key_auth: UserAPIKeyAuth | None, + mcp_servers: Sequence[str] | None, + client_ip: str | None = None, +) -> list[MCPServer]: + """Return allowed MCP servers for a request after applying filters. + + Args: + user_api_key_auth: The authenticated user's API key info. + mcp_servers: Optional list of server names to filter to. + client_ip: Client IP for IP-based access control. If None, falls back to + auth context. Pass explicitly from request handlers for safety. + Note: If client_ip is None and auth context is not set, IP filtering is skipped. + This is intentional for internal callers but may indicate a bug if called + from a request handler without proper context setup. + """ + allowed_mcp_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) + ( + allowed_mcp_server_ids, + _ip_blocked, + ) = global_mcp_server_manager.filter_server_ids_by_ip_with_info(allowed_mcp_server_ids, client_ip) + verbose_logger.debug( + "MCP IP filter: client_ip=%s, allowed_server_ids=%s", + client_ip, + allowed_mcp_server_ids, + ) + if _ip_blocked > 0: + verbose_logger.debug( + "MCP IP filtering: %d server(s) are not accessible from client IP %s " + "because they are restricted to internal networks. " + "No tools from those servers will be returned. " + "To expose a server externally, set 'available_on_public_internet: true' " + "in its configuration.", + _ip_blocked, + client_ip, + ) + allowed_mcp_servers: list[MCPServer] = [] + for allowed_mcp_server_id in allowed_mcp_server_ids: + mcp_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id) + if mcp_server is not None: + # Apply the request-time oauth2_flow backstop for legacy null rows. + mcp_server = MCPServerManager.resolve_oauth2_flow_for_request(mcp_server) + allowed_mcp_servers.append(mcp_server) + + if mcp_servers is not None: + allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( + mcp_servers=mcp_servers, + allowed_mcp_servers=allowed_mcp_servers, + ) + + return allowed_mcp_servers + + +def _client_has_per_server_auth_header( + server: MCPServer, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, +) -> bool: + """True if the request carries a per-server ``x-mcp-{alias}-authorization`` + header for this server. This is the multi-server binding: it names one + upstream, so it is unambiguously the caller's upstream token regardless of + auth mode (never the LiteLLM admission credential). + + Resolves through the same ``lookup_mcp_server_auth_in_headers`` egress uses, so + the connect gate and egress agree on which per-server header names match: a + dashboard client sends ``x-mcp-{sanitize_mcp_alias_for_header(alias)}-authorization``, + and matching only the raw alias here would 401 a token egress would forward. + """ + if not mcp_server_auth_headers: + return False + from litellm.proxy._experimental.mcp_server.utils import ( + lookup_mcp_server_auth_in_headers, + ) + + server_headers: Final = lookup_mcp_server_auth_in_headers( + mcp_server_auth_headers, + alias=server.alias, + server_name=server.server_name, + access_groups=server.access_groups, + ) + if isinstance(server_headers, str): + return bool(server_headers.strip()) + if isinstance(server_headers, dict): + return any(isinstance(hk, str) and hk.lower() == "authorization" for hk in server_headers) + return False + + +def _client_has_passthrough_authorization( + server: MCPServer, + oauth2_headers: dict[str, str] | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, +) -> bool: + """True if the incoming request already carries an ``Authorization`` + header the gateway will forward to this pass-through server. + + The client may supply the bearer as either the top-level + ``Authorization`` header (surfaced via ``oauth2_headers``) or a + per-server ``x-mcp-auth-`` style header (surfaced via + ``mcp_server_auth_headers``). Either form skips the pre-emptive 401. + """ + if oauth2_headers: + for k in oauth2_headers: + if k.lower() == "authorization": + return True + return _client_has_per_server_auth_header(server, mcp_server_auth_headers) + + +async def _get_user_oauth_extra_headers_from_db( + server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, + prefetched_creds: 'Mapping[str, "OAuthCredentialPayload"] | None' = None, +) -> dict[str, str] | None: + """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None. + + Thin wrapper over ``resolve_user_oauth_access_token`` (Redis cache, else DB + refresh); + ``prefetched_creds`` skips the per-server Redis/DB lookups for the batch path. + """ + if server.auth_type != MCPAuth.oauth2 or user_api_key_auth is None: + return None + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + resolve_user_oauth_access_token, + ) + + token: Final = await resolve_user_oauth_access_token( + getattr(user_api_key_auth, "user_id", None), server, prefetched_creds + ) + return {"Authorization": f"Bearer {token}"} if token else None + + +async def _prefetch_oauth_creds_for_user( + user_api_key_auth: UserAPIKeyAuth | None, +) -> dict[str, "OAuthCredentialPayload"]: + """Fetch all OAuth2 credentials for the user in one DB query. + + Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. + """ + user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None + if not user_id: + return {} + try: + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + list_user_oauth_credentials, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 + + prisma_client: Final = get_prisma_client_or_throw( + "Database not connected. Connect a database to use OAuth2 MCP tools." + ) + creds: Final = await list_user_oauth_credentials(prisma_client, user_id) + return {c["server_id"]: c for c in creds if "server_id" in c} + except Exception: + verbose_logger.warning("_prefetch_oauth_creds_for_user: failed to prefetch OAuth credentials") + return {} + + +def _prepare_mcp_server_headers( + server: MCPServer, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + mcp_auth_header: str | None, + oauth2_headers: dict[str, str] | None, + raw_headers: dict[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None = None, + scope_servers: list[MCPServer] | None = None, +) -> tuple[dict[str, str] | str | None, dict[str, str] | None]: + """Build auth and extra headers for a server. + + ``scope_servers`` is the full server list a fan-out handler iterates. Passing it lets the + client-forwarded token modes withhold the caller's request-wide ``Authorization`` when + another server in the scope would also receive it (``_caller_authorization_fans_out``); + explicitly-addressed operations leave it None. Per-server ``x-mcp-{alias}-authorization`` + headers are unaffected — they bind one token to one server and are the multi-server shape. + """ + server_auth_header: dict[str, str] | str | None = None + if mcp_server_auth_headers: + from litellm.proxy._experimental.mcp_server.utils import ( + lookup_mcp_server_auth_in_headers, + ) + + server_auth_header = lookup_mcp_server_auth_in_headers( + mcp_server_auth_headers, + alias=server.alias, + server_name=server.server_name, + access_groups=server.access_groups, + ) + + extra_headers: dict[str, str] | None = None + is_client_forwarded_mode: Final = server.is_client_forwarded_token + # In a multi-server listing scope the request-wide Authorization can only carry one token, + # so it is withheld from a client-forwarded server when another server in scope also consumes + # it (RFC 9700 cross-resource replay); such scopes must bind per-server via + # x-mcp-{alias}-authorization. The decision is computed once so BOTH the forwarding branch and + # the extra_headers copy loop below honor it — otherwise a server that lists Authorization in + # extra_headers would re-copy the withheld bearer from raw_headers and replay it anyway. + withhold_forwarded_authorization: Final = is_client_forwarded_mode and _caller_authorization_fans_out( + server, scope_servers + ) + if server.auth_type == MCPAuth.oauth2: + # For OAuth2 M2M servers, upstream Authorization must come from + # client_credentials token fetch, never from caller headers. + if server.has_client_credentials: + extra_headers = None + else: + # Copy to avoid mutating the original dict (important for parallel fetching) + extra_headers = oauth2_headers.copy() if oauth2_headers else None + # Migrated authorization_code: the v2 resolver injects the stored per-user + # token, so drop the caller-forwarded Authorization (apply-if-absent would + # otherwise let it shadow the resolved token). Delegate keeps it. Centralized + # via _should_strip_caller_authorization to match _call_regular_mcp_tool. + if extra_headers and _should_strip_caller_authorization( + mcp_server=server, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + ): + extra_headers = without_header(extra_headers, DEFAULT_CREDENTIAL_HEADER) + elif is_client_forwarded_mode: + if not withhold_forwarded_authorization: + extra_headers = _client_forwarded_authorization_headers( + mcp_server=server, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + ) + + if server.extra_headers and raw_headers: + if extra_headers is None: + extra_headers = {} + + normalized_raw_headers: Final = {str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str)} + + # Centralized strip decision shared with + # ``MCPServerManager._call_regular_mcp_tool`` so the two + # code paths cannot drift on this security-sensitive choice. + # See ``_should_strip_caller_authorization`` for the rules. + strip_caller_authorization: Final = _should_strip_caller_authorization( + mcp_server=server, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + ) + + for header in server.extra_headers: + if not isinstance(header, str): + continue + if header.lower() == "authorization" and (strip_caller_authorization or withhold_forwarded_authorization): + continue + header_value = normalized_raw_headers.get(header.lower()) + if header_value is None: + continue + extra_headers[header] = header_value + + # Reset to None if no headers were actually added + if extra_headers is not None and len(extra_headers) == 0: + extra_headers = None + + if server_auth_header is None: + server_auth_header = mcp_auth_header + + return server_auth_header, extra_headers + + +def _merge_gateway_initialize_instructions( + allowed_mcp_servers: list[MCPServer], +) -> str | None: + """YAML/DB override, else upstream text (prefetch on init, or list_tools / health_check / call_tool cache).""" + if not allowed_mcp_servers: + return None + + texts: Final[list[tuple[str, str]]] = [] + for server in allowed_mcp_servers: + label = server.alias or server.server_name or server.name or server.server_id or "mcp" + if server.instructions and server.instructions.strip(): + texts.append((label, server.instructions.strip())) + continue + if server.spec_path: + continue + cached = global_mcp_server_manager._upstream_initialize_instructions_by_server_id.get(server.server_id) + if cached and cached.strip(): + texts.append((label, cached.strip())) + + if not texts: + return None + if len(texts) == 1: + return texts[0][1] + return "\n\n---\n\n".join(f"[{lbl}]\n{txt}" for lbl, txt in texts) + + +async def _raise_if_initialize_grants_no_mcp_servers( + allowed: Sequence[MCPServer], + user_api_key_auth: UserAPIKeyAuth | None, + mcp_servers: Sequence[str] | None, + client_ip: str | None, +) -> None: + if allowed or user_api_key_auth is None or not user_api_key_auth.api_key: + return + if mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) + no_servers_denial: Final[_McpDeniedDetail] = { + "error": ( + "The key has no MCP servers granted, or none of its granted servers is loaded and allowed for " + "this client IP. Grant servers or access groups to the key, its team, or its organization " + "(object_permission.mcp_servers), check the server's allowed IPs, and reconnect." + ) + } + raise HTTPException(status_code=403, detail=no_servers_denial) + + +def _aggregate_server_key(server: MCPServer) -> str: + """The client-visible key for a server in listing outcomes and spend metadata: the same + display prefix (alias, or the short prefix when that mode is enabled) the caller already + sees on the tool names. Canonical internal server names never key a caller-readable + surface; when the display naming deliberately hides them, the outcome keys must too.""" + return get_server_prefix(server) or "unknown" + + +async def _get_tools_from_mcp_servers( + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None, + mcp_servers: list[str] | 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, + log_list_tools_to_spendlogs: bool = False, + list_tools_log_source: str | None = None, + litellm_trace_id: str | None = None, + request_tags: list[str] | None = None, + client_ip: str | None = None, + mcp_proxy_mode: bool = False, +) -> AggregateToolListing: + """ + Helper method to fetch tools from MCP servers based on server filtering criteria. + + Args: + user_api_key_auth: User authentication info for access control + mcp_auth_header: Optional auth header for MCP server (deprecated) + mcp_servers: Optional list of server names/aliases to filter by + mcp_server_auth_headers: Optional dict of server-specific auth headers + oauth2_headers: Optional dict of oauth2 headers + + Returns: + AggregateToolListing: Combined tools from filtered servers plus each server's + classified listing outcome + """ + + list_tools_start_time: Final = datetime.now() + litellm_logging_obj: LiteLLMLoggingObj | None = None + list_tools_request_data: dict[str, object] = {} + + if log_list_tools_to_spendlogs: + # This is intentionally minimal: only async_success_handler / post_call_failure_hook + rules_obj: Final = Rules() + list_tools_call_id: Final = str(uuid.uuid4()) + # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) + effective_litellm_trace_id: Final = litellm_trace_id or get_chain_id_from_headers(raw_headers) + spend_logs_metadata: Final[dict[str, object]] = { + "mcp_operation": "list_tools", + } + if isinstance(list_tools_log_source, str): + spend_logs_metadata["source"] = list_tools_log_source + if isinstance(mcp_servers, list): + spend_logs_metadata["requested_mcp_servers"] = mcp_servers + + list_tools_request_data = { + "model": "MCP: list_tools", + "call_type": CallTypes.list_mcp_tools.value, + "litellm_call_id": list_tools_call_id, + "litellm_trace_id": effective_litellm_trace_id, + "metadata": { + "spend_logs_metadata": spend_logs_metadata, + "headers": logging_safe_mcp_headers(raw_headers), + **({"tags": request_tags} if request_tags else {}), + }, + # Provide a small input payload for standard logging + "input": [ + { + "role": "system", + "content": { + "mcp_operation": "list_tools", + "requested_mcp_servers": mcp_servers, + }, + } + ], + } + + # Attach user identifiers using the standard helper + if user_api_key_auth is not None: + LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data=list_tools_request_data, + user_api_key_dict=user_api_key_auth, + _metadata_variable_name="metadata", + ) + + user_identifier: Final = getattr(user_api_key_auth, "end_user_id", None) or getattr( + user_api_key_auth, "user_id", None + ) + if user_identifier: + list_tools_request_data["user"] = user_identifier + + try: + litellm_logging_obj, _ = function_setup( + original_function="list_mcp_tools", + is_async_call=False, + rules_obj=rules_obj, + start_time=list_tools_start_time, + **list_tools_request_data, + ) + if litellm_logging_obj: + litellm_logging_obj.call_type = CallTypes.list_mcp_tools.value + litellm_logging_obj.model = "MCP: list_tools" + except Exception as logging_error: + verbose_logger.debug("Failed to initialize logging for MCP list_tools: %s", logging_error) + litellm_logging_obj = None + + try: + allowed_mcp_servers: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_servers=mcp_servers, + client_ip=client_ip, + ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) + + # Pre-fetch OAuth credentials only when at least one server uses OAuth2, + # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. + _has_oauth2_server = any(getattr(s, "auth_type", None) == MCPAuth.oauth2 for s in allowed_mcp_servers) + _prefetched_oauth_creds: Final = ( + await _prefetch_oauth_creds_for_user(user_api_key_auth) if _has_oauth2_server else {} + ) + + async def _fetch_and_filter_server_tools( + server: MCPServer, + ) -> "tuple[list[MCPTool], ServerOutcome]": + """Fetch and filter tools from a single server, classifying any failure into that + server's outcome so the aggregate can keep serving the healthy subset without a + broken server masquerading as an empty one.""" + if server is None: + return [], ServerListOk(tool_count=0) + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=mcp_server_auth_headers, + mcp_auth_header=mcp_auth_header, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + scope_servers=allowed_mcp_servers, + ) + + # Prefer server-stored per-user OAuth when configured, so a stale + # Authorization header from the MCP client cannot override Redis/DB + # (same issue as call_tool in mcp_server_manager: VS Code caches tokens). + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 + to_server_spec, + ) + + # A server migrated to the v2 resolver gets its token from the resolver at connect + # time; building it here would double-resolve and be shadowed by the v2 graft. The + # preemptive 401 already challenged a missing token, so one exists for the connect. + migrated_to_v2: Final = to_server_spec(server) is not None + if ( + not migrated_to_v2 + and server.auth_type == MCPAuth.oauth2 + and getattr(server, "needs_user_oauth_token", False) + and user_api_key_auth is not None + ): + db_headers: Final = await _get_user_oauth_extra_headers_from_db( + server, + user_api_key_auth, + prefetched_creds=_prefetched_oauth_creds, + ) + if db_headers: + extra_headers = db_headers + + # If still no OAuth2 token, fall back to pre-fetched creds (non-stale-client path) + elif not migrated_to_v2 and extra_headers is None and server.auth_type == MCPAuth.oauth2: + extra_headers = await _get_user_oauth_extra_headers_from_db( + server, + user_api_key_auth, + prefetched_creds=_prefetched_oauth_creds, + ) + + if server.is_byok and server.auth_type != MCPAuth.oauth2 and server_auth_header is None: + server_auth_header = await _get_byok_credential(server, user_api_key_auth) + + try: + tools: Final = await global_mcp_server_manager._get_tools_from_server( + server=server, + mcp_auth_header=server_auth_header, + extra_headers=extra_headers, + add_prefix=True, # Always add server prefix + raw_headers=raw_headers, + client_ip=client_ip, + user_api_key_auth=user_api_key_auth, + oauth2_headers=oauth2_headers, + ) + filtered_tools = filter_tools_by_allowed_tools(tools, server) + + filtered_tools = await filter_tools_by_key_team_permissions( + tools=filtered_tools, + server_id=server.server_id, + user_api_key_auth=user_api_key_auth, + ) + + if mcp_proxy_mode: + from litellm.proxy._experimental.mcp_server.tool_search import with_mcp_proxy_identity + + filtered_tools = [ # mutable-ok: MCP tool pipeline + with_mcp_proxy_identity(tool, server.server_id) for tool in filtered_tools + ] + else: + filtered_tools = apply_tool_overrides(filtered_tools, server) + + verbose_logger.debug( + "Successfully fetched %s tools from server %s, %s after filtering", + len(tools), + server.name, + len(filtered_tools), + ) + return filtered_tools, ServerListOk(tool_count=len(filtered_tools)) + except MCPUpstreamAuthError as e: + # Absorb so one unauthenticated server does not empty every other server's + # tools. Surfacing the upstream 401 to the client as a re-auth challenge is + # intentionally not done here: raising from this list handler cannot produce a + # 401 + WWW-Authenticate (the MCP session manager serializes it as a JSON-RPC + # error). Single-server routes surface it via the request-scope preemptive + # check in _raise_preemptive_401_for_unauthenticated_servers instead. + verbose_logger.debug("MCP list_tools: omitting %s; it needs upstream auth", server.name) + return [], classify_list_exception(e) + except Exception as e: + verbose_logger.exception("Error getting tools from server %s: %s", server.name, e) + return [], classify_list_exception(e) + + # Fetch tools from all servers in parallel + tasks: Final = [_fetch_and_filter_server_tools(server) for server in allowed_mcp_servers] + results: Final = await asyncio.gather(*tasks) + + # Flatten results into single list + all_tools: Final[list[MCPTool]] = [tool for tools, _ in results for tool in tools] + server_outcomes: Final[dict[str, ServerOutcome]] = { + _aggregate_server_key(server): outcome + for server, (_, outcome) in zip(allowed_mcp_servers, results) + if server is not None + } + + # If logging is enabled, enrich spend_logs_metadata with counts + if litellm_logging_obj: + per_server_tool_counts: Final[dict[str, int]] = { + _aggregate_server_key(server): len(server_tools) + for server, (server_tools, _) in zip(allowed_mcp_servers, results) + if server is not None + } + + metadata_dict: Final = litellm_logging_obj.model_call_details.get("metadata") + if isinstance(metadata_dict, dict): + spend_meta = metadata_dict.get("spend_logs_metadata") + if not isinstance(spend_meta, dict): + spend_meta = {} + metadata_dict["spend_logs_metadata"] = spend_meta + spend_meta["allowed_server_count"] = len(allowed_mcp_servers) + spend_meta["tool_count_total"] = len(all_tools) + spend_meta["per_server_tool_counts"] = per_server_tool_counts + spend_meta["per_server_list_outcomes"] = { + key: outcome_wire_value(outcome) for key, outcome in server_outcomes.items() + } + + end_time: Final = datetime.now() + try: + await litellm_logging_obj.async_success_handler( + result=[tool.model_dump(mode="json") if isinstance(tool, MCPTool) else tool for tool in all_tools], + start_time=list_tools_start_time, + end_time=end_time, + ) + except Exception as log_exc: + # list_tools responses must not be dropped due to non-blocking + # observability/serialization failures. + verbose_logger.warning( + "MCP list_tools success logging failed (continuing): %s", + log_exc, + ) + + verbose_logger.info("Successfully fetched %s tools total from all MCP servers", len(all_tools)) + + return AggregateToolListing(tools=all_tools, outcomes=server_outcomes) + except Exception as e: + # Only fire failure hook if logging was requested for this list-tools execution + if log_list_tools_to_spendlogs and user_api_key_auth is not None: + try: + from litellm.proxy.proxy_server import proxy_logging_obj + + if proxy_logging_obj: + traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + await proxy_logging_obj.post_call_failure_hook( + request_data=list_tools_request_data or {}, + original_exception=e, + user_api_key_dict=user_api_key_auth, + route="/mcp/list_tools", + traceback_str=traceback_str, + ) + except Exception: + verbose_logger.debug("Failed to log MCP list_tools failure via post_call_failure_hook") + raise + + +async def _get_prompts_from_mcp_servers( + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None, + mcp_servers: list[str] | 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, + client_ip: str | None = None, +) -> list[Prompt]: + """ + Helper method to fetch prompt from MCP servers based on server filtering criteria. + + Args: + user_api_key_auth: User authentication info for access control + mcp_auth_header: Optional auth header for MCP server (deprecated) + mcp_servers: Optional list of server names/aliases to filter by + mcp_server_auth_headers: Optional dict of server-specific auth headers + oauth2_headers: Optional dict of oauth2 headers + + Returns: + List[Prompt]: Combined list of prompts from filtered servers + """ + + allowed_mcp_servers: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_servers=mcp_servers, + client_ip=client_ip, + ) + + # Get prompts from each allowed server + all_prompts: Final = [] + for server in allowed_mcp_servers: + if server is None: + continue + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=mcp_server_auth_headers, + mcp_auth_header=mcp_auth_header, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + scope_servers=allowed_mcp_servers, + ) + + try: + prompts = await global_mcp_server_manager.get_prompts_from_server( + server=server, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=server_auth_header, + extra_headers=extra_headers, + add_prefix=True, # Always add server prefix + raw_headers=raw_headers, + client_ip=client_ip, + ) + + all_prompts.extend(prompts) + + verbose_logger.debug("Successfully fetched %s prompts from server %s", len(prompts), server.name) + except Exception as e: + verbose_logger.exception("Error getting prompts from server %s: %s", server.name, e) + # Continue with other servers instead of failing completely + + verbose_logger.info("Successfully fetched %s prompts total from all MCP servers", len(all_prompts)) + + return all_prompts + + +async def _get_resources_from_mcp_servers( + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None, + mcp_servers: list[str] | 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, + client_ip: str | None = None, +) -> list[Resource]: + """Fetch resources from allowed MCP servers.""" + + allowed_mcp_servers: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_servers=mcp_servers, + client_ip=client_ip, + ) + + all_resources: Final[list[Resource]] = [] + for server in allowed_mcp_servers: + if server is None: + continue + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=mcp_server_auth_headers, + mcp_auth_header=mcp_auth_header, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + scope_servers=allowed_mcp_servers, + ) + + try: + resources = await global_mcp_server_manager.get_resources_from_server( + server=server, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=server_auth_header, + extra_headers=extra_headers, + add_prefix=True, # Always add server prefix + raw_headers=raw_headers, + client_ip=client_ip, + ) + all_resources.extend(resources) + + verbose_logger.debug("Successfully fetched %s resources from server %s", len(resources), server.name) + except Exception as e: + verbose_logger.exception("Error getting resources from server %s: %s", server.name, e) + + verbose_logger.info("Successfully fetched %s resources total from all MCP servers", len(all_resources)) + + return all_resources + + +async def _get_resource_templates_from_mcp_servers( + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None, + mcp_servers: list[str] | 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, + client_ip: str | None = None, +) -> list[ResourceTemplate]: + """Fetch resource templates from allowed MCP servers.""" + + allowed_mcp_servers: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_servers=mcp_servers, + client_ip=client_ip, + ) + + all_resource_templates: Final[list[ResourceTemplate]] = [] + for server in allowed_mcp_servers: + if server is None: + continue + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=mcp_server_auth_headers, + mcp_auth_header=mcp_auth_header, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + scope_servers=allowed_mcp_servers, + ) + + try: + resource_templates = await global_mcp_server_manager.get_resource_templates_from_server( + server=server, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=server_auth_header, + extra_headers=extra_headers, + add_prefix=True, # Always add server prefix + raw_headers=raw_headers, + client_ip=client_ip, + ) + all_resource_templates.extend(resource_templates) + verbose_logger.debug( + "Successfully fetched %s resource templates from server %s", + len(resource_templates), + server.name, + ) + except Exception as e: + verbose_logger.exception( + "Error getting resource templates from server %s: %s", + server.name, + str(e), + ) + + verbose_logger.info( + "Successfully fetched %s resource templates total from all MCP servers", + len(all_resource_templates), + ) + + return all_resource_templates + + +async def filter_tools_by_key_team_permissions( + tools: list[MCPTool], + server_id: str, + user_api_key_auth: UserAPIKeyAuth | None, +) -> list[MCPTool]: + """ + Filter tools based on key/team mcp_tool_permissions. + + Note: Tool names in the DB are stored without server prefixes, + but tool names from MCP servers are prefixed. We need to strip + the prefix before comparing. + """ + # Filter by key/team tool-level permissions + allowed_tool_names: Final = await MCPRequestHandler.get_allowed_tools_for_server( + server_id=server_id, + user_api_key_auth=user_api_key_auth, + ) + + # Tools arrive prefixed with the server's own prefix; strip exactly that + # prefix (resolved from the server) rather than the first separator, so a + # prefix containing the separator still reduces to the stored bare name. + server: Final = global_mcp_server_manager.get_mcp_server_by_id(server_id) + return [ + t + for t in tools + if MCPRequestHandler.tool_is_granted(strip_known_server_prefix(t.name, server), allowed_tool_names) + ] + + +async def _list_mcp_tools( + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[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, + log_list_tools_to_spendlogs: bool = False, + list_tools_log_source: str | None = None, + client_ip: str | None = None, + mcp_proxy_mode: bool = False, +) -> AggregateToolListing: + """ + List all available MCP tools. + + Args: + user_api_key_auth: User authentication info for access control + mcp_auth_header: Optional auth header for MCP server (deprecated) + mcp_servers: Optional list of server names/aliases to filter by + mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} + client_ip: Client IP for IP-based server access control + + Returns: + AggregateToolListing: Combined tools from all accessible servers plus each server's + classified listing outcome + """ + + try: + listing: Final = await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + log_list_tools_to_spendlogs=log_list_tools_to_spendlogs, + list_tools_log_source=list_tools_log_source, + client_ip=client_ip, + mcp_proxy_mode=mcp_proxy_mode, + ) + verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools)) + return listing + except HTTPException: + raise + except Exception as e: + verbose_logger.exception("Error getting tools from managed MCP servers: %s", e) + # Continue with an empty listing instead of failing completely + return AggregateToolListing(tools=[], outcomes={}) + + +async def _list_mcp_prompts( + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[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, + client_ip: str | None = None, +) -> list[Prompt]: + """ + List all available MCP prompts. + + Args: + user_api_key_auth: User authentication info for access control + mcp_auth_header: Optional auth header for MCP server (deprecated) + mcp_servers: Optional list of server names/aliases to filter by + mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} + + Returns: + List[Prompt]: Combined list of tools from all accessible servers + """ + # Get tools from managed MCP servers with error handling + managed_prompts = [] + try: + managed_prompts = await _get_prompts_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + verbose_logger.debug("Successfully fetched %s prompts from managed MCP servers", len(managed_prompts)) + except Exception as e: + verbose_logger.exception("Error getting tools from managed MCP servers: %s", e) + # Continue with empty managed tools list instead of failing completely + + return managed_prompts + + +async def _list_mcp_resources( + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[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, + client_ip: str | None = None, +) -> list[Resource]: + """List all available MCP resources.""" + + managed_resources: list[Resource] = [] + try: + managed_resources = await _get_resources_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + verbose_logger.debug("Successfully fetched %s resources from managed MCP servers", len(managed_resources)) + except Exception as e: + verbose_logger.exception("Error getting resources from managed MCP servers: %s", e) + + return managed_resources + + +async def _list_mcp_resource_templates( + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[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, + client_ip: str | None = None, +) -> list[ResourceTemplate]: + """List all available MCP resource templates.""" + + managed_resource_templates: list[ResourceTemplate] = [] + try: + managed_resource_templates = await _get_resource_templates_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + verbose_logger.debug( + "Successfully fetched %s resource templates from managed MCP servers", + len(managed_resource_templates), + ) + except Exception as e: + verbose_logger.exception( + "Error getting resource templates from managed MCP servers: %s", + str(e), + ) + + return managed_resource_templates + + +def _resolve_display_name_to_original( + name: str, + allowed_mcp_servers: list[MCPServer], +) -> str: + """Translate a display-name override back to the original prefixed tool name. + + When a client received a customised display name from tools/list (e.g. + "Get Pet") it will call tools/call with that same string. We need to + reverse-map it to the original prefixed name (e.g. + "petstore_mcp-getPetById") before any routing or permission logic runs. + """ + for server in allowed_mcp_servers: + display_map = server.tool_name_to_display_name or {} + for unprefixed_name, display_name in display_map.items(): + if display_name == name: + return add_server_prefix_to_name(unprefixed_name, get_server_prefix(server)) + return name + + +async def _get_byok_credential( + mcp_server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, +) -> str | None: + """Retrieve the stored BYOK credential for a user+server pair, served from the worker cache within its TTL.""" + if not mcp_server.is_byok: + return None + user_id: Final = (user_api_key_auth.user_id if user_api_key_auth else None) or "" + if not user_id: + return None + + cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id) + if cached is not None: + return cached.credential + + from litellm.proxy._experimental.mcp_server.db import get_user_credential + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return None + credential: Final = await get_user_credential( + prisma_client=prisma_client, + user_id=user_id, + server_id=mcp_server.server_id, + ) + cache_byok_credential(user_id, mcp_server.server_id, credential) + return credential + + +async def _check_byok_credential( + mcp_server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, +) -> None: + """ + If the MCP server is BYOK-enabled, verify that the requesting user has a + stored credential. When no credential is found, raise an HTTP 401 with a + WWW-Authenticate header that points the MCP client to our OAuth metadata + endpoint so it can drive the authorization flow. + """ + if not mcp_server.is_byok: + return + + user_id: Final = (user_api_key_auth.user_id if user_api_key_auth else None) or "" + if not user_id: + raise HTTPException( + status_code=401, + detail={ + "error": "byok_auth_required", + "server_id": mcp_server.server_id, + "server_name": mcp_server.server_name or mcp_server.name, + "message": "User identity is required for BYOK servers", + }, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, + ) + + cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id) + if cached is not None: + if cached.credential is None: + raise HTTPException( + status_code=401, + detail={ + "error": "byok_auth_required", + "server_id": mcp_server.server_id, + "server_name": mcp_server.server_name or mcp_server.name, + "message": ( + "No stored credential found for this BYOK server. " + "Complete the OAuth authorization flow to provide your API key." + ), + }, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, + ) + return + + from litellm.proxy._experimental.mcp_server.db import get_user_credential + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + # Fail closed on DB unavailability: returning here previously + # bypassed the ownership check and let any proxy-authenticated + # caller invoke BYOK tools during outage windows. + raise HTTPException( + status_code=503, + detail={ + "error": "byok_auth_unavailable", + "server_id": mcp_server.server_id, + "server_name": mcp_server.server_name or mcp_server.name, + "message": "BYOK credential check requires a database connection.", + }, + ) + + credential: Final = await get_user_credential( + prisma_client=prisma_client, + user_id=user_id, + server_id=mcp_server.server_id, + ) + cache_byok_credential(user_id, mcp_server.server_id, credential) + if credential is None: + raise HTTPException( + status_code=401, + detail={ + "error": "byok_auth_required", + "server_id": mcp_server.server_id, + "server_name": mcp_server.server_name or mcp_server.name, + "message": ( + "No stored credential found for this BYOK server. " + "Complete the OAuth authorization flow to provide your API key." + ), + }, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, + ) + + +async def _list_tools_before_first_call( + server: MCPServer | None, + tool_name: str, + allowed_mcp_servers: list[MCPServer], + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + oauth2_headers: dict[str, str] | None, + raw_headers: dict[str, str] | None, + client_ip: str | None = None, +) -> None: + """List ``server`` with the caller's own credentials when it does not yet expose ``tool_name`` here. + + The startup fill skips a server whose upstream wants the caller's token, and mcp 2 no + longer lists before an uncached tools/call, so a worker that has not served tools/list + for this caller would otherwise answer 404 for a tool the caller can see. Gating on the + requested tool, not on any prior listing, keeps callers with different upstream catalogs + from masking each other. + """ + if server is None or global_mcp_server_manager.server_exposes_tool(server, tool_name): + return + if all(allowed.server_id != server.server_id for allowed in allowed_mcp_servers): + return + try: + await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=[server.server_id], + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + except Exception as e: # noqa: BLE001 # best effort: resolution below answers as it did before + verbose_logger.debug("MCP tools/call: listing %s before its first call failed: %s", server.name, e) + + +async def execute_mcp_tool( + name: str, + arguments: dict[str, object], + allowed_mcp_servers: list[MCPServer], + start_time: datetime, + user_api_key_auth: UserAPIKeyAuth | None = 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, + host_progress_callback: ProgressCallback | None = None, + guardrail_context: Mapping[str, object] | None = None, + client_ip: str | None = None, + **kwargs: object, # kwargs-ok: preserves the existing REST and decorated logging call contract +) -> CallToolResult: + context: Final = prepare_context( + user_api_key_auth=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, + client_ip=client_ip, + ) + operation: Final = AuthorizedToolCall( + name=name, + arguments=arguments, + allowed_mcp_servers=tuple(allowed_mcp_servers), + start_time=start_time, + host_progress_callback=host_progress_callback, + guardrail_context=guardrail_context, + logging_data=types.MappingProxyType(kwargs), + ) + return await GatewayOperations().execute(operation, context) + + +async def _execute_mcp_tool( + name: str, + arguments: dict[str, object], + allowed_mcp_servers: list[MCPServer], + start_time: datetime, + user_api_key_auth: UserAPIKeyAuth | None = 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, + host_progress_callback: ProgressCallback | None = None, + guardrail_context: Mapping[str, object] | None = None, + client_ip: str | None = None, + **kwargs: Any, +) -> CallToolResult: + """ + Execute MCP tool. + + This function assumes permission checks have already been performed. + + Args: + name: Tool name (may include server prefix) + arguments: Tool arguments + allowed_mcp_servers: Pre-validated list of servers the user can access + start_time: Start time for logging + user_api_key_auth: Optional user API key auth for logging + mcp_auth_header: Optional MCP auth header + mcp_server_auth_headers: Optional server-specific auth headers + oauth2_headers: Optional OAuth2 headers + raw_headers: Optional raw HTTP headers + **kwargs: Additional arguments (e.g., litellm_logging_obj) + + Returns: + CallToolResult: Tool execution result + """ + # Track resolved MCP server for both permission checks and dispatch + mcp_server: MCPServer | None = None + requested_server_id: Final[str | None] = kwargs.get("requested_server_id") + + # If the client called with a display-name override (e.g. "Get Pet"), + # translate it back to the original prefixed name before any routing. + name = _resolve_display_name_to_original(name, allowed_mcp_servers) + + # Remove prefix from tool name for logging and processing + original_tool_name, server_name = split_server_prefix_from_name(name) + + requested_server: MCPServer | None = None + if requested_server_id: + requested_server = next( + (s for s in allowed_mcp_servers if s.server_id == requested_server_id), + None, + ) + + name_is_prefixed = False + if requested_server is not None and MCP_TOOL_PREFIX_SEPARATOR in name: + all_registry_prefixes: Final[set[str]] = set() + for registry_server in global_mcp_server_manager.get_registry().values(): + for known_prefix in iter_known_server_prefixes(registry_server): + all_registry_prefixes.add(normalize_server_name(known_prefix)) + name_is_prefixed = is_tool_name_prefixed(name, known_server_prefixes=all_registry_prefixes) + + first_call_target: Final = ( + requested_server + if requested_server is not None and not name_is_prefixed + else global_mcp_server_manager.server_owning_tool_name_prefix(name) + ) + first_call_tool_name: Final = ( + name + if first_call_target is None or (requested_server is not None and not name_is_prefixed) + else strip_known_server_prefix(name, first_call_target) + ) + await _list_tools_before_first_call( + server=first_call_target, + tool_name=first_call_tool_name, + allowed_mcp_servers=allowed_mcp_servers, + user_api_key_auth=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, + client_ip=client_ip, + ) + + if requested_server is not None and not name_is_prefixed: + # REST callers may pass server_id with the upstream tool name (no + # LiteLLM prefix). The first segment is not a registered server + # prefix, so the whole string is the upstream tool name and may + # legitimately contain the separator (e.g. "text-to-speech"). + # server_id is authoritative for routing and auth. + mcp_server = requested_server + server_name = requested_server.name + original_tool_name = name + else: + # Resolve from tool name (MCP JSON-RPC or prefixed REST tool names). + mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + if mcp_server is None and requested_server is not None: + for known_prefix in iter_known_server_prefixes(requested_server): + candidate = global_mcp_server_manager._get_mcp_server_from_tool_name( + add_server_prefix_to_name(name, known_prefix) + ) + if candidate is not None: + mcp_server = candidate + break + if mcp_server is not None: + server_name = mcp_server.name + original_tool_name = strip_known_server_prefix(name, mcp_server) + + if requested_server is not None: + if mcp_server is not None and mcp_server.server_id != requested_server.server_id: + raise HTTPException( + status_code=403, + detail={ + "error": "tool_server_mismatch", + "message": ( + f"Tool '{name}' belongs to MCP server " + f"'{mcp_server.name}' but request specified " + f"server_id for '{requested_server.name}'." + ), + }, + ) + if mcp_server is None: + mcp_server = requested_server + server_name = requested_server.name + original_tool_name = strip_known_server_prefix(name, requested_server) + + # Only enforce server-level permissions when we can resolve a server + if server_name: + if not MCPRequestHandler.is_tool_allowed( + allowed_mcp_servers=[server.name for server in allowed_mcp_servers], + server_name=server_name, + ): + raise HTTPException( + status_code=403, + detail="User not allowed to call this tool.", + ) + + standard_logging_mcp_tool_call: Final[StandardLoggingMCPToolCall] = _get_standard_logging_mcp_tool_call( + name=original_tool_name, # Use original name for logging + arguments=arguments, + server_name=server_name, + session_id=_mcp_session_id_from_headers(raw_headers), + ) + litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None) + 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: {name}" + litellm_logging_obj.model_call_details["model"] = f"MCP: {name}" + # Resolve the MCP server early so BYOK checks and credential injection + # apply to ALL dispatch paths (local tool registry AND managed MCP server). + if mcp_server is None: + mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + + if mcp_server: + standard_logging_mcp_tool_call["mcp_server_cost_info"] = (mcp_server.mcp_info or {}).get("mcp_server_cost_info") + if litellm_logging_obj: + litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call + + # BYOK: retrieve the stored per-user credential. A single DB call + # both checks existence and fetches the value, avoiding a double query. + if mcp_server.is_byok and not mcp_auth_header: + byok_cred: Final = await _get_byok_credential(mcp_server, user_api_key_auth) + if byok_cred is None: + raise HTTPException( + status_code=401, + detail={ + "error": "byok_auth_required", + "server_id": mcp_server.server_id, + "server_name": mcp_server.server_name or mcp_server.name, + "message": ( + "No stored credential found for this BYOK server. " + "Complete the OAuth authorization flow to provide your API key." + ), + }, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, + ) + mcp_auth_header = byok_cred + elif mcp_server.is_byok: + # External auth header supplied; still enforce user-identity check. + await _check_byok_credential(mcp_server, user_api_key_auth) + + # Check if tool exists in local registry first (for OpenAPI-based tools) + # These tools are registered with their prefixed names + ######################################################### + local_tool: Final = global_mcp_tool_registry.get_tool(name) + if local_tool: + # OpenAPI-backed tools used to bypass `pre_call_tool_check` — + # only the managed path ran allowed/banned-tool checks, key/team + # tool permissions, and parameter validation. Run the same checks + # before dispatching to the local registry. Refuse the call if + # we cannot resolve a server: tools registered via + # openapi_to_mcp_generator are always tied to a server, so a + # missing mcp_server here means the tool->server mapping has + # not finished initializing or the registry entry is orphaned. + # Skipping the check would re-open the same authorization gap. + if mcp_server is None: + raise HTTPException( + status_code=503, + detail=( + f"MCP server for tool '{name}' is not available; " + "refusing to dispatch without authorization checks. " + "Retry once the server is registered." + ), + ) + + # `pre_call_tool_check` calls into `proxy_logging_obj` for the + # pre-call guardrail hooks, so source it from the canonical + # `proxy_server` module the same way `_handle_managed_mcp_tool` + # does. `kwargs.get("proxy_logging_obj")` is None on the MCP + # entry path and would crash with AttributeError after the + # security checks pass. + from litellm.proxy.proxy_server import proxy_logging_obj + + hook_result = await global_mcp_server_manager.pre_call_tool_check( + name=original_tool_name, + arguments=arguments or {}, + server_name=server_name or mcp_server.name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=mcp_server, + raw_headers=raw_headers, + litellm_logging_obj=litellm_logging_obj, + guardrail_context=guardrail_context, + ) + # `pre_call_tool_check` may return guardrail-modified + # arguments; honor them on the local path too. + if isinstance(hook_result, dict) and "arguments" in hook_result: + arguments = hook_result["arguments"] + + verbose_logger.debug("Executing local registry tool: %s", name) + # The credential rides ContextVars because the tool function has its + # headers baked into the closure at registration time. + auth_header_value, openapi_forwarded_headers, upstream_credential = _resolve_openapi_tool_auth( + mcp_server=mcp_server, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + ) + ( + resolved_auth_headers, + forwarded_headers, + ) = await global_mcp_server_manager.resolve_openapi_upstream_auth( + mcp_server=mcp_server, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + mcp_auth_header=upstream_credential, + user_api_key_auth=user_api_key_auth, + forwarded_headers=openapi_forwarded_headers, + ) + + _auth_token: Final = _request_auth_header.set(auth_header_value) + _extra_token: Final = _request_extra_headers.set(forwarded_headers) + _resolved_token: Final = _request_resolved_auth_headers.set(resolved_auth_headers) + try: + response = await _handle_local_mcp_tool(name, arguments) + finally: + _request_auth_header.reset(_auth_token) + _request_extra_headers.reset(_extra_token) + _request_resolved_auth_headers.reset(_resolved_token) + + # Try managed MCP server tool (the name is bare; the prefix boundary was + # already resolved above against this server's registered prefixes) + # Primary and recommended way to use external MCP servers + ######################################################### + elif mcp_server: + response = await _handle_managed_mcp_tool( + server_name=server_name, + name=original_tool_name, + arguments=arguments, + user_api_key_auth=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, + client_ip=client_ip, + litellm_logging_obj=litellm_logging_obj, + guardrail_context=guardrail_context, + host_progress_callback=host_progress_callback, + ) + + # Fall back to local tool registry with original name (legacy support) + ######################################################### + # Deprecated: Local MCP Server Tool + ######################################################### + else: + # Gate only what can actually dispatch. When the unprefixed name is + # not in the registry either, `_handle_local_mcp_tool` below reports + # 404 and nothing runs, so demanding a server here would turn every + # unknown tool name into a misleading 503. + if global_mcp_tool_registry.get_tool(original_tool_name) is not None: + # `mcp_server` is None here because the tool name is not in the + # tool -> server mapping, but the name still carries a prefix + # that the server-level check above compared against the + # caller's `allowed_mcp_servers` by exact `name`. So the named + # server is in that list and can carry the tool-level checks, + # even with the mapping cold. Resolve it from + # `allowed_mcp_servers` rather than the registry: the registry + # would happily return a server the caller holds no grant for, + # and matching anything other than `name` would accept a server + # the check never validated. + prefix_server: Final = next( + (candidate for candidate in allowed_mcp_servers if candidate.name == server_name), + None, + ) + if prefix_server is None: + # A non-empty prefix that passed the server-level check + # always matches here, so this arm only fires when the + # prefix was empty, which is exactly the case that check + # skips. Fail closed rather than dispatch with no server to + # evaluate a tool ceiling against. + raise HTTPException( + status_code=503, + detail=( + f"MCP server for tool '{original_tool_name}' is not available; " + "refusing to dispatch without authorization checks. " + "Retry once the server is registered." + ), + ) + + from litellm.proxy.proxy_server import proxy_logging_obj + + hook_result = await global_mcp_server_manager.pre_call_tool_check( + name=original_tool_name, + arguments=arguments, + server_name=server_name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=prefix_server, + raw_headers=raw_headers, + litellm_logging_obj=litellm_logging_obj, + guardrail_context=guardrail_context, + ) + if "arguments" in hook_result: + arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args + + response = await _handle_local_mcp_tool(original_tool_name, arguments) + + return await _run_post_mcp_call_guardrails( + result=response, + litellm_logging_obj=litellm_logging_obj, + user_api_key_auth=user_api_key_auth, + request_data=kwargs, + ) + + +async def _run_post_mcp_call_guardrails( + result: CallToolResult, + litellm_logging_obj: LiteLLMLoggingObj | None, + user_api_key_auth: UserAPIKeyAuth | None, + request_data: Mapping[str, object], +) -> CallToolResult: + """Run ``post_mcp_call`` guardrails over an executed tool result. + + Lives on ``execute_mcp_tool``'s return path rather than inside + ``_fire_mcp_tool_call_logging`` so enforcement never depends on logging + being configured, and so every dispatch route gets it: the MCP protocol + handler, the REST endpoint, and tool search all funnel through here. + A guardrail that rejects the result raises, matching ``pre_mcp_call``. + """ + from litellm.proxy.proxy_server import proxy_logging_obj + + if proxy_logging_obj is None: + return result + return await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data=( + litellm_logging_obj.model_call_details if litellm_logging_obj is not None else dict(request_data) + ), + user_api_key_dict=user_api_key_auth, + ) + + +async def _fire_mcp_tool_call_logging( + logging_obj: LiteLLMLoggingObj, + result: CallToolResult, + start_time: datetime, + end_time: datetime, + user_api_key_auth: UserAPIKeyAuth | None = None, + request_data: Mapping[str, object] | None = None, +) -> CallToolResult: + """Fire post-call logging for an executed MCP tool call, returning the result to send. + + The returned result is what the caller must forward to the client: a + ``post_mcp_call`` guardrail may rewrite the tool output (e.g. mask + sensitive values) or reject it, in which case its exception propagates. + Guardrails run before the success/failure logging so the masked text, not + the raw one, is what gets logged. + + A result with ``is_error=True`` is logged as a failure (``status="failure"`` + payload, so OTel marks the span ERROR) while the HTTP wire behavior stays + 200 + ``isError: true`` per the MCP spec. The error check runs after + ``async_post_mcp_tool_call_hook`` because guardrails may flip the result + to ``is_error=True`` in that hook. Raised exceptions never reach here (the + ``@client`` wrapper and ``call_mcp_tool``'s except path log those), so + this cannot double-log a failure. + + ``request_data`` may carry credential-bearing fields (the REST path puts + ``raw_headers``, ``mcp_auth_header``, ``mcp_server_auth_headers``, and + ``oauth2_headers`` at the top level of its data dict), so those are + stripped before the dict is handed to ``post_call_failure_hook`` + callbacks. + """ + from litellm.proxy.proxy_server import proxy_logging_obj + + logging_obj.post_call(original_response=result) + await logging_obj.async_post_mcp_tool_call_hook( + kwargs=logging_obj.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + logging_obj.call_type = CallTypes.call_mcp_tool.value + error_message: Final = extract_mcp_tool_result_error_message(result) + if error_message is None: + await logging_obj.async_success_handler(result=result, start_time=start_time, end_time=end_time) + return result + + logging_obj.has_run_logging(event_type="sync_success") + logging_obj.has_run_logging(event_type="async_success") + tool_error: Final = MCPToolResultError(error_message) + logging_obj.failure_handler(tool_error, "", start_time, end_time) + await logging_obj.async_failure_handler(tool_error, "", start_time, end_time) + + if user_api_key_auth is None: + return result + + if proxy_logging_obj: + sanitized_request_data: Final = { + key: value for key, value in (request_data or {}).items() if key not in _MCP_CREDENTIAL_REQUEST_FIELDS + } + await proxy_logging_obj.post_call_failure_hook( + request_data=sanitized_request_data, + original_exception=tool_error, + user_api_key_dict=user_api_key_auth, + route="/mcp/call_tool", + ) + return result + + +async def fire_mcp_tool_call_failure_logging( + logging_obj: LiteLLMLoggingObj | None, + exception: Exception, + start_time: datetime, + user_api_key_auth: UserAPIKeyAuth | None, + request_data: Mapping[str, object], +) -> None: + """Failure logging shared by the ``/mcp`` path and the REST endpoint. Call from + inside the ``except`` block so the traceback is still available. + + The failure handlers run first because ``_ProxyDBLogger.async_post_call_failure_hook`` + builds the failure spend-log row from the ``standard_logging_object`` they produce; + both gate on ``should_run_logging``, so the ``@client`` wrapper does not log twice. + A relayed upstream 401 (``MCPUpstreamAuthError``) is an expected caller-must-reauth + signal and skips ``post_call_failure_hook``, which fires the ``llm_exceptions`` alert. + """ + from litellm.proxy.proxy_server import proxy_logging_obj + + traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + if logging_obj is not None: + end_time: Final = datetime.now() # noqa: DTZ005 # naive to match `start_time`, which it is subtracted from + logging_obj.failure_handler(exception, traceback_str, start_time, end_time) + await logging_obj.async_failure_handler(exception, traceback_str, start_time, end_time) + + if isinstance(exception, MCPUpstreamAuthError) or not proxy_logging_obj or user_api_key_auth is None: + return + sanitized_request_data: Final = { + key: value for key, value in request_data.items() if key not in _MCP_CREDENTIAL_REQUEST_FIELDS + } + await proxy_logging_obj.post_call_failure_hook( + request_data=sanitized_request_data, + original_exception=exception, + user_api_key_dict=user_api_key_auth, + route="/mcp/call_tool", + traceback_str=traceback_str, + ) + + +@client +async def call_mcp_tool( + name: str, + arguments: dict[str, object] | None = None, + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[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, + client_ip: str | None = None, + **kwargs: Any, +) -> CallToolResult: + """ + Call a specific tool with the provided arguments (handles prefixed tool names). + """ + start_time: Final = datetime.now() + litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None) + + try: + if arguments is None: + raise HTTPException(status_code=400, detail="Request arguments are required") + + ## CHECK IF USER IS ALLOWED TO CALL THIS TOOL + allowed_mcp_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + ) + + allowed_mcp_servers: list[MCPServer] = [] + for allowed_mcp_server_id in allowed_mcp_server_ids: + allowed_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id) + if allowed_server is not None: + # Same request-time oauth2_flow backstop the listing path applies, + # so a null-flow M2M-shape row is treated as M2M on tool calls too. + allowed_server = MCPServerManager.resolve_oauth2_flow_for_request(allowed_server) + allowed_mcp_servers.append(allowed_server) + + allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( + mcp_servers=mcp_servers, + allowed_mcp_servers=allowed_mcp_servers, + ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) + if not allowed_mcp_servers: + raise HTTPException( + status_code=403, + detail="User not allowed to call this tool.", + ) + + # Delegate to execute_mcp_tool for execution + response = await execute_mcp_tool( + name=name, + arguments=arguments, + allowed_mcp_servers=allowed_mcp_servers, + start_time=start_time, + user_api_key_auth=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, + client_ip=client_ip, + **kwargs, + ) + except Exception as e: + await fire_mcp_tool_call_failure_logging(litellm_logging_obj, e, start_time, user_api_key_auth, kwargs) + raise + + if litellm_logging_obj: + response = await _fire_mcp_tool_call_logging( + logging_obj=litellm_logging_obj, + result=response, + start_time=start_time, + end_time=datetime.now(), + user_api_key_auth=user_api_key_auth, + request_data=kwargs, + ) + return response + + +async def mcp_get_prompt( + name: str, + arguments: dict[str, str] | None = None, + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[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, + client_ip: str | None = None, +) -> GetPromptResult: + """ + Fetch a specific MCP prompt, handling both prefixed and unprefixed names. + """ + allowed_mcp_servers: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_servers=mcp_servers, + client_ip=client_ip, + ) + + if not allowed_mcp_servers: + raise HTTPException( + status_code=403, + detail="User not allowed to get this prompt.", + ) + + # Extract server name from prefixed prompt name + original_prompt_name, server_name = split_server_prefix_from_name(name) + + server: Final = next((s for s in allowed_mcp_servers if s.name == server_name), None) + if server is None: + raise HTTPException( + status_code=403, + detail="User not allowed to get this prompt.", + ) + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=mcp_server_auth_headers, + mcp_auth_header=mcp_auth_header, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + ) + + return await global_mcp_server_manager.get_prompt_from_server( + server=server, + user_api_key_auth=user_api_key_auth, + prompt_name=original_prompt_name, + arguments=arguments, + mcp_auth_header=server_auth_header, + extra_headers=extra_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + + +async def mcp_read_resource( + url: AnyUrl, + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: list[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, + client_ip: str | None = None, +) -> ReadResourceResult: + """Read resource contents from upstream MCP servers.""" + + allowed_mcp_servers: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_servers=mcp_servers, + client_ip=client_ip, + ) + + if not allowed_mcp_servers: + raise HTTPException( + status_code=403, + detail="User not allowed to read this resource.", + ) + + if len(allowed_mcp_servers) != 1: + raise HTTPException( + status_code=400, + detail=("Multiple MCP servers configured; read_resource currently supports exactly one allowed server."), + ) + + server: Final = allowed_mcp_servers[0] + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=mcp_server_auth_headers, + mcp_auth_header=mcp_auth_header, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + ) + + return await global_mcp_server_manager.read_resource_from_server( + server=server, + user_api_key_auth=user_api_key_auth, + url=url, + mcp_auth_header=server_auth_header, + extra_headers=extra_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + + +def _get_standard_logging_mcp_tool_call( + name: str, + arguments: dict[str, object], + server_name: str | None, + session_id: str | None = None, +) -> StandardLoggingMCPToolCall: + mcp_server: Final = global_mcp_server_manager._get_mcp_server_from_tool_name( + add_server_prefix_to_name(name, server_name) if server_name else name + ) + namespaced_tool_name: Final = f"{server_name}/{name}" if server_name else name + if mcp_server: + mcp_info: Final = mcp_server.mcp_info or {} + return StandardLoggingMCPToolCall( + name=name, + arguments=arguments, + mcp_server_name=mcp_info.get("server_name"), + mcp_server_logo_url=mcp_info.get("logo_url"), + namespaced_tool_name=namespaced_tool_name, + mcp_session_id=session_id, + mcp_auth_mode=mcp_server.auth_type, + mcp_server_resource=_redact_mcp_resource_url(mcp_server.url), + ) + else: + return StandardLoggingMCPToolCall( + name=name, + arguments=arguments, + namespaced_tool_name=namespaced_tool_name, + mcp_session_id=session_id, + ) + + +async def _handle_managed_mcp_tool( + server_name: str, + name: str, + arguments: dict[str, object], + user_api_key_auth: UserAPIKeyAuth | None = 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_logging_obj: LiteLLMLoggingObj | None = None, + host_progress_callback: ProgressCallback | None = None, + guardrail_context: Mapping[str, object] | None = None, + client_ip: str | None = None, +) -> CallToolResult: + """Handle tool execution for managed server tools""" + # Import here to avoid circular import + from litellm.proxy.proxy_server import proxy_logging_obj + + call_tool_result: Final = await global_mcp_server_manager.call_tool( + server_name=server_name, + name=name, + arguments=arguments, + user_api_key_auth=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, + client_ip=client_ip, + proxy_logging_obj=proxy_logging_obj, + host_progress_callback=host_progress_callback, + litellm_logging_obj=litellm_logging_obj, + guardrail_context=guardrail_context, + ) + verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) + return call_tool_result + + +async def _handle_local_mcp_tool(name: str, arguments: dict[str, object]) -> CallToolResult: + """Execute a local-registry tool and report whether it succeeded. + + Returns the result rather than bare content because the verdict is part of it: the content + alone cannot say whether the handler failed, so callers used to stamp is_error=False on every + outcome and an upstream rejection was served as tool output. + + A failure is reported as ``is_error=True`` here rather than raised, because the REST surface + turns an unrecognized exception into a 500 and an upstream 403 or 429 is not a gateway crash. + ``MCPUpstreamAuthError`` is the exception: it propagates so the caller is told to + re-authenticate, which both renderers already know how to say. + + Note: Local tools don't use prefixes, so we use the original name + """ + import inspect + + tool: Final = global_mcp_tool_registry.get_tool(name) + if not tool: + raise HTTPException(status_code=404, detail=f"Tool '{name}' not found") + + try: + if inspect.iscoroutinefunction(tool.handler): + result = await tool.handler(**arguments) + else: + result = tool.handler(**arguments) + except MCPUpstreamAuthError: + raise + except Exception as e: + verbose_logger.exception("Error executing local tool %s: %s", name, e) + return CallToolResult( + content=[TextContent(text=f"Error: {e}", type="text")], # mutable-ok: MCP result content + is_error=True, + ) + return CallToolResult( + content=[TextContent(text=str(result), type="text")], # mutable-ok: MCP result content + is_error=False, + ) + + +_MCP_CREDENTIAL_REQUEST_FIELDS: Final = frozenset( + { + "raw_headers", + "mcp_auth_header", + "mcp_server_auth_headers", + "oauth2_headers", + "user_api_key_auth", + } +) + + +class _McpDeniedDetail(TypedDict): + error: ReadOnly[str] + + +async def _execute_handle_list_tools( + context: OperationContext, params: PaginatedRequestParams, host_progress_callback: ProgressCallback | None = None +) -> ListToolsResult: + try: + ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) = context.legacy_auth() + verbose_logger.debug("MCP list_tools - User API Key Auth from context: %s", user_api_key_auth) + verbose_logger.debug("MCP list_tools - MCP servers from context: %s", mcp_servers) + verbose_logger.debug( + "MCP list_tools - MCP server auth headers: %s", + list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, + ) + from mcp.types import Tool + + from litellm.proxy._experimental.mcp_server.tool_search import ( + get_mcp_proxy_tool_definitions, + get_virtual_tool_definitions, + ) + + if context.mcp_proxy_mode: + return ListToolsResult(tools=[Tool.model_validate(d) for d in get_mcp_proxy_tool_definitions()]) + if getattr( + getattr(user_api_key_auth, "object_permission", None), + "mcp_tool_search_enabled", + False, + ): + return ListToolsResult(tools=[Tool.model_validate(d) for d in get_virtual_tool_definitions()]) + + # Get mcp_servers from context variable + verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") + listing: Final = await _list_mcp_tools( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + log_list_tools_to_spendlogs=True, + list_tools_log_source="mcp_protocol", + client_ip=_client_ip, + ) + verbose_logger.info("MCP list_tools - Successfully returned %s tools", len(listing.tools)) + if not listing.outcomes: + return ListToolsResult(tools=listing.tools) + outcome_meta: Final = { + SERVER_OUTCOMES_META_KEY: {key: outcome_wire_value(outcome) for key, outcome in listing.outcomes.items()} + } + return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta}) + except HTTPException as e: + from mcp.shared.exceptions import MCPError + from mcp.types import INVALID_REQUEST + + raise MCPError(code=INVALID_REQUEST, message=_http_detail_message(e.detail)) from e + except Exception as e: + verbose_logger.exception("Error in list_tools endpoint: %s", e) + # Return empty list instead of failing completely + # This prevents the HTTP stream from failing and allows the client to get a response + return ListToolsResult(tools=[]) # mutable-ok: MCP result payload + + +async def _execute_mcp_server_tool_call( + context: OperationContext, params: CallToolRequestParams, host_progress_callback: ProgressCallback | None = None +) -> CallToolResult: + from mcp.types import CallToolResult + + from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + from litellm.proxy.proxy_server import proxy_config + + ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) = context.legacy_auth() + verbose_logger.debug( + "MCP mcp_server_tool_call - user_api_key_auth=%s, user_role=%s", + user_api_key_auth, + getattr(user_api_key_auth, "user_role", "N/A"), + ) + + verbose_logger.debug("MCP mcp_server_tool_call - User API Key Auth from context: %s", user_api_key_auth) + + try: + # Inside this try so virtual-tool errors convert to isError + # CallToolResult instead of raising out of the protocol handler. + virtual_tool_result: Final = await _dispatch_virtual_mcp_tool( + name=params.name, + arguments=params.arguments, + user_api_key_auth=user_api_key_auth, + client_ip=_client_ip, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + mcp_proxy_mode=context.mcp_proxy_mode, + ) + if virtual_tool_result is not None: + return virtual_tool_result + + # Create a body date for logging + body_data: Final = {"name": params.name, "arguments": params.arguments} # mutable-ok: logging payload + # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A) + chain_id: Final = get_chain_id_from_headers(raw_headers) + if chain_id: + body_data["litellm_trace_id"] = chain_id + body_data["litellm_session_id"] = chain_id + + request: Final = build_synthetic_mcp_request( + path="/mcp/tools/call", + raw_headers=raw_headers, + client_ip=_client_ip, + ) + if user_api_key_auth is not None: + data = await add_litellm_data_to_request( + data=body_data, + request=request, + # Bill a team-derived call to the team that granted it. A keyless admitted + # subject carries no team_id, so spend skipped team updates entirely and + # charged the user's PRIMARY org — the granting team's budget never + # accumulated (so it could never begin to block) and, cross-org, the wrong + # organization was charged. This is the ACCOUNTING half; the enforcement + # half (an already-over-budget team stops granting) lives in the source gate. + # Authorization is unaffected: it ran before this, and the union is resolved + # from the untouched auth object passed to call_mcp_tool below. + user_api_key_dict=await MCPRequestHandler.billing_auth_for_tool_call( + user_api_key_auth, tool_name=params.name + ), + proxy_config=proxy_config, + ) + else: + data = body_data + + response: Final = await call_mcp_tool( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=_client_ip, + host_progress_callback=host_progress_callback, + **data, # for logging + ) + except MCPMissingUserEnvVarsError as e: + verbose_logger.info( + "MCP mcp_server_tool_call missing per-user env vars: server_id=%s missing=%s", + e.server_id, + e.missing, + ) + return CallToolResult( + content=[TextContent(text=str(e), type="text")], + is_error=True, + ) + except BlockedPiiEntityError as e: + verbose_logger.error("BlockedPiiEntityError in MCP tool call: %s", e) + return CallToolResult( + content=[ + TextContent( + text=f"Error: Blocked PII entity detected - {e}", + type="text", + ) + ], + is_error=True, + ) + except GuardrailRaisedException as e: + verbose_logger.error("GuardrailRaisedException in MCP tool call: %s", e) + return CallToolResult( + content=[TextContent(text=f"Error: Guardrail violation - {e}", type="text")], + is_error=True, + ) + except HTTPException as e: + verbose_logger.error("HTTPException in MCP tool call: %s", e) + return CallToolResult( + content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")], + is_error=True, + ) + except MCPUpstreamAuthError as e: + # The MCP session manager serializes handler exceptions as JSON-RPC errors, so a + # mid-session tool call cannot emit a raw 401 + WWW-Authenticate the way the REST + # call path and the connect-time preemptive check do. Return an explicit isError + # naming the upstream status (at info level, not a traceback) so the client still + # learns it must re-authenticate upstream and expected pass-through 401s don't spam. + verbose_logger.info("Upstream auth failure calling MCP tool: HTTP %s", e.status_code) + return CallToolResult( + content=[ + TextContent( + text=f"Error: upstream authentication required (HTTP {e.status_code})", + type="text", + ) + ], + is_error=True, + ) + except Exception as e: + verbose_logger.exception("MCP mcp_server_tool_call - error: %s", e) + return CallToolResult( + content=[TextContent(text=f"Error: {e}", type="text")], + is_error=True, + ) + + return response + + +async def _execute_list_prompts( + context: OperationContext, params: PaginatedRequestParams, host_progress_callback: ProgressCallback | None = None +) -> ListPromptsResult: + if context.mcp_proxy_mode: + _reject_mcp_proxy_operation() + try: + ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) = context.legacy_auth() + verbose_logger.debug("MCP list_prompts - User API Key Auth from context: %s", user_api_key_auth) + verbose_logger.debug("MCP list_prompts - MCP servers from context: %s", mcp_servers) + verbose_logger.debug( + "MCP list_prompts - MCP server auth headers: %s", + list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, + ) + # Get mcp_servers from context variable + verbose_logger.debug("MCP list_prompts - Calling _list_prompts") + prompts: Final = await _list_mcp_prompts( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=_client_ip, + ) + verbose_logger.info("MCP list_prompts - Successfully returned %s prompts", len(prompts)) + return ListPromptsResult(prompts=prompts) + except Exception as e: + verbose_logger.exception("Error in list_prompts endpoint: %s", e) + # Return empty list instead of failing completely + # This prevents the HTTP stream from failing and allows the client to get a response + return ListPromptsResult(prompts=[]) # mutable-ok: MCP result payload + + +async def _execute_get_prompt( + context: OperationContext, params: GetPromptRequestParams, host_progress_callback: ProgressCallback | None = None +) -> GetPromptResult: + if context.mcp_proxy_mode: + _reject_mcp_proxy_operation() + ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) = context.legacy_auth() + + verbose_logger.debug("MCP mcp_server_tool_call - User API Key Auth from context: %s", user_api_key_auth) + return await mcp_get_prompt( + name=params.name, + arguments=params.arguments, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=_client_ip, + ) + + +async def _execute_list_resources( + context: OperationContext, params: PaginatedRequestParams, host_progress_callback: ProgressCallback | None = None +) -> ListResourcesResult: + if context.mcp_proxy_mode: + _reject_mcp_proxy_operation() + try: + ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) = context.legacy_auth() + verbose_logger.debug("MCP list_resources - User API Key Auth from context: %s", user_api_key_auth) + verbose_logger.debug("MCP list_resources - MCP servers from context: %s", mcp_servers) + verbose_logger.debug( + "MCP list_resources - MCP server auth headers: %s", + list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, + ) + + resources: Final = await _list_mcp_resources( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=_client_ip, + ) + verbose_logger.info("MCP list_resources - Successfully returned %s resources", len(resources)) + return ListResourcesResult(resources=resources) + except Exception as e: + verbose_logger.exception("Error in list_resources endpoint: %s", e) + return ListResourcesResult(resources=[]) # mutable-ok: MCP result payload + + +async def _execute_list_resource_templates( + context: OperationContext, params: PaginatedRequestParams, host_progress_callback: ProgressCallback | None = None +) -> ListResourceTemplatesResult: + if context.mcp_proxy_mode: + _reject_mcp_proxy_operation() + try: + ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) = context.legacy_auth() + verbose_logger.debug("MCP list_resource_templates - User API Key Auth from context: %s", user_api_key_auth) + verbose_logger.debug("MCP list_resource_templates - MCP servers from context: %s", mcp_servers) + verbose_logger.debug( + "MCP list_resource_templates - MCP server auth headers: %s", + list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, + ) + + resource_templates: Final = await _list_mcp_resource_templates( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=_client_ip, + ) + verbose_logger.info( + "MCP list_resource_templates - Successfully returned %s resource templates", len(resource_templates) + ) + return ListResourceTemplatesResult(resource_templates=resource_templates) + except Exception as e: + verbose_logger.exception("Error in list_resource_templates endpoint: %s", e) + return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload + + +async def _execute_read_resource( + context: OperationContext, params: ReadResourceRequestParams, host_progress_callback: ProgressCallback | None = None +) -> ReadResourceResult: + if context.mcp_proxy_mode: + _reject_mcp_proxy_operation() + ( + user_api_key_auth, + mcp_auth_header, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + _client_ip, + ) = context.legacy_auth() + + read_resource_result: Final = await mcp_read_resource( + url=params.uri, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=_client_ip, + ) + + return read_resource_result + + +def _reject_mcp_proxy_operation() -> NoReturn: + from mcp.shared.exceptions import MCPError + from mcp.types import METHOD_NOT_FOUND + + raise MCPError(code=METHOD_NOT_FOUND, message="Operation unavailable on /mcp/proxy") + + +def prepare_context( + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_auth_header: str | None = None, + mcp_servers: Sequence[str] | None = None, + mcp_server_auth_headers: Mapping[str, Mapping[str, str]] | None = None, + oauth2_headers: Mapping[str, str] | None = None, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, + mcp_proxy_mode: bool = False, +) -> OperationContext: + return OperationContext( + _caller=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=tuple(mcp_servers) if mcp_servers is not None else None, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + mcp_proxy_mode=mcp_proxy_mode, + ) + + +GatewayOperation: TypeAlias = ( + AuthorizedToolCall + | ListToolsRequest + | CallToolRequest + | ListPromptsRequest + | GetPromptRequest + | ListResourcesRequest + | ListResourceTemplatesRequest + | ReadResourceRequest +) +GatewayResult: TypeAlias = ( + ListToolsResult + | CallToolResult + | ListPromptsResult + | GetPromptResult + | ListResourcesResult + | ListResourceTemplatesResult + | ReadResourceResult +) + + +class GatewayOperations: + def __init__(self, host_progress_callback: ProgressCallback | None = None) -> None: + self._host_progress_callback = host_progress_callback + + @overload + async def execute(self, operation: AuthorizedToolCall, context: OperationContext) -> CallToolResult: ... + + @overload + async def execute(self, operation: ListToolsRequest, context: OperationContext) -> ListToolsResult: ... + + @overload + async def execute(self, operation: CallToolRequest, context: OperationContext) -> CallToolResult: ... + + @overload + async def execute(self, operation: ListPromptsRequest, context: OperationContext) -> ListPromptsResult: ... + + @overload + async def execute(self, operation: GetPromptRequest, context: OperationContext) -> GetPromptResult: ... + + @overload + async def execute(self, operation: ListResourcesRequest, context: OperationContext) -> ListResourcesResult: ... + + @overload + async def execute( + self, operation: ListResourceTemplatesRequest, context: OperationContext + ) -> ListResourceTemplatesResult: ... + + @overload + async def execute(self, operation: ReadResourceRequest, context: OperationContext) -> ReadResourceResult: ... + + async def execute(self, operation: GatewayOperation, context: OperationContext) -> GatewayResult: + match operation: + case AuthorizedToolCall(): + auth, token, _servers, server_headers, oauth_headers, headers, _client_ip = context.legacy_auth() + return await _execute_mcp_tool( + name=operation.name, + arguments=dict(operation.arguments), # mutable-ok: existing tool hooks own mutable argument data + allowed_mcp_servers=list( + operation.allowed_mcp_servers + ), # mutable-ok: legacy dispatch list contract + start_time=operation.start_time, + user_api_key_auth=auth, + mcp_auth_header=token, + mcp_server_auth_headers=server_headers, + oauth2_headers=oauth_headers, + raw_headers=headers, + client_ip=_client_ip, + host_progress_callback=operation.host_progress_callback, + guardrail_context=operation.guardrail_context, + **operation.logging_data, + ) + case ListToolsRequest(params=params): + return await _execute_handle_list_tools( + context, params or PaginatedRequestParams(), self._host_progress_callback + ) + case CallToolRequest(params=params): + return await _execute_mcp_server_tool_call(context, params, self._host_progress_callback) + case ListPromptsRequest(params=params): + return await _execute_list_prompts( + context, params or PaginatedRequestParams(), self._host_progress_callback + ) + case GetPromptRequest(params=params): + return await _execute_get_prompt(context, params, self._host_progress_callback) + case ListResourcesRequest(params=params): + return await _execute_list_resources( + context, params or PaginatedRequestParams(), self._host_progress_callback + ) + case ListResourceTemplatesRequest(params=params): + return await _execute_list_resource_templates( + context, params or PaginatedRequestParams(), self._host_progress_callback + ) + case ReadResourceRequest(params=params): + return await _execute_read_resource(context, params, self._host_progress_callback) + case _: + return assert_never(operation) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 15f97a15b73..c2f7bf7d531 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -203,17 +203,19 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_request_base_url, ) - from litellm.proxy._experimental.mcp_server.server import ( + from litellm.proxy._experimental.mcp_server.operations import ( ListMCPToolsRestAPIResponseObject, MCPInfo, MCPServer, - _aggregate_server_key, # pyright: ignore[reportPrivateUsage] # same per-server key as the tools/list _meta outcomes - _apply_toolset_scope, + _aggregate_server_key, _fire_mcp_tool_call_logging, execute_mcp_tool, filter_tools_by_allowed_tools, filter_tools_by_key_team_permissions, fire_mcp_tool_call_failure_logging, + ) + from litellm.proxy._experimental.mcp_server.server import ( + _apply_toolset_scope, reject_disallowed_mcp_client, ) @@ -670,6 +672,7 @@ if MCP_AVAILABLE: user_api_key_auth: UserAPIKeyAuth | None = None, extra_headers: dict[str, str] | None = None, apply_tool_filters: bool = True, + client_ip: str | None = None, ): """Helper function to get tools for a single server. @@ -684,6 +687,7 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=False, raw_headers=raw_headers, + client_ip=client_ip, user_api_key_auth=user_api_key_auth, ) @@ -797,6 +801,7 @@ if MCP_AVAILABLE: user_api_key_dict, extra_headers=user_oauth_extra_headers, apply_tool_filters=apply_tool_filters, + client_ip=rest_client_ip, ) except MCPUpstreamAuthError: # Surface the upstream 401/403 to the caller so it can emit the @@ -1016,6 +1021,7 @@ if MCP_AVAILABLE: user_api_key_dict, extra_headers=user_oauth_extra_headers, apply_tool_filters=apply_tool_filters, + client_ip=_rest_client_ip, ) except Exception as e: verbose_logger.warning( @@ -1193,6 +1199,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers=data.get("mcp_server_auth_headers"), oauth2_headers=user_oauth_extra_headers or data.get("oauth2_headers"), raw_headers=data.get("raw_headers"), + client_ip=IPAddressUtils.get_mcp_client_ip(request), litellm_logging_obj=data.get("litellm_logging_obj"), guardrail_context=MCPRequestContext.resolve_guardrail_context(data), requested_server_id=canonical_server_id, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 397a82cfa45..433b693fcae 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -11,28 +11,22 @@ import hashlib import json import os import time -import traceback import types -import uuid from collections import Counter -from collections.abc import AsyncIterator, Callable, Iterable, Mapping, Sequence -from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, NoReturn, Protocol import httpx from fastapi import FastAPI, HTTPException -from pydantic import AnyUrl, ConfigDict, Field, TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError from starlette.requests import Request as StarletteRequest from starlette.responses import JSONResponse from starlette.types import Message, Receive, Scope, Send -from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.constants import ( - MAXIMUM_TRACEBACK_LINES_TO_LOG, MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH, ) -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -41,12 +35,6 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, _is_mcp_admitted_user_subject, ) -from litellm.proxy._experimental.mcp_server.byok_credential_cache import ( - byok_credential_cache, - byok_credential_cache_key, - cache_byok_credential, - get_cached_byok_credential, -) from litellm.proxy._experimental.mcp_server.client_allowlist import ( MCPClientAllowlist, check_mcp_client_allowed, @@ -56,7 +44,6 @@ from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, ) from litellm.proxy._experimental.mcp_server.exceptions import ( - MCPToolResultError, MCPUpstreamAuthError, ) from litellm.proxy._experimental.mcp_server.mcp_context import ( @@ -74,7 +61,6 @@ from litellm.proxy._experimental.mcp_server.mcp_debug import ( ) from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, - get_byok_www_authenticate, get_passthrough_www_authenticate, get_route_relative_request_path, well_known_root_suffix, @@ -84,14 +70,6 @@ from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, LITELLM_MCP_SERVER_VERSION, - MCPMissingUserEnvVarsError, - add_server_prefix_to_name, - build_synthetic_mcp_request, - extract_mcp_tool_result_error_message, - get_server_prefix, - iter_known_server_prefixes, - logging_safe_mcp_headers, - match_known_tool_name, ) from litellm.proxy._types import ( ProxyException, @@ -99,13 +77,6 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils -from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( - publish_auth_cache_invalidation, -) -from litellm.proxy.litellm_pre_call_utils import ( - LiteLLMProxyRequestSetup, - get_chain_id_from_headers, -) from litellm.types.mcp import ( MCPAuth, MCPGatewaySession, @@ -114,14 +85,11 @@ from litellm.types.mcp import ( MCPGatewaySessionsTerminateResponse, MCPSpecVersion, ) -from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer -from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall -from litellm.utils import Rules, client, function_setup +from litellm.types.mcp_server.mcp_server_manager import MCPServer if TYPE_CHECKING: from mcp.server.session import ServerSession as _McpServerSession - from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS: Final = 30 * 60 # Upper bound on concurrent stateful sessions a single caller may hold. Each @@ -159,13 +127,6 @@ def unsupported_protocol_version(scope: Scope) -> str | None: return None -async def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: - """Drop a stored-or-deleted BYOK credential from this worker's cache and from every peer worker's.""" - cache_key: Final = byok_credential_cache_key(user_id, server_id) - byok_credential_cache.delete_cache(cache_key) - await publish_auth_cache_invalidation(cache_key=cache_key) - - # Check if MCP is available # "mcp" requires python 3.10 or higher, but several litellm users use python 3.8 # We're making this conditional import to avoid breaking users who use python 3.8. @@ -210,19 +171,6 @@ _SESSION_MANAGERS_INITIALIZED = False _INITIALIZATION_LOCK: Final = asyncio.Lock() -def _mcp_session_id_from_headers( - raw_headers: dict[str, str] | None, -) -> str | None: - """The ``mcp-session-id`` of a stateful MCP session, read case-insensitively - from the request headers. ``None`` for stateless calls (no such header).""" - if not raw_headers: - return None - for key, value in raw_headers.items(): - if isinstance(key, str) and key.lower() == "mcp-session-id": - return value or None - return None - - def _jsonrpc_text_has_top_level_method(text: str) -> bool: """Whether a (possibly truncated) JSON-RPC envelope has a ``method`` key at the root object's top level. @@ -466,6 +414,59 @@ def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: if MCP_AVAILABLE: + __all__ = ( + "_MCP_CREDENTIAL_REQUEST_FIELDS", + "BlobResourceContents", + "ListMCPToolsRestAPIResponseObject", + "ResourceTemplate", + "TextResourceContents", + "_McpDeniedDetail", + "_aggregate_server_key", + "_build_virtual_call_logging_obj", + "_check_byok_credential", + "_client_has_passthrough_authorization", + "_client_has_per_server_auth_header", + "_dispatch_virtual_mcp_tool", + "_fire_mcp_tool_call_logging", + "_get_allowed_mcp_servers", + "_get_allowed_mcp_servers_from_mcp_server_names", + "_get_byok_credential", + "_get_prompts_from_mcp_servers", + "_get_resource_templates_from_mcp_servers", + "_get_resources_from_mcp_servers", + "_get_standard_logging_mcp_tool_call", + "_get_tools_from_mcp_servers", + "_get_user_oauth_extra_headers_from_db", + "_handle_local_mcp_tool", + "_handle_managed_mcp_tool", + "_http_detail_message", + "_invalidate_byok_cred_cache", + "_list_mcp_prompts", + "_list_mcp_resource_templates", + "_list_mcp_resources", + "_list_mcp_tools", + "_list_tools_before_first_call", + "_mcp_session_id_from_headers", + "_merge_gateway_initialize_instructions", + "_prefetch_oauth_creds_for_user", + "_prepare_mcp_server_headers", + "_raise_if_initialize_grants_no_mcp_servers", + "_redact_mcp_resource_url", + "_resolve_display_name_to_original", + "_run_post_mcp_call_guardrails", + "_server_answers_to", + "_tool_name_matches", + "apply_tool_overrides", + "call_mcp_tool", + "execute_mcp_tool", + "filter_tools_by_allowed_tools", + "filter_tools_by_key_team_permissions", + "fire_mcp_tool_call_failure_logging", + "global_mcp_server_manager", + "mcp_get_prompt", + "mcp_read_resource", + "raise_denied_scoped_mcp_access", + ) from mcp.server import Server # Import auth context variables and middleware @@ -476,6 +477,23 @@ if MCP_AVAILABLE: from mcp.server.context import ServerRequestContext from mcp.server.lowlevel.server import NotificationOptions from mcp.server.models import InitializationOptions + from mcp.shared.exceptions import MCPError + from mcp.types import ( + CallToolRequest, + GetPromptRequest, + ListPromptsRequest, + ListResourcesRequest, + ListResourceTemplatesRequest, + ListToolsRequest, + ReadResourceRequest, + ) + + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.contracts import OperationContext + from litellm.proxy._experimental.mcp_server.operations import ( + _invalidate_byok_cred_cache, + _mcp_session_id_from_headers, + ) try: from mcp.server.streamable_http_manager import StreamableHTTPSessionManager @@ -493,62 +511,27 @@ if MCP_AVAILABLE: ListResourceTemplatesResult, ListToolsResult, PaginatedRequestParams, - Prompt, ReadResourceRequestParams, - TextContent, ) - from mcp.types import Tool as MCPTool from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( MCPAuthenticatedUser, ) - from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( - SERVER_OUTCOMES_META_KEY, - AggregateToolListing, - ServerListOk, - ServerOutcome, - classify_list_exception, - outcome_wire_value, - ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, - _caller_authorization_fans_out, - _client_forwarded_authorization_headers, - _resolve_openapi_tool_auth, - _should_strip_caller_authorization, global_mcp_server_manager, ) - from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( - _request_auth_header, - _request_extra_headers, - _request_resolved_auth_headers, - ) - from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport - from litellm.proxy._experimental.mcp_server.tool_registry import ( - global_mcp_tool_registry, - ) - from litellm.proxy._experimental.mcp_server.utils import ( - MCP_TOOL_PREFIX_SEPARATOR, - is_tool_name_prefixed, - normalize_server_name, - split_server_prefix_from_name, - strip_known_server_prefix, - ) - from litellm.types.mcp import DEFAULT_CREDENTIAL_HEADER, without_header ###################################################### ############ MCP Tools List REST API Response Object # # Defined here because we don't want to add `mcp` as a # required dependency for `litellm` pip package ###################################################### - class ListMCPToolsRestAPIResponseObject(MCPTool): - """ - Object returned by the /tools/list REST API route. - """ - - mcp_info: MCPInfo | None = Field(default=None, alias="mcp_info") - model_config = ConfigDict(arbitrary_types_allowed=True) + from litellm.proxy._experimental.mcp_server.operations import ( + ListMCPToolsRestAPIResponseObject, + ) + from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport def _gateway_create_initialization_options( self, @@ -818,94 +801,45 @@ if MCP_AVAILABLE: ############### MCP Server Routes ####################### ######################################################## - async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListToolsResult: - """ - List all available tools, with each server's listing outcome attached to the result's - ``_meta`` (SERVER_OUTCOMES_META_KEY) so a broken upstream is distinguishable from a healthy - server with no tools. Returning a ListToolsResult (rather than a bare list) makes the MCP SDK - pass the result through unwrapped, which is what lets the ``_meta`` survive to the client. - Also captures the active session for propagation to callbacks. - """ - req_ctx: Final = ctx - _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) - _session_reset_token: Final = active_mcp_session_var.set(ctx.session) - _trace_token = None - _transport_token = None - _destinations_token = None - - try: - _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) - _transport_token = _otel_set_mcp_transport_span(_otel_transport_span_from_message(req_ctx)) - _destinations_token = _otel_set_mcp_request_destinations(req_ctx) - # Get user authentication from context variable + @contextlib.asynccontextmanager + async def _legacy_operation_context(ctx: ServerRequestContext, *, trace: bool) -> AsyncGenerator[OperationContext]: + with contextlib.ExitStack() as cleanup: + cleanup.callback(active_mcp_request_ctx_var.reset, active_mcp_request_ctx_var.set(ctx)) + cleanup.callback(active_mcp_session_var.reset, active_mcp_session_var.set(ctx.session)) + if trace: + cleanup.callback( + _otel_reset_mcp_trace_carrier, _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(ctx)) + ) + cleanup.callback( + _otel_reset_mcp_transport_span, _otel_set_mcp_transport_span(_otel_transport_span_from_message(ctx)) + ) + cleanup.callback(_otel_reset_mcp_request_destinations, _otel_set_mcp_request_destinations(ctx)) ( - user_api_key_auth, - mcp_auth_header, - mcp_servers, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - _client_ip, + auth, + token, + servers, + server_headers, + oauth_headers, + headers, + client_ip, ) = await get_or_extract_auth_context() - verbose_logger.debug("MCP list_tools - User API Key Auth from context: %s", user_api_key_auth) - verbose_logger.debug("MCP list_tools - MCP servers from context: %s", mcp_servers) - verbose_logger.debug( - "MCP list_tools - MCP server auth headers: %s", - list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, - ) - from mcp.types import Tool - - from litellm.proxy._experimental.mcp_server.tool_search import ( - get_mcp_proxy_tool_definitions, - get_virtual_tool_definitions, + yield operations.prepare_context( + auth, token, servers, server_headers, oauth_headers, headers, client_ip, _mcp_proxy_mode.get() ) - if _mcp_proxy_mode.get(): - return ListToolsResult(tools=[Tool.model_validate(d) for d in get_mcp_proxy_tool_definitions()]) - if getattr( - getattr(user_api_key_auth, "object_permission", None), - "mcp_tool_search_enabled", - False, - ): - return ListToolsResult(tools=[Tool.model_validate(d) for d in get_virtual_tool_definitions()]) - - # Get mcp_servers from context variable - verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") - listing: Final = await _list_mcp_tools( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - log_list_tools_to_spendlogs=True, - list_tools_log_source="mcp_protocol", - ) - verbose_logger.info("MCP list_tools - Successfully returned %s tools", len(listing.tools)) - if not listing.outcomes: - return ListToolsResult(tools=listing.tools) - outcome_meta: Final = { - SERVER_OUTCOMES_META_KEY: { - key: outcome_wire_value(outcome) for key, outcome in listing.outcomes.items() - } - } - return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta}) - except HTTPException as e: - from mcp.shared.exceptions import MCPError - from mcp.types import INVALID_REQUEST - - raise MCPError(code=INVALID_REQUEST, message=_http_detail_message(e.detail)) from e - except Exception as e: - verbose_logger.exception("Error in list_tools endpoint: %s", e) - # Return empty list instead of failing completely - # This prevents the HTTP stream from failing and allows the client to get a response - return ListToolsResult(tools=[]) # mutable-ok: MCP result payload - finally: - _otel_reset_mcp_request_destinations(_destinations_token) - _otel_reset_mcp_transport_span(_transport_token) - _otel_reset_mcp_trace_carrier(_trace_token) - active_mcp_session_var.reset(_session_reset_token) - active_mcp_request_ctx_var.reset(_ctx_reset_token) + async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListToolsResult: + try: + async with _legacy_operation_context(ctx, trace=True) as context: + return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute( + ListToolsRequest(params=params), context + ) + except MCPError: + raise + except HTTPException as exc: + raise MCPError(code=INVALID_REQUEST, message=operations._http_detail_message(exc.detail)) from exc + except Exception as exc: # noqa: BLE001 # preserve native listing fallback for ingress failures + verbose_logger.exception("Error in list_tools endpoint: %s", exc) + return ListToolsResult(tools=[]) def _capture_host_progress_callback(ctx: ServerRequestContext) -> Callable | None: """Return a progress-forwarding callback bound to the host MCP session. @@ -942,581 +876,71 @@ if MCP_AVAILABLE: raise MCPError(code=METHOD_NOT_FOUND, message="Operation unavailable on /mcp/proxy") - async def _build_virtual_call_logging_obj( - name: str, - arguments: dict[str, object], - user_api_key_auth: UserAPIKeyAuth, - raw_headers: Mapping[str, str] | None = None, - client_ip: str | None = None, - ) -> LiteLLMLoggingObj | None: - """Run the pre-call pipeline (guardrails + logging setup) for a virtual - mcp_tool_call so the SSE path spend-logs like the REST path.""" - from litellm.proxy.common_request_processing import ( - ProxyBaseLLMRequestProcessing, - ) - from litellm.proxy.proxy_server import ( - general_settings, - proxy_config, - proxy_logging_obj, - ) - - request: Final = build_synthetic_mcp_request( - path="/mcp/tools/call", - raw_headers=raw_headers, - client_ip=client_ip, - ) - _, virtual_logging_obj = await ProxyBaseLLMRequestProcessing( - data={"name": name, "arguments": arguments} - ).common_processing_pre_call_logic( - request=request, - user_api_key_dict=user_api_key_auth, - proxy_config=proxy_config, - route_type=CallTypes.call_mcp_tool.value, - proxy_logging_obj=proxy_logging_obj, - general_settings=general_settings, - ) - return virtual_logging_obj - - async def _dispatch_virtual_mcp_tool( - name: str, - arguments: dict[str, object] | None, - user_api_key_auth: UserAPIKeyAuth | None, - client_ip: str | None, - mcp_servers: list[str] | None = 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, - ) -> CallToolResult | None: - """Handle the mcp_tool_search / mcp_tool_call virtual tools. - - Returns a CallToolResult when ``name`` is a virtual tool, else ``None`` so - the caller falls through to normal tool routing. - """ - from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K - from litellm.proxy._experimental.mcp_server.tool_search import ( - AGENT_SEARCH_TOOL_NAME, - DEFAULT_AGENT_SEARCH_TOP_K, - MCP_PROXY_CALL_TOOL_NAME, - MCP_PROXY_TOOL_NAMES, - MCP_TOOL_SEARCH_TOOL_NAME, - SKILL_SEARCH_TOOL_NAME, - VIRTUAL_TOOL_NAMES, - coerce_top_k, - handle_agent_search, - handle_mcp_proxy_tool, - handle_mcp_tool_call, - handle_mcp_tool_search, - handle_skill_search, - ) - - if _mcp_proxy_mode.get() and name not in MCP_PROXY_TOOL_NAMES: - return CallToolResult( - content=[ # mutable-ok: MCP result content - TextContent(type="text", text=f"Tool {name} is unavailable on /mcp/proxy") - ], - is_error=True, - ) - - if _mcp_proxy_mode.get() and name in MCP_PROXY_TOOL_NAMES: - assert user_api_key_auth is not None - proxy_call_start: Final = datetime.now() # noqa: DTZ005 # logging pipeline uses naive datetimes - proxy_logging_obj: Final = ( - await _build_virtual_call_logging_obj( - name=name, - arguments=arguments or {}, # mutable-ok: logging pipeline payload - user_api_key_auth=user_api_key_auth, - raw_headers=raw_headers, - client_ip=client_ip, - ) - if name == MCP_PROXY_CALL_TOOL_NAME - else None - ) - try: - proxy_result: Final = await handle_mcp_proxy_tool( - name=name, - arguments=arguments or {}, # mutable-ok: proxy handler payload - user_api_key_dict=user_api_key_auth, - client_ip=client_ip, - mcp_servers=mcp_servers, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - litellm_logging_obj=proxy_logging_obj, - ) - except Exception as exc: - if proxy_logging_obj is not None: - from litellm.proxy.proxy_server import proxy_logging_obj as request_logging_obj - - failure_end: Final = datetime.now() # noqa: DTZ005 # matches the logging pipeline start time - failure_traceback: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) - try: - proxy_logging_obj.failure_handler(exc, failure_traceback, proxy_call_start, failure_end) - await proxy_logging_obj.async_failure_handler( - exc, failure_traceback, proxy_call_start, failure_end - ) - if not isinstance(exc, MCPUpstreamAuthError): - await request_logging_obj.post_call_failure_hook( - request_data={ # mutable-ok: failure hook mutates its request payload - "name": name, - "arguments": arguments, - "litellm_logging_obj": proxy_logging_obj, - }, - original_exception=exc, - user_api_key_dict=user_api_key_auth, - route="/mcp/call_tool", - traceback_str=failure_traceback, - ) - except Exception: # noqa: BLE001 # a failing failure hook must not mask the tool call's own error - verbose_logger.exception("Error logging failed MCP proxy tool call") - raise - if proxy_logging_obj is not None: - return await _fire_mcp_tool_call_logging( - logging_obj=proxy_logging_obj, - result=proxy_result, - start_time=proxy_call_start, - end_time=datetime.now(), # noqa: DTZ005 # matches the logging pipeline start time - user_api_key_auth=user_api_key_auth, - request_data=types.MappingProxyType({"name": name, "arguments": arguments}), - ) - return proxy_result - - if name not in VIRTUAL_TOOL_NAMES: - return None - - if not getattr( - getattr(user_api_key_auth, "object_permission", None), - "mcp_tool_search_enabled", - False, - ): - return CallToolResult( - content=[ - TextContent( - type="text", - text=f"Tool {name} requires mcp_tool_search_enabled on the key", - ) - ], - is_error=True, - ) - - args: Final = arguments or {} - if name == MCP_TOOL_SEARCH_TOOL_NAME: - return await handle_mcp_tool_search( - query=args.get("query", ""), - top_k=coerce_top_k(args.get("top_k", 5)), - user_api_key_dict=user_api_key_auth, - client_ip=client_ip, - mcp_servers=mcp_servers, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) - - assert user_api_key_auth is not None # guaranteed by the flag check above - if name == AGENT_SEARCH_TOOL_NAME: - return await handle_agent_search( - query=str(args.get("query", "")), - top_k=coerce_top_k(args.get("top_k", DEFAULT_AGENT_SEARCH_TOP_K), default=DEFAULT_AGENT_SEARCH_TOP_K), - user_api_key_dict=user_api_key_auth, - ) - if name == SKILL_SEARCH_TOOL_NAME: - return await handle_skill_search( - query=str(args.get("query", "")), - top_k=coerce_top_k(args.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K), - user_api_key_dict=user_api_key_auth, - ) - virtual_logging_obj: Final = await _build_virtual_call_logging_obj( - name=name, - arguments=args, - user_api_key_auth=user_api_key_auth, - raw_headers=raw_headers, - client_ip=client_ip, - ) - return await handle_mcp_tool_call( - tool_name=args.get("tool_name", ""), - arguments=args.get("arguments") or {}, - user_api_key_dict=user_api_key_auth, - client_ip=client_ip, - mcp_servers=mcp_servers, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - litellm_logging_obj=virtual_logging_obj, - ) + from litellm.proxy._experimental.mcp_server.operations import ( + _build_virtual_call_logging_obj, + _dispatch_virtual_mcp_tool, + ) async def mcp_server_tool_call(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult: - """ - Call a specific tool with the provided arguments - Args: - ctx: SDK request context carrying the client session and HTTP request - params (CallToolRequestParams): Tool name and arguments - Returns: - CallToolResult: Tool execution results - """ - from mcp.types import CallToolResult - - from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException - from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request - from litellm.proxy.proxy_server import proxy_config - - req_ctx: Final = ctx - _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) - _session_reset_token: Final = active_mcp_session_var.set(ctx.session) - _trace_token = None - _transport_token = None - _destinations_token = None - - try: - _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) - _transport_token = _otel_set_mcp_transport_span(_otel_transport_span_from_message(req_ctx)) - _destinations_token = _otel_set_mcp_request_destinations(req_ctx) - # Validate arguments - ( - user_api_key_auth, - mcp_auth_header, - mcp_servers, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - _client_ip, - ) = await get_or_extract_auth_context() - verbose_logger.debug( - "MCP mcp_server_tool_call - user_api_key_auth=%s, user_role=%s", - user_api_key_auth, - getattr(user_api_key_auth, "user_role", "N/A"), + async with _legacy_operation_context(ctx, trace=True) as context: + return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute( + CallToolRequest(params=params), context ) - verbose_logger.debug("MCP mcp_server_tool_call - User API Key Auth from context: %s", user_api_key_auth) - - try: - # Inside this try so virtual-tool errors convert to isError - # CallToolResult instead of raising out of the protocol handler. - virtual_tool_result: Final = await _dispatch_virtual_mcp_tool( - name=params.name, - arguments=params.arguments, - user_api_key_auth=user_api_key_auth, - client_ip=_client_ip, - mcp_servers=mcp_servers, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) - if virtual_tool_result is not None: - return virtual_tool_result - - host_progress_callback: Final = _capture_host_progress_callback(ctx) - # Create a body date for logging - body_data: Final = {"name": params.name, "arguments": params.arguments} # mutable-ok: logging payload - # Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A) - chain_id: Final = get_chain_id_from_headers(raw_headers) - if chain_id: - body_data["litellm_trace_id"] = chain_id - body_data["litellm_session_id"] = chain_id - - request: Final = build_synthetic_mcp_request( - path="/mcp/tools/call", - raw_headers=raw_headers, - client_ip=_client_ip, - ) - if user_api_key_auth is not None: - data = await add_litellm_data_to_request( - data=body_data, - request=request, - # Bill a team-derived call to the team that granted it. A keyless admitted - # subject carries no team_id, so spend skipped team updates entirely and - # charged the user's PRIMARY org — the granting team's budget never - # accumulated (so it could never begin to block) and, cross-org, the wrong - # organization was charged. This is the ACCOUNTING half; the enforcement - # half (an already-over-budget team stops granting) lives in the source gate. - # Authorization is unaffected: it ran before this, and the union is resolved - # from the untouched auth object passed to call_mcp_tool below. - user_api_key_dict=await MCPRequestHandler.billing_auth_for_tool_call( - user_api_key_auth, tool_name=params.name - ), - proxy_config=proxy_config, - ) - else: - data = body_data - - response: Final = await call_mcp_tool( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - client_ip=_client_ip, - host_progress_callback=host_progress_callback, - **data, # for logging - ) - except MCPMissingUserEnvVarsError as e: - verbose_logger.info( - "MCP mcp_server_tool_call missing per-user env vars: server_id=%s missing=%s", - e.server_id, - e.missing, - ) - return CallToolResult( - content=[TextContent(text=str(e), type="text")], - is_error=True, - ) - except BlockedPiiEntityError as e: - verbose_logger.error("BlockedPiiEntityError in MCP tool call: %s", e) - return CallToolResult( - content=[ - TextContent( - text=f"Error: Blocked PII entity detected - {e}", - type="text", - ) - ], - is_error=True, - ) - except GuardrailRaisedException as e: - verbose_logger.error("GuardrailRaisedException in MCP tool call: %s", e) - return CallToolResult( - content=[TextContent(text=f"Error: Guardrail violation - {e}", type="text")], - is_error=True, - ) - except HTTPException as e: - verbose_logger.error("HTTPException in MCP tool call: %s", e) - return CallToolResult( - content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")], - is_error=True, - ) - except MCPUpstreamAuthError as e: - # The MCP session manager serializes handler exceptions as JSON-RPC errors, so a - # mid-session tool call cannot emit a raw 401 + WWW-Authenticate the way the REST - # call path and the connect-time preemptive check do. Return an explicit isError - # naming the upstream status (at info level, not a traceback) so the client still - # learns it must re-authenticate upstream and expected pass-through 401s don't spam. - verbose_logger.info("Upstream auth failure calling MCP tool: HTTP %s", e.status_code) - return CallToolResult( - content=[ - TextContent( - text=f"Error: upstream authentication required (HTTP {e.status_code})", - type="text", - ) - ], - is_error=True, - ) - except Exception as e: - verbose_logger.exception("MCP mcp_server_tool_call - error: %s", e) - return CallToolResult( - content=[TextContent(text=f"Error: {e}", type="text")], - is_error=True, - ) - - return response - finally: - _otel_reset_mcp_request_destinations(_destinations_token) - _otel_reset_mcp_transport_span(_transport_token) - _otel_reset_mcp_trace_carrier(_trace_token) - active_mcp_session_var.reset(_session_reset_token) - active_mcp_request_ctx_var.reset(_ctx_reset_token) - async def list_prompts(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListPromptsResult: - """ - List all available prompts - """ if _mcp_proxy_mode.get(): _reject_mcp_proxy_operation() - _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) - _session_reset_token: Final = active_mcp_session_var.set(ctx.session) - try: - # Get user authentication from context variable - ( - user_api_key_auth, - mcp_auth_header, - mcp_servers, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - _client_ip, - ) = await get_or_extract_auth_context() - verbose_logger.debug("MCP list_prompts - User API Key Auth from context: %s", user_api_key_auth) - verbose_logger.debug("MCP list_prompts - MCP servers from context: %s", mcp_servers) - verbose_logger.debug( - "MCP list_prompts - MCP server auth headers: %s", - list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, - ) - # Get mcp_servers from context variable - verbose_logger.debug("MCP list_prompts - Calling _list_prompts") - prompts: Final = await _list_mcp_prompts( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) - verbose_logger.info("MCP list_prompts - Successfully returned %s prompts", len(prompts)) - return ListPromptsResult(prompts=prompts) - except Exception as e: - verbose_logger.exception("Error in list_prompts endpoint: %s", e) - # Return empty list instead of failing completely - # This prevents the HTTP stream from failing and allows the client to get a response - return ListPromptsResult(prompts=[]) # mutable-ok: MCP result payload - finally: - active_mcp_session_var.reset(_session_reset_token) - active_mcp_request_ctx_var.reset(_ctx_reset_token) + async with _legacy_operation_context(ctx, trace=False) as context: + return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute( + ListPromptsRequest(params=params), context + ) + except Exception as exc: # noqa: BLE001 # preserve native listing fallback for ingress failures + verbose_logger.exception("Error in list_prompts endpoint: %s", exc) + return ListPromptsResult(prompts=[]) async def get_prompt(ctx: ServerRequestContext, params: GetPromptRequestParams) -> GetPromptResult: - """ - Get a specific prompt with the provided arguments - """ if _mcp_proxy_mode.get(): _reject_mcp_proxy_operation() - _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) - _session_reset_token: Final = active_mcp_session_var.set(ctx.session) - - try: - ( - user_api_key_auth, - mcp_auth_header, - mcp_servers, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - _client_ip, - ) = await get_or_extract_auth_context() - - verbose_logger.debug("MCP mcp_server_tool_call - User API Key Auth from context: %s", user_api_key_auth) - return await mcp_get_prompt( - name=params.name, - arguments=params.arguments, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, + async with _legacy_operation_context(ctx, trace=False) as context: + return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute( + GetPromptRequest(params=params), context ) - finally: - active_mcp_session_var.reset(_session_reset_token) - active_mcp_request_ctx_var.reset(_ctx_reset_token) async def list_resources(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListResourcesResult: - """List all available resources.""" if _mcp_proxy_mode.get(): _reject_mcp_proxy_operation() - _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) - _session_reset_token: Final = active_mcp_session_var.set(ctx.session) - try: - ( - user_api_key_auth, - mcp_auth_header, - mcp_servers, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - _client_ip, - ) = await get_or_extract_auth_context() - verbose_logger.debug("MCP list_resources - User API Key Auth from context: %s", user_api_key_auth) - verbose_logger.debug("MCP list_resources - MCP servers from context: %s", mcp_servers) - verbose_logger.debug( - "MCP list_resources - MCP server auth headers: %s", - list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, - ) - - resources: Final = await _list_mcp_resources( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) - verbose_logger.info("MCP list_resources - Successfully returned %s resources", len(resources)) - return ListResourcesResult(resources=resources) - except Exception as e: - verbose_logger.exception("Error in list_resources endpoint: %s", e) - return ListResourcesResult(resources=[]) # mutable-ok: MCP result payload - finally: - active_mcp_session_var.reset(_session_reset_token) - active_mcp_request_ctx_var.reset(_ctx_reset_token) + async with _legacy_operation_context(ctx, trace=False) as context: + return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute( + ListResourcesRequest(params=params), context + ) + except Exception as exc: # noqa: BLE001 # preserve native listing fallback for ingress failures + verbose_logger.exception("Error in list_resources endpoint: %s", exc) + return ListResourcesResult(resources=[]) async def list_resource_templates( ctx: ServerRequestContext, params: PaginatedRequestParams ) -> ListResourceTemplatesResult: - """List all available resource templates.""" if _mcp_proxy_mode.get(): _reject_mcp_proxy_operation() - _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) - _session_reset_token: Final = active_mcp_session_var.set(ctx.session) - try: - ( - user_api_key_auth, - mcp_auth_header, - mcp_servers, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - _client_ip, - ) = await get_or_extract_auth_context() - verbose_logger.debug("MCP list_resource_templates - User API Key Auth from context: %s", user_api_key_auth) - verbose_logger.debug("MCP list_resource_templates - MCP servers from context: %s", mcp_servers) - verbose_logger.debug( - "MCP list_resource_templates - MCP server auth headers: %s", - list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None, - ) - - resource_templates: Final = await _list_mcp_resource_templates( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) - verbose_logger.info( - "MCP list_resource_templates - Successfully returned %s resource templates", len(resource_templates) - ) - return ListResourceTemplatesResult(resource_templates=resource_templates) - except Exception as e: - verbose_logger.exception("Error in list_resource_templates endpoint: %s", e) - return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload - finally: - active_mcp_session_var.reset(_session_reset_token) - active_mcp_request_ctx_var.reset(_ctx_reset_token) + async with _legacy_operation_context(ctx, trace=False) as context: + return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute( + ListResourceTemplatesRequest(params=params), context + ) + except Exception as exc: # noqa: BLE001 # preserve native listing fallback for ingress failures + verbose_logger.exception("Error in list_resource_templates endpoint: %s", exc) + return ListResourceTemplatesResult(resource_templates=[]) async def read_resource(ctx: ServerRequestContext, params: ReadResourceRequestParams) -> ReadResourceResult: if _mcp_proxy_mode.get(): _reject_mcp_proxy_operation() - _ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx) - _session_reset_token: Final = active_mcp_session_var.set(ctx.session) - - try: - ( - user_api_key_auth, - mcp_auth_header, - mcp_servers, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - _client_ip, - ) = await get_or_extract_auth_context() - - read_resource_result: Final = await mcp_read_resource( - url=params.uri, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, + async with _legacy_operation_context(ctx, trace=False) as context: + return await operations.GatewayOperations(_capture_host_progress_callback(ctx)).execute( + ReadResourceRequest(params=params), context ) - return read_resource_result - finally: - active_mcp_session_var.reset(_session_reset_token) - active_mcp_request_ctx_var.reset(_ctx_reset_token) - server.add_request_handler("tools/list", PaginatedRequestParams, handle_list_tools) server.add_request_handler("tools/call", CallToolRequestParams, mcp_server_tool_call) server.add_request_handler("prompts/list", PaginatedRequestParams, list_prompts) @@ -1533,527 +957,24 @@ if MCP_AVAILABLE: ############ Helper Functions ########################## ######################################################## - async def _get_allowed_mcp_servers_from_mcp_server_names( - mcp_servers: Sequence[str] | None, - allowed_mcp_servers: list[MCPServer], - ) -> list[MCPServer]: - """ - Get the filtered MCP servers from the MCP server names. - - Fails closed when ``mcp_servers`` is explicitly provided (path- or - header-derived) but none of the names resolve to a server alias or - access group the caller can access. The previous behavior returned - the full ``allowed_mcp_servers`` set, which silently widened scope - when a client targeted ``/mcp//`` and made URL/header - namespacing appear to work when it did not. - """ - - filtered_server: Final[dict[str, MCPServer]] = {} - # Filter servers based on mcp_servers parameter if provided - if mcp_servers is not None: - for server_or_group in mcp_servers: - server_name_matched = False - - for server in allowed_mcp_servers: - if server and _server_answers_to(server, server_or_group): - filtered_server[server.server_id] = server - server_name_matched = True - break - - if not server_name_matched: - try: - access_group_server_ids = await MCPRequestHandler._get_mcp_servers_from_access_groups( - [server_or_group] - ) - # Only include servers that the user has access to - for server_id in access_group_server_ids: - for server in allowed_mcp_servers: - if server_id == server.server_id: - filtered_server[server.server_id] = server - except Exception as e: - verbose_logger.debug("Could not resolve '%s' as access group: %s", server_or_group, e) - - if filtered_server: - return list(filtered_server.values()) - - if mcp_servers is not None: - # Caller asked for a specific scope but nothing resolved. Fail - # closed so URL/header namespacing cannot silently fall back to - # the caller's full allowed-server set. - verbose_logger.debug( - "MCP scope filter resolved to no servers for requested names %s; returning empty list (fail-closed).", - mcp_servers, - ) - return [] - - return allowed_mcp_servers - - def _http_detail_message(detail: object) -> str: - return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail) - - def _server_answers_to(server: MCPServer, name: str) -> bool: - requested: Final = name.lower() - return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known) - - class _McpDeniedDetail(TypedDict): - error: ReadOnly[str] - - async def raise_denied_scoped_mcp_access( - requested_names: Sequence[str], - user_api_key_auth: UserAPIKeyAuth | None, - client_ip: str | None = None, - ) -> None: - """A scoped request (``/mcp/`` path or ``x-mcp-servers`` header) resolved to zero - allowed servers, so the denial must be loud: a silent 200 with no tools reads as a healthy - server with no tools. Unknown, unauthorized, and access-group names all share one generic - error so scoping cannot probe which servers exist; the agent variant fires only when the - same request resolves once the agent binding is stripped, proving the binding caused the veto.""" - agent_id: Final = user_api_key_auth.agent_id if user_api_key_auth else None - if user_api_key_auth is not None and agent_id: - resolved_without_agent: Final = await _get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None})), - mcp_servers=requested_names, - client_ip=client_ip, - ) - - def _resolved_to_server(name: str) -> bool: - return any(_server_answers_to(server, name) for server in resolved_without_agent) - - vetoed_server: Final = next((name for name in requested_names if _resolved_to_server(name)), None) - if vetoed_server is not None: - agent_denial: Final[_McpDeniedDetail] = { - "error": ( - f"MCP server '{vetoed_server}' is not available to this key: the key is bound to " - f"agent '{agent_id}', whose MCP grants do not include this server. Add the server " - f"to the agent's object_permission.mcp_servers (edit the agent in the Admin UI or " - f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent." - ) - } - raise HTTPException(status_code=403, detail=agent_denial) - vetoed_group: Final = next( - ( - name - for name in requested_names - if not _resolved_to_server(name) - and any(name in (server.access_groups or ()) for server in resolved_without_agent) - ), - None, - ) - if vetoed_group is not None: - group_denial: Final[_McpDeniedDetail] = { - "error": ( - f"MCP access group '{vetoed_group}' is not available to this key: the key is bound to " - f"agent '{agent_id}', whose MCP grants do not include it. Add the group to the " - f"agent's object_permission.mcp_access_groups (edit the agent in the Admin UI or " - f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent." - ) - } - raise HTTPException(status_code=403, detail=group_denial) - generic_denial: Final[_McpDeniedDetail] = { - "error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}" - } - raise HTTPException(status_code=403, detail=generic_denial) - - def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool: - """ - Check if a tool name matches any name in the filter list. - - Reads the same owner the server-level permission checks use, so discovery hides - exactly what dispatch refuses. ``mcp_server`` is required: guessing the boundary - at the first separator mismatches every tool on a server whose prefix contains - the separator. - """ - bare_name: Final = strip_known_server_prefix(tool_name, mcp_server) - return match_known_tool_name(bare_name, mcp_server, filter_list) is not None - - def filter_tools_by_allowed_tools( - tools: list[MCPTool], - mcp_server: MCPServer, - ) -> list[MCPTool]: - """ - Filter tools by allowed/disallowed tools configuration. - - If allowed_tools is set, only tools in that list are returned. - If disallowed_tools is set, tools in that list are excluded. - Tool names are matched with and without server prefixes for flexibility. - - Args: - tools: List of tools to filter - mcp_server: Server configuration with allowed_tools/disallowed_tools - - Returns: - Filtered list of tools - """ - from litellm.proxy._experimental.mcp_server.utils import ( - server_applies_tool_allowlist, - ) - - tools_to_return = tools - - # Filter by allowed_tools (whitelist) - if server_applies_tool_allowlist(mcp_server): - if not mcp_server.allowed_tools: - return [] - tools_to_return = [ - tool for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools, mcp_server) - ] - - # Filter by disallowed_tools (blacklist) - if mcp_server.disallowed_tools: - tools_to_return = [ - tool - for tool in tools_to_return - if not _tool_name_matches(tool.name, mcp_server.disallowed_tools, mcp_server) - ] - - return tools_to_return - - def apply_tool_overrides( - tools: list[MCPTool], - mcp_server: MCPServer, - ) -> list[MCPTool]: - """Apply admin-configured display name/description overrides to tools. - - Overrides are keyed by the unprefixed tool name, same convention as - allowed_tools configuration. - """ - display_name_map: Final = mcp_server.tool_name_to_display_name or {} - description_map: Final = mcp_server.tool_name_to_description or {} - if not display_name_map and not description_map: - return tools - - for tool in tools: - unprefixed = strip_known_server_prefix(tool.name, mcp_server) - lookup_key = unprefixed or tool.name - if lookup_key in display_name_map: - tool.name = display_name_map[lookup_key] - if lookup_key in description_map: - tool.description = description_map[lookup_key] - return tools - - def _get_client_ip_from_context() -> str | None: - """ - Extract client_ip from auth context. - Returns None if context not set (caller should handle this as "no IP filtering"). - """ - try: - auth_user: Final = auth_context_var.get() - if auth_user and isinstance(auth_user, MCPAuthenticatedUser): - return auth_user.client_ip - except Exception: - pass - return None - - async def _get_allowed_mcp_servers( - user_api_key_auth: UserAPIKeyAuth | None, - mcp_servers: Sequence[str] | None, - client_ip: str | None = None, - ) -> list[MCPServer]: - """Return allowed MCP servers for a request after applying filters. - - Args: - user_api_key_auth: The authenticated user's API key info. - mcp_servers: Optional list of server names to filter to. - client_ip: Client IP for IP-based access control. If None, falls back to - auth context. Pass explicitly from request handlers for safety. - Note: If client_ip is None and auth context is not set, IP filtering is skipped. - This is intentional for internal callers but may indicate a bug if called - from a request handler without proper context setup. - """ - # Use explicit client_ip if provided, otherwise try auth context - if client_ip is None: - client_ip = _get_client_ip_from_context() - if client_ip is None: - verbose_logger.debug( - "MCP _get_allowed_mcp_servers called without client_ip and no auth context. " - "IP filtering will be skipped. This is expected for internal calls." - ) - - allowed_mcp_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) - ( - allowed_mcp_server_ids, - _ip_blocked, - ) = global_mcp_server_manager.filter_server_ids_by_ip_with_info(allowed_mcp_server_ids, client_ip) - verbose_logger.debug( - "MCP IP filter: client_ip=%s, allowed_server_ids=%s", - client_ip, - allowed_mcp_server_ids, - ) - if _ip_blocked > 0: - verbose_logger.debug( - "MCP IP filtering: %d server(s) are not accessible from client IP %s " - "because they are restricted to internal networks. " - "No tools from those servers will be returned. " - "To expose a server externally, set 'available_on_public_internet: true' " - "in its configuration.", - _ip_blocked, - client_ip, - ) - allowed_mcp_servers: list[MCPServer] = [] - for allowed_mcp_server_id in allowed_mcp_server_ids: - mcp_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id) - if mcp_server is not None: - # Apply the request-time oauth2_flow backstop for legacy null rows. - mcp_server = MCPServerManager.resolve_oauth2_flow_for_request(mcp_server) - allowed_mcp_servers.append(mcp_server) - - if mcp_servers is not None: - allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( - mcp_servers=mcp_servers, - allowed_mcp_servers=allowed_mcp_servers, - ) - - return allowed_mcp_servers - - def _client_has_per_server_auth_header( - server: MCPServer, - mcp_server_auth_headers: dict[str, dict[str, str]] | None, - ) -> bool: - """True if the request carries a per-server ``x-mcp-{alias}-authorization`` - header for this server. This is the multi-server binding: it names one - upstream, so it is unambiguously the caller's upstream token regardless of - auth mode (never the LiteLLM admission credential). - - Resolves through the same ``lookup_mcp_server_auth_in_headers`` egress uses, so - the connect gate and egress agree on which per-server header names match: a - dashboard client sends ``x-mcp-{sanitize_mcp_alias_for_header(alias)}-authorization``, - and matching only the raw alias here would 401 a token egress would forward. - """ - if not mcp_server_auth_headers: - return False - from litellm.proxy._experimental.mcp_server.utils import ( - lookup_mcp_server_auth_in_headers, - ) - - server_headers: Final = lookup_mcp_server_auth_in_headers( - mcp_server_auth_headers, - alias=server.alias, - server_name=server.server_name, - access_groups=server.access_groups, - ) - if isinstance(server_headers, str): - return bool(server_headers.strip()) - if isinstance(server_headers, dict): - return any(isinstance(hk, str) and hk.lower() == "authorization" for hk in server_headers) - return False - - def _client_has_passthrough_authorization( - server: MCPServer, - oauth2_headers: dict[str, str] | None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None, - ) -> bool: - """True if the incoming request already carries an ``Authorization`` - header the gateway will forward to this pass-through server. - - The client may supply the bearer as either the top-level - ``Authorization`` header (surfaced via ``oauth2_headers``) or a - per-server ``x-mcp-auth-`` style header (surfaced via - ``mcp_server_auth_headers``). Either form skips the pre-emptive 401. - """ - if oauth2_headers: - for k in oauth2_headers: - if k.lower() == "authorization": - return True - return _client_has_per_server_auth_header(server, mcp_server_auth_headers) - - async def _get_user_oauth_extra_headers_from_db( - server: MCPServer, - user_api_key_auth: UserAPIKeyAuth | None, - prefetched_creds: 'Mapping[str, "OAuthCredentialPayload"] | None' = None, - ) -> dict[str, str] | None: - """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None. - - Thin wrapper over ``resolve_user_oauth_access_token`` (Redis cache, else DB + refresh); - ``prefetched_creds`` skips the per-server Redis/DB lookups for the batch path. - """ - if server.auth_type != MCPAuth.oauth2 or user_api_key_auth is None: - return None - from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 - resolve_user_oauth_access_token, - ) - - token: Final = await resolve_user_oauth_access_token( - getattr(user_api_key_auth, "user_id", None), server, prefetched_creds - ) - return {"Authorization": f"Bearer {token}"} if token else None - - async def _prefetch_oauth_creds_for_user( - user_api_key_auth: UserAPIKeyAuth | None, - ) -> dict[str, "OAuthCredentialPayload"]: - """Fetch all OAuth2 credentials for the user in one DB query. - - Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. - """ - user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None - if not user_id: - return {} - try: - from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 - list_user_oauth_credentials, - ) - from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 - - prisma_client: Final = get_prisma_client_or_throw( - "Database not connected. Connect a database to use OAuth2 MCP tools." - ) - creds: Final = await list_user_oauth_credentials(prisma_client, user_id) - return {c["server_id"]: c for c in creds if "server_id" in c} - except Exception as e: - verbose_logger.warning("_prefetch_oauth_creds_for_user: failed to prefetch for user=%s: %s", user_id, e) - return {} - - def _prepare_mcp_server_headers( - server: MCPServer, - mcp_server_auth_headers: dict[str, dict[str, str]] | None, - mcp_auth_header: str | None, - oauth2_headers: dict[str, str] | None, - raw_headers: dict[str, str] | None, - user_api_key_auth: UserAPIKeyAuth | None = None, - scope_servers: list[MCPServer] | None = None, - ) -> tuple[dict[str, str] | str | None, dict[str, str] | None]: - """Build auth and extra headers for a server. - - ``scope_servers`` is the full server list a fan-out handler iterates. Passing it lets the - client-forwarded token modes withhold the caller's request-wide ``Authorization`` when - another server in the scope would also receive it (``_caller_authorization_fans_out``); - explicitly-addressed operations leave it None. Per-server ``x-mcp-{alias}-authorization`` - headers are unaffected — they bind one token to one server and are the multi-server shape. - """ - server_auth_header: dict[str, str] | str | None = None - if mcp_server_auth_headers: - from litellm.proxy._experimental.mcp_server.utils import ( - lookup_mcp_server_auth_in_headers, - ) - - server_auth_header = lookup_mcp_server_auth_in_headers( - mcp_server_auth_headers, - alias=server.alias, - server_name=server.server_name, - access_groups=server.access_groups, - ) - - extra_headers: dict[str, str] | None = None - is_client_forwarded_mode: Final = server.is_client_forwarded_token - # In a multi-server listing scope the request-wide Authorization can only carry one token, - # so it is withheld from a client-forwarded server when another server in scope also consumes - # it (RFC 9700 cross-resource replay); such scopes must bind per-server via - # x-mcp-{alias}-authorization. The decision is computed once so BOTH the forwarding branch and - # the extra_headers copy loop below honor it — otherwise a server that lists Authorization in - # extra_headers would re-copy the withheld bearer from raw_headers and replay it anyway. - withhold_forwarded_authorization: Final = is_client_forwarded_mode and _caller_authorization_fans_out( - server, scope_servers - ) - if server.auth_type == MCPAuth.oauth2: - # For OAuth2 M2M servers, upstream Authorization must come from - # client_credentials token fetch, never from caller headers. - if server.has_client_credentials: - extra_headers = None - else: - # Copy to avoid mutating the original dict (important for parallel fetching) - extra_headers = oauth2_headers.copy() if oauth2_headers else None - # Migrated authorization_code: the v2 resolver injects the stored per-user - # token, so drop the caller-forwarded Authorization (apply-if-absent would - # otherwise let it shadow the resolved token). Delegate keeps it. Centralized - # via _should_strip_caller_authorization to match _call_regular_mcp_tool. - if extra_headers and _should_strip_caller_authorization( - mcp_server=server, - raw_headers=raw_headers, - user_api_key_auth=user_api_key_auth, - ): - extra_headers = without_header(extra_headers, DEFAULT_CREDENTIAL_HEADER) - elif is_client_forwarded_mode: - if not withhold_forwarded_authorization: - extra_headers = _client_forwarded_authorization_headers( - mcp_server=server, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - user_api_key_auth=user_api_key_auth, - ) - - if server.extra_headers and raw_headers: - if extra_headers is None: - extra_headers = {} - - normalized_raw_headers: Final = {str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str)} - - # Centralized strip decision shared with - # ``MCPServerManager._call_regular_mcp_tool`` so the two - # code paths cannot drift on this security-sensitive choice. - # See ``_should_strip_caller_authorization`` for the rules. - strip_caller_authorization: Final = _should_strip_caller_authorization( - mcp_server=server, - raw_headers=raw_headers, - user_api_key_auth=user_api_key_auth, - ) - - for header in server.extra_headers: - if not isinstance(header, str): - continue - if header.lower() == "authorization" and ( - strip_caller_authorization or withhold_forwarded_authorization - ): - continue - header_value = normalized_raw_headers.get(header.lower()) - if header_value is None: - continue - extra_headers[header] = header_value - - # Reset to None if no headers were actually added - if extra_headers is not None and len(extra_headers) == 0: - extra_headers = None - - if server_auth_header is None: - server_auth_header = mcp_auth_header - - return server_auth_header, extra_headers - - def _merge_gateway_initialize_instructions( - allowed_mcp_servers: list[MCPServer], - ) -> str | None: - """YAML/DB override, else upstream text (prefetch on init, or list_tools / health_check / call_tool cache).""" - if not allowed_mcp_servers: - return None - - texts: Final[list[tuple[str, str]]] = [] - for server in allowed_mcp_servers: - label = server.alias or server.server_name or server.name or server.server_id or "mcp" - if server.instructions and server.instructions.strip(): - texts.append((label, server.instructions.strip())) - continue - if server.spec_path: - continue - cached = global_mcp_server_manager._upstream_initialize_instructions_by_server_id.get(server.server_id) - if cached and cached.strip(): - texts.append((label, cached.strip())) - - if not texts: - return None - if len(texts) == 1: - return texts[0][1] - return "\n\n---\n\n".join(f"[{lbl}]\n{txt}" for lbl, txt in texts) - - async def _raise_if_initialize_grants_no_mcp_servers( - allowed: Sequence[MCPServer], - user_api_key_auth: UserAPIKeyAuth | None, - mcp_servers: Sequence[str] | None, - client_ip: str | None, - ) -> None: - if allowed or user_api_key_auth is None or not user_api_key_auth.api_key: - return - if mcp_servers: - await raise_denied_scoped_mcp_access( - requested_names=mcp_servers, - user_api_key_auth=user_api_key_auth, - client_ip=client_ip, - ) - no_servers_denial: Final[_McpDeniedDetail] = { - "error": ( - "The key has no MCP servers granted, or none of its granted servers is loaded and allowed for " - "this client IP. Grant servers or access groups to the key, its team, or its organization " - "(object_permission.mcp_servers), check the server's allowed IPs, and reconnect." - ) - } - raise HTTPException(status_code=403, detail=no_servers_denial) + from litellm.proxy._experimental.mcp_server.operations import ( + _client_has_passthrough_authorization, + _client_has_per_server_auth_header, + _get_allowed_mcp_servers, + _get_allowed_mcp_servers_from_mcp_server_names, + _get_user_oauth_extra_headers_from_db, + _http_detail_message, + _McpDeniedDetail, + _merge_gateway_initialize_instructions, + _prefetch_oauth_creds_for_user, + _prepare_mcp_server_headers, + _raise_if_initialize_grants_no_mcp_servers, + _server_answers_to, + _tool_name_matches, + apply_tool_overrides, + filter_tools_by_allowed_tools, + raise_denied_scoped_mcp_access, + ) @contextlib.asynccontextmanager async def _gateway_initialize_instructions_request_scope( @@ -2063,26 +984,28 @@ if MCP_AVAILABLE: scoped_server_endpoint: bool = False, is_initialize: bool = False, ) -> AsyncIterator[None]: - allowed: Final = await _get_allowed_mcp_servers( + allowed: Final = await operations._get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip, ) if is_initialize: - await _raise_if_initialize_grants_no_mcp_servers(allowed, user_api_key_auth, mcp_servers, client_ip) + await operations._raise_if_initialize_grants_no_mcp_servers( + allowed, user_api_key_auth, mcp_servers, client_ip + ) if allowed: # return_exceptions=True: a per-server probe failure (incl. CancelledError # bubbled from anyio task group teardown on connection refused) must not # cancel sibling probes or 500 the gateway initialize request. await asyncio.gather( *[ - global_mcp_server_manager._ensure_upstream_initialize_instructions_cached(s) + operations.global_mcp_server_manager._ensure_upstream_initialize_instructions_cached(s) for s in allowed if s is not None ], return_exceptions=True, ) - merged: Final = _merge_gateway_initialize_instructions(allowed_mcp_servers=allowed) + merged: Final = operations._merge_gateway_initialize_instructions(allowed_mcp_servers=allowed) scoped_server_name = None if scoped_server_endpoint and len(allowed) == 1: scoped_server: Final = allowed[0] @@ -2097,1599 +1020,34 @@ if MCP_AVAILABLE: _mcp_gateway_initialize_instructions.reset(instructions_token) _mcp_gateway_server_name.reset(server_name_token) - def _aggregate_server_key(server: MCPServer) -> str: - """The client-visible key for a server in listing outcomes and spend metadata: the same - display prefix (alias, or the short prefix when that mode is enabled) the caller already - sees on the tool names. Canonical internal server names never key a caller-readable - surface; when the display naming deliberately hides them, the outcome keys must too.""" - return get_server_prefix(server) or "unknown" - - async def _get_tools_from_mcp_servers( - user_api_key_auth: UserAPIKeyAuth | None, - mcp_auth_header: str | None, - mcp_servers: list[str] | 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, - log_list_tools_to_spendlogs: bool = False, - list_tools_log_source: str | None = None, - litellm_trace_id: str | None = None, - request_tags: list[str] | None = None, - client_ip: str | None = None, - mcp_proxy_mode: bool = False, - ) -> AggregateToolListing: - """ - Helper method to fetch tools from MCP servers based on server filtering criteria. - - Args: - user_api_key_auth: User authentication info for access control - mcp_auth_header: Optional auth header for MCP server (deprecated) - mcp_servers: Optional list of server names/aliases to filter by - mcp_server_auth_headers: Optional dict of server-specific auth headers - oauth2_headers: Optional dict of oauth2 headers - - Returns: - AggregateToolListing: Combined tools from filtered servers plus each server's - classified listing outcome - """ - if not MCP_AVAILABLE: - return AggregateToolListing(tools=[], outcomes={}) - - list_tools_start_time: Final = datetime.now() - litellm_logging_obj: LiteLLMLoggingObj | None = None - list_tools_request_data: dict[str, object] = {} - - if log_list_tools_to_spendlogs: - # This is intentionally minimal: only async_success_handler / post_call_failure_hook - rules_obj: Final = Rules() - list_tools_call_id: Final = str(uuid.uuid4()) - # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) - effective_litellm_trace_id: Final = litellm_trace_id or get_chain_id_from_headers(raw_headers) - spend_logs_metadata: Final[dict[str, object]] = { - "mcp_operation": "list_tools", - } - if isinstance(list_tools_log_source, str): - spend_logs_metadata["source"] = list_tools_log_source - if isinstance(mcp_servers, list): - spend_logs_metadata["requested_mcp_servers"] = mcp_servers - - list_tools_request_data = { - "model": "MCP: list_tools", - "call_type": CallTypes.list_mcp_tools.value, - "litellm_call_id": list_tools_call_id, - "litellm_trace_id": effective_litellm_trace_id, - "metadata": { - "spend_logs_metadata": spend_logs_metadata, - "headers": logging_safe_mcp_headers(raw_headers), - **({"tags": request_tags} if request_tags else {}), - }, - # Provide a small input payload for standard logging - "input": [ - { - "role": "system", - "content": { - "mcp_operation": "list_tools", - "requested_mcp_servers": mcp_servers, - }, - } - ], - } - - # Attach user identifiers using the standard helper - if user_api_key_auth is not None: - LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( - data=list_tools_request_data, - user_api_key_dict=user_api_key_auth, - _metadata_variable_name="metadata", - ) - - user_identifier: Final = getattr(user_api_key_auth, "end_user_id", None) or getattr( - user_api_key_auth, "user_id", None - ) - if user_identifier: - list_tools_request_data["user"] = user_identifier - - try: - litellm_logging_obj, _ = function_setup( - original_function="list_mcp_tools", - rules_obj=rules_obj, - start_time=list_tools_start_time, - **list_tools_request_data, - ) - if litellm_logging_obj: - litellm_logging_obj.call_type = CallTypes.list_mcp_tools.value - litellm_logging_obj.model = "MCP: list_tools" - except Exception as logging_error: - verbose_logger.debug("Failed to initialize logging for MCP list_tools: %s", logging_error) - litellm_logging_obj = None - - try: - allowed_mcp_servers: Final = await _get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_servers=mcp_servers, - client_ip=client_ip, - ) - if mcp_servers and not allowed_mcp_servers: - await raise_denied_scoped_mcp_access( - requested_names=mcp_servers, - user_api_key_auth=user_api_key_auth, - client_ip=client_ip, - ) - - # Pre-fetch OAuth credentials only when at least one server uses OAuth2, - # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. - _has_oauth2_server = any(getattr(s, "auth_type", None) == MCPAuth.oauth2 for s in allowed_mcp_servers) - _prefetched_oauth_creds: Final = ( - await _prefetch_oauth_creds_for_user(user_api_key_auth) if _has_oauth2_server else {} - ) - - async def _fetch_and_filter_server_tools( - server: MCPServer, - ) -> "tuple[list[MCPTool], ServerOutcome]": - """Fetch and filter tools from a single server, classifying any failure into that - server's outcome so the aggregate can keep serving the healthy subset without a - broken server masquerading as an empty one.""" - if server is None: - return [], ServerListOk(tool_count=0) - - server_auth_header, extra_headers = _prepare_mcp_server_headers( - server=server, - mcp_server_auth_headers=mcp_server_auth_headers, - mcp_auth_header=mcp_auth_header, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - user_api_key_auth=user_api_key_auth, - scope_servers=allowed_mcp_servers, - ) - - # Prefer server-stored per-user OAuth when configured, so a stale - # Authorization header from the MCP client cannot override Redis/DB - # (same issue as call_tool in mcp_server_manager: VS Code caches tokens). - from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 - to_server_spec, - ) - - # A server migrated to the v2 resolver gets its token from the resolver at connect - # time; building it here would double-resolve and be shadowed by the v2 graft. The - # preemptive 401 already challenged a missing token, so one exists for the connect. - migrated_to_v2: Final = to_server_spec(server) is not None - if ( - not migrated_to_v2 - and server.auth_type == MCPAuth.oauth2 - and getattr(server, "needs_user_oauth_token", False) - and user_api_key_auth is not None - ): - db_headers: Final = await _get_user_oauth_extra_headers_from_db( - server, - user_api_key_auth, - prefetched_creds=_prefetched_oauth_creds, - ) - if db_headers: - extra_headers = db_headers - - # If still no OAuth2 token, fall back to pre-fetched creds (non-stale-client path) - elif not migrated_to_v2 and extra_headers is None and server.auth_type == MCPAuth.oauth2: - extra_headers = await _get_user_oauth_extra_headers_from_db( - server, - user_api_key_auth, - prefetched_creds=_prefetched_oauth_creds, - ) - - if server.is_byok and server.auth_type != MCPAuth.oauth2 and server_auth_header is None: - server_auth_header = await _get_byok_credential(server, user_api_key_auth) - - try: - tools: Final = await global_mcp_server_manager._get_tools_from_server( - server=server, - mcp_auth_header=server_auth_header, - extra_headers=extra_headers, - add_prefix=True, # Always add server prefix - raw_headers=raw_headers, - user_api_key_auth=user_api_key_auth, - oauth2_headers=oauth2_headers, - ) - filtered_tools = filter_tools_by_allowed_tools(tools, server) - - filtered_tools = await filter_tools_by_key_team_permissions( - tools=filtered_tools, - server_id=server.server_id, - user_api_key_auth=user_api_key_auth, - ) - - if mcp_proxy_mode: - from litellm.proxy._experimental.mcp_server.tool_search import with_mcp_proxy_identity - - filtered_tools = [ # mutable-ok: MCP tool pipeline - with_mcp_proxy_identity(tool, server.server_id) for tool in filtered_tools - ] - else: - filtered_tools = apply_tool_overrides(filtered_tools, server) - - verbose_logger.debug( - "Successfully fetched %s tools from server %s, %s after filtering", - len(tools), - server.name, - len(filtered_tools), - ) - return filtered_tools, ServerListOk(tool_count=len(filtered_tools)) - except MCPUpstreamAuthError as e: - # Absorb so one unauthenticated server does not empty every other server's - # tools. Surfacing the upstream 401 to the client as a re-auth challenge is - # intentionally not done here: raising from this list handler cannot produce a - # 401 + WWW-Authenticate (the MCP session manager serializes it as a JSON-RPC - # error). Single-server routes surface it via the request-scope preemptive - # check in _raise_preemptive_401_for_unauthenticated_servers instead. - verbose_logger.debug("MCP list_tools: omitting %s; it needs upstream auth", server.name) - return [], classify_list_exception(e) - except Exception as e: - verbose_logger.exception("Error getting tools from server %s: %s", server.name, e) - return [], classify_list_exception(e) - - # Fetch tools from all servers in parallel - tasks: Final = [_fetch_and_filter_server_tools(server) for server in allowed_mcp_servers] - results: Final = await asyncio.gather(*tasks) - - # Flatten results into single list - all_tools: Final[list[MCPTool]] = [tool for tools, _ in results for tool in tools] - server_outcomes: Final[dict[str, ServerOutcome]] = { - _aggregate_server_key(server): outcome - for server, (_, outcome) in zip(allowed_mcp_servers, results) - if server is not None - } - - # If logging is enabled, enrich spend_logs_metadata with counts - if litellm_logging_obj: - per_server_tool_counts: Final[dict[str, int]] = { - _aggregate_server_key(server): len(server_tools) - for server, (server_tools, _) in zip(allowed_mcp_servers, results) - if server is not None - } - - metadata_dict: Final = litellm_logging_obj.model_call_details.get("metadata") - if isinstance(metadata_dict, dict): - spend_meta = metadata_dict.get("spend_logs_metadata") - if not isinstance(spend_meta, dict): - spend_meta = {} - metadata_dict["spend_logs_metadata"] = spend_meta - spend_meta["allowed_server_count"] = len(allowed_mcp_servers) - spend_meta["tool_count_total"] = len(all_tools) - spend_meta["per_server_tool_counts"] = per_server_tool_counts - spend_meta["per_server_list_outcomes"] = { - key: outcome_wire_value(outcome) for key, outcome in server_outcomes.items() - } - - end_time: Final = datetime.now() - try: - await litellm_logging_obj.async_success_handler( - result=[ - tool.model_dump(mode="json") if isinstance(tool, MCPTool) else tool for tool in all_tools - ], - start_time=list_tools_start_time, - end_time=end_time, - ) - except Exception as log_exc: - # list_tools responses must not be dropped due to non-blocking - # observability/serialization failures. - verbose_logger.warning( - "MCP list_tools success logging failed (continuing): %s", - log_exc, - ) - - verbose_logger.info("Successfully fetched %s tools total from all MCP servers", len(all_tools)) - - return AggregateToolListing(tools=all_tools, outcomes=server_outcomes) - except Exception as e: - # Only fire failure hook if logging was requested for this list-tools execution - if log_list_tools_to_spendlogs and user_api_key_auth is not None: - try: - from litellm.proxy.proxy_server import proxy_logging_obj - - if proxy_logging_obj: - traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) - await proxy_logging_obj.post_call_failure_hook( - request_data=list_tools_request_data or {}, - original_exception=e, - user_api_key_dict=user_api_key_auth, - route="/mcp/list_tools", - traceback_str=traceback_str, - ) - except Exception: - verbose_logger.debug("Failed to log MCP list_tools failure via post_call_failure_hook") - raise - - async def _get_prompts_from_mcp_servers( - user_api_key_auth: UserAPIKeyAuth | None, - mcp_auth_header: str | None, - mcp_servers: list[str] | 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, - ) -> list[Prompt]: - """ - Helper method to fetch prompt from MCP servers based on server filtering criteria. - - Args: - user_api_key_auth: User authentication info for access control - mcp_auth_header: Optional auth header for MCP server (deprecated) - mcp_servers: Optional list of server names/aliases to filter by - mcp_server_auth_headers: Optional dict of server-specific auth headers - oauth2_headers: Optional dict of oauth2 headers - - Returns: - List[Prompt]: Combined list of prompts from filtered servers - """ - if not MCP_AVAILABLE: - return [] - - allowed_mcp_servers: Final = await _get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_servers=mcp_servers, - ) - - # Get prompts from each allowed server - all_prompts: Final = [] - for server in allowed_mcp_servers: - if server is None: - continue - - server_auth_header, extra_headers = _prepare_mcp_server_headers( - server=server, - mcp_server_auth_headers=mcp_server_auth_headers, - mcp_auth_header=mcp_auth_header, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - user_api_key_auth=user_api_key_auth, - scope_servers=allowed_mcp_servers, - ) - - try: - prompts = await global_mcp_server_manager.get_prompts_from_server( - server=server, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=server_auth_header, - extra_headers=extra_headers, - add_prefix=True, # Always add server prefix - raw_headers=raw_headers, - ) - - all_prompts.extend(prompts) - - verbose_logger.debug("Successfully fetched %s prompts from server %s", len(prompts), server.name) - except Exception as e: - verbose_logger.exception("Error getting prompts from server %s: %s", server.name, e) - # Continue with other servers instead of failing completely - - verbose_logger.info("Successfully fetched %s prompts total from all MCP servers", len(all_prompts)) - - return all_prompts - - async def _get_resources_from_mcp_servers( - user_api_key_auth: UserAPIKeyAuth | None, - mcp_auth_header: str | None, - mcp_servers: list[str] | 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, - ) -> list[Resource]: - """Fetch resources from allowed MCP servers.""" - - if not MCP_AVAILABLE: - return [] - - allowed_mcp_servers: Final = await _get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_servers=mcp_servers, - ) - - all_resources: Final[list[Resource]] = [] - for server in allowed_mcp_servers: - if server is None: - continue - - server_auth_header, extra_headers = _prepare_mcp_server_headers( - server=server, - mcp_server_auth_headers=mcp_server_auth_headers, - mcp_auth_header=mcp_auth_header, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - user_api_key_auth=user_api_key_auth, - scope_servers=allowed_mcp_servers, - ) - - try: - resources = await global_mcp_server_manager.get_resources_from_server( - server=server, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=server_auth_header, - extra_headers=extra_headers, - add_prefix=True, # Always add server prefix - raw_headers=raw_headers, - ) - all_resources.extend(resources) - - verbose_logger.debug("Successfully fetched %s resources from server %s", len(resources), server.name) - except Exception as e: - verbose_logger.exception("Error getting resources from server %s: %s", server.name, e) - - verbose_logger.info("Successfully fetched %s resources total from all MCP servers", len(all_resources)) - - return all_resources - - async def _get_resource_templates_from_mcp_servers( - user_api_key_auth: UserAPIKeyAuth | None, - mcp_auth_header: str | None, - mcp_servers: list[str] | 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, - ) -> list[ResourceTemplate]: - """Fetch resource templates from allowed MCP servers.""" - - if not MCP_AVAILABLE: - return [] - - allowed_mcp_servers: Final = await _get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_servers=mcp_servers, - ) - - all_resource_templates: Final[list[ResourceTemplate]] = [] - for server in allowed_mcp_servers: - if server is None: - continue - - server_auth_header, extra_headers = _prepare_mcp_server_headers( - server=server, - mcp_server_auth_headers=mcp_server_auth_headers, - mcp_auth_header=mcp_auth_header, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - user_api_key_auth=user_api_key_auth, - scope_servers=allowed_mcp_servers, - ) - - try: - resource_templates = await global_mcp_server_manager.get_resource_templates_from_server( - server=server, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=server_auth_header, - extra_headers=extra_headers, - add_prefix=True, # Always add server prefix - raw_headers=raw_headers, - ) - all_resource_templates.extend(resource_templates) - verbose_logger.debug( - "Successfully fetched %s resource templates from server %s", - len(resource_templates), - server.name, - ) - except Exception as e: - verbose_logger.exception( - "Error getting resource templates from server %s: %s", - server.name, - str(e), - ) - - verbose_logger.info( - "Successfully fetched %s resource templates total from all MCP servers", - len(all_resource_templates), - ) - - return all_resource_templates - - async def filter_tools_by_key_team_permissions( - tools: list[MCPTool], - server_id: str, - user_api_key_auth: UserAPIKeyAuth | None, - ) -> list[MCPTool]: - """ - Filter tools based on key/team mcp_tool_permissions. - - Note: Tool names in the DB are stored without server prefixes, - but tool names from MCP servers are prefixed. We need to strip - the prefix before comparing. - """ - # Filter by key/team tool-level permissions - allowed_tool_names: Final = await MCPRequestHandler.get_allowed_tools_for_server( - server_id=server_id, - user_api_key_auth=user_api_key_auth, - ) - - # Tools arrive prefixed with the server's own prefix; strip exactly that - # prefix (resolved from the server) rather than the first separator, so a - # prefix containing the separator still reduces to the stored bare name. - server: Final = global_mcp_server_manager.get_mcp_server_by_id(server_id) - return [ - t - for t in tools - if MCPRequestHandler.tool_is_granted(strip_known_server_prefix(t.name, server), allowed_tool_names) - ] - - async def _list_mcp_tools( - user_api_key_auth: UserAPIKeyAuth | None = None, - mcp_auth_header: str | None = None, - mcp_servers: list[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, - log_list_tools_to_spendlogs: bool = False, - list_tools_log_source: str | None = None, - client_ip: str | None = None, - mcp_proxy_mode: bool = False, - ) -> AggregateToolListing: - """ - List all available MCP tools. - - Args: - user_api_key_auth: User authentication info for access control - mcp_auth_header: Optional auth header for MCP server (deprecated) - mcp_servers: Optional list of server names/aliases to filter by - mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} - client_ip: Client IP for IP-based server access control - - Returns: - AggregateToolListing: Combined tools from all accessible servers plus each server's - classified listing outcome - """ - if not MCP_AVAILABLE: - return AggregateToolListing(tools=[], outcomes={}) - - try: - listing: Final = await _get_tools_from_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - log_list_tools_to_spendlogs=log_list_tools_to_spendlogs, - list_tools_log_source=list_tools_log_source, - client_ip=client_ip, - mcp_proxy_mode=mcp_proxy_mode, - ) - verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools)) - return listing - except HTTPException: - raise - except Exception as e: - verbose_logger.exception("Error getting tools from managed MCP servers: %s", e) - # Continue with an empty listing instead of failing completely - return AggregateToolListing(tools=[], outcomes={}) - - async def _list_mcp_prompts( - user_api_key_auth: UserAPIKeyAuth | None = None, - mcp_auth_header: str | None = None, - mcp_servers: list[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, - ) -> list[Prompt]: - """ - List all available MCP prompts. - - Args: - user_api_key_auth: User authentication info for access control - mcp_auth_header: Optional auth header for MCP server (deprecated) - mcp_servers: Optional list of server names/aliases to filter by - mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} - - Returns: - List[Prompt]: Combined list of tools from all accessible servers - """ - if not MCP_AVAILABLE: - return [] - # Get tools from managed MCP servers with error handling - managed_prompts = [] - try: - managed_prompts = await _get_prompts_from_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) - verbose_logger.debug("Successfully fetched %s prompts from managed MCP servers", len(managed_prompts)) - except Exception as e: - verbose_logger.exception("Error getting tools from managed MCP servers: %s", e) - # Continue with empty managed tools list instead of failing completely - - return managed_prompts - - async def _list_mcp_resources( - user_api_key_auth: UserAPIKeyAuth | None = None, - mcp_auth_header: str | None = None, - mcp_servers: list[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, - ) -> list[Resource]: - """List all available MCP resources.""" - - if not MCP_AVAILABLE: - return [] - - managed_resources: list[Resource] = [] - try: - managed_resources = await _get_resources_from_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) - verbose_logger.debug("Successfully fetched %s resources from managed MCP servers", len(managed_resources)) - except Exception as e: - verbose_logger.exception("Error getting resources from managed MCP servers: %s", e) - - return managed_resources - - async def _list_mcp_resource_templates( - user_api_key_auth: UserAPIKeyAuth | None = None, - mcp_auth_header: str | None = None, - mcp_servers: list[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, - ) -> list[ResourceTemplate]: - """List all available MCP resource templates.""" - - if not MCP_AVAILABLE: - return [] - - managed_resource_templates: list[ResourceTemplate] = [] - try: - managed_resource_templates = await _get_resource_templates_from_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) - verbose_logger.debug( - "Successfully fetched %s resource templates from managed MCP servers", - len(managed_resource_templates), - ) - except Exception as e: - verbose_logger.exception( - "Error getting resource templates from managed MCP servers: %s", - str(e), - ) - - return managed_resource_templates - - def _resolve_display_name_to_original( - name: str, - allowed_mcp_servers: list[MCPServer], - ) -> str: - """Translate a display-name override back to the original prefixed tool name. - - When a client received a customised display name from tools/list (e.g. - "Get Pet") it will call tools/call with that same string. We need to - reverse-map it to the original prefixed name (e.g. - "petstore_mcp-getPetById") before any routing or permission logic runs. - """ - for server in allowed_mcp_servers: - display_map = server.tool_name_to_display_name or {} - for unprefixed_name, display_name in display_map.items(): - if display_name == name: - return add_server_prefix_to_name(unprefixed_name, get_server_prefix(server)) - return name - - async def _get_byok_credential( - mcp_server: MCPServer, - user_api_key_auth: UserAPIKeyAuth | None, - ) -> str | None: - """Retrieve the stored BYOK credential for a user+server pair, served from the worker cache within its TTL.""" - if not mcp_server.is_byok: - return None - user_id: Final = (user_api_key_auth.user_id if user_api_key_auth else None) or "" - if not user_id: - return None - - cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id) - if cached is not None: - return cached.credential - - from litellm.proxy._experimental.mcp_server.db import get_user_credential - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - return None - credential: Final = await get_user_credential( - prisma_client=prisma_client, - user_id=user_id, - server_id=mcp_server.server_id, - ) - cache_byok_credential(user_id, mcp_server.server_id, credential) - return credential - - async def _check_byok_credential( - mcp_server: MCPServer, - user_api_key_auth: UserAPIKeyAuth | None, - ) -> None: - """ - If the MCP server is BYOK-enabled, verify that the requesting user has a - stored credential. When no credential is found, raise an HTTP 401 with a - WWW-Authenticate header that points the MCP client to our OAuth metadata - endpoint so it can drive the authorization flow. - """ - if not mcp_server.is_byok: - return - - user_id: Final = (user_api_key_auth.user_id if user_api_key_auth else None) or "" - if not user_id: - raise HTTPException( - status_code=401, - detail={ - "error": "byok_auth_required", - "server_id": mcp_server.server_id, - "server_name": mcp_server.server_name or mcp_server.name, - "message": "User identity is required for BYOK servers", - }, - headers={"WWW-Authenticate": get_byok_www_authenticate()}, - ) - - cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id) - if cached is not None: - if cached.credential is None: - raise HTTPException( - status_code=401, - detail={ - "error": "byok_auth_required", - "server_id": mcp_server.server_id, - "server_name": mcp_server.server_name or mcp_server.name, - "message": ( - "No stored credential found for this BYOK server. " - "Complete the OAuth authorization flow to provide your API key." - ), - }, - headers={"WWW-Authenticate": get_byok_www_authenticate()}, - ) - return - - from litellm.proxy._experimental.mcp_server.db import get_user_credential - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - # Fail closed on DB unavailability: returning here previously - # bypassed the ownership check and let any proxy-authenticated - # caller invoke BYOK tools during outage windows. - raise HTTPException( - status_code=503, - detail={ - "error": "byok_auth_unavailable", - "server_id": mcp_server.server_id, - "server_name": mcp_server.server_name or mcp_server.name, - "message": "BYOK credential check requires a database connection.", - }, - ) - - credential: Final = await get_user_credential( - prisma_client=prisma_client, - user_id=user_id, - server_id=mcp_server.server_id, - ) - cache_byok_credential(user_id, mcp_server.server_id, credential) - if credential is None: - raise HTTPException( - status_code=401, - detail={ - "error": "byok_auth_required", - "server_id": mcp_server.server_id, - "server_name": mcp_server.server_name or mcp_server.name, - "message": ( - "No stored credential found for this BYOK server. " - "Complete the OAuth authorization flow to provide your API key." - ), - }, - headers={"WWW-Authenticate": get_byok_www_authenticate()}, - ) - - async def _list_tools_before_first_call( - server: MCPServer | None, - tool_name: str, - allowed_mcp_servers: list[MCPServer], - user_api_key_auth: UserAPIKeyAuth | None, - mcp_auth_header: str | None, - mcp_server_auth_headers: dict[str, dict[str, str]] | None, - oauth2_headers: dict[str, str] | None, - raw_headers: dict[str, str] | None, - ) -> None: - """List ``server`` with the caller's own credentials when it does not yet expose ``tool_name`` here. - - The startup fill skips a server whose upstream wants the caller's token, and mcp 2 no - longer lists before an uncached tools/call, so a worker that has not served tools/list - for this caller would otherwise answer 404 for a tool the caller can see. Gating on the - requested tool, not on any prior listing, keeps callers with different upstream catalogs - from masking each other. - """ - if server is None or global_mcp_server_manager.server_exposes_tool(server, tool_name): - return - if all(allowed.server_id != server.server_id for allowed in allowed_mcp_servers): - return - try: - await _get_tools_from_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=[server.server_id], - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) - except Exception as e: # noqa: BLE001 # best effort: resolution below answers as it did before - verbose_logger.debug("MCP tools/call: listing %s before its first call failed: %s", server.name, e) - - async def execute_mcp_tool( - name: str, - arguments: dict[str, object], - allowed_mcp_servers: list[MCPServer], - start_time: datetime, - user_api_key_auth: UserAPIKeyAuth | None = 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, - host_progress_callback: Callable | None = None, - guardrail_context: Mapping[str, object] | None = None, - **kwargs: Any, - ) -> CallToolResult: - """ - Execute MCP tool. - - This function assumes permission checks have already been performed. - - Args: - name: Tool name (may include server prefix) - arguments: Tool arguments - allowed_mcp_servers: Pre-validated list of servers the user can access - start_time: Start time for logging - user_api_key_auth: Optional user API key auth for logging - mcp_auth_header: Optional MCP auth header - mcp_server_auth_headers: Optional server-specific auth headers - oauth2_headers: Optional OAuth2 headers - raw_headers: Optional raw HTTP headers - **kwargs: Additional arguments (e.g., litellm_logging_obj) - - Returns: - CallToolResult: Tool execution result - """ - # Track resolved MCP server for both permission checks and dispatch - mcp_server: MCPServer | None = None - requested_server_id: Final[str | None] = kwargs.get("requested_server_id") - - # If the client called with a display-name override (e.g. "Get Pet"), - # translate it back to the original prefixed name before any routing. - name = _resolve_display_name_to_original(name, allowed_mcp_servers) - - # Remove prefix from tool name for logging and processing - original_tool_name, server_name = split_server_prefix_from_name(name) - - requested_server: MCPServer | None = None - if requested_server_id: - requested_server = next( - (s for s in allowed_mcp_servers if s.server_id == requested_server_id), - None, - ) - - name_is_prefixed = False - if requested_server is not None and MCP_TOOL_PREFIX_SEPARATOR in name: - all_registry_prefixes: Final[set[str]] = set() - for registry_server in global_mcp_server_manager.get_registry().values(): - for known_prefix in iter_known_server_prefixes(registry_server): - all_registry_prefixes.add(normalize_server_name(known_prefix)) - name_is_prefixed = is_tool_name_prefixed(name, known_server_prefixes=all_registry_prefixes) - - first_call_target: Final = ( - requested_server - if requested_server is not None and not name_is_prefixed - else global_mcp_server_manager.server_owning_tool_name_prefix(name) - ) - first_call_tool_name: Final = ( - name - if first_call_target is None or (requested_server is not None and not name_is_prefixed) - else strip_known_server_prefix(name, first_call_target) - ) - await _list_tools_before_first_call( - server=first_call_target, - tool_name=first_call_tool_name, - allowed_mcp_servers=allowed_mcp_servers, - user_api_key_auth=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, - ) - - if requested_server is not None and not name_is_prefixed: - # REST callers may pass server_id with the upstream tool name (no - # LiteLLM prefix). The first segment is not a registered server - # prefix, so the whole string is the upstream tool name and may - # legitimately contain the separator (e.g. "text-to-speech"). - # server_id is authoritative for routing and auth. - mcp_server = requested_server - server_name = requested_server.name - original_tool_name = name - else: - # Resolve from tool name (MCP JSON-RPC or prefixed REST tool names). - mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) - if mcp_server is None and requested_server is not None: - for known_prefix in iter_known_server_prefixes(requested_server): - candidate = global_mcp_server_manager._get_mcp_server_from_tool_name( - add_server_prefix_to_name(name, known_prefix) - ) - if candidate is not None: - mcp_server = candidate - break - if mcp_server is not None: - server_name = mcp_server.name - original_tool_name = strip_known_server_prefix(name, mcp_server) - - if requested_server is not None: - if mcp_server is not None and mcp_server.server_id != requested_server.server_id: - raise HTTPException( - status_code=403, - detail={ - "error": "tool_server_mismatch", - "message": ( - f"Tool '{name}' belongs to MCP server " - f"'{mcp_server.name}' but request specified " - f"server_id for '{requested_server.name}'." - ), - }, - ) - if mcp_server is None: - mcp_server = requested_server - server_name = requested_server.name - original_tool_name = strip_known_server_prefix(name, requested_server) - - # Only enforce server-level permissions when we can resolve a server - if server_name: - if not MCPRequestHandler.is_tool_allowed( - allowed_mcp_servers=[server.name for server in allowed_mcp_servers], - server_name=server_name, - ): - raise HTTPException( - status_code=403, - detail="User not allowed to call this tool.", - ) - - standard_logging_mcp_tool_call: Final[StandardLoggingMCPToolCall] = _get_standard_logging_mcp_tool_call( - name=original_tool_name, # Use original name for logging - arguments=arguments, - server_name=server_name, - session_id=_mcp_session_id_from_headers(raw_headers), - ) - litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None) - 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: {name}" - litellm_logging_obj.model_call_details["model"] = f"MCP: {name}" - # Resolve the MCP server early so BYOK checks and credential injection - # apply to ALL dispatch paths (local tool registry AND managed MCP server). - if mcp_server is None: - mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) - - if mcp_server: - standard_logging_mcp_tool_call["mcp_server_cost_info"] = (mcp_server.mcp_info or {}).get( - "mcp_server_cost_info" - ) - if litellm_logging_obj: - litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call - - # BYOK: retrieve the stored per-user credential. A single DB call - # both checks existence and fetches the value, avoiding a double query. - if mcp_server.is_byok and not mcp_auth_header: - byok_cred: Final = await _get_byok_credential(mcp_server, user_api_key_auth) - if byok_cred is None: - raise HTTPException( - status_code=401, - detail={ - "error": "byok_auth_required", - "server_id": mcp_server.server_id, - "server_name": mcp_server.server_name or mcp_server.name, - "message": ( - "No stored credential found for this BYOK server. " - "Complete the OAuth authorization flow to provide your API key." - ), - }, - headers={"WWW-Authenticate": get_byok_www_authenticate()}, - ) - mcp_auth_header = byok_cred - elif mcp_server.is_byok: - # External auth header supplied; still enforce user-identity check. - await _check_byok_credential(mcp_server, user_api_key_auth) - - # Check if tool exists in local registry first (for OpenAPI-based tools) - # These tools are registered with their prefixed names - ######################################################### - local_tool: Final = global_mcp_tool_registry.get_tool(name) - if local_tool: - # OpenAPI-backed tools used to bypass `pre_call_tool_check` — - # only the managed path ran allowed/banned-tool checks, key/team - # tool permissions, and parameter validation. Run the same checks - # before dispatching to the local registry. Refuse the call if - # we cannot resolve a server: tools registered via - # openapi_to_mcp_generator are always tied to a server, so a - # missing mcp_server here means the tool->server mapping has - # not finished initializing or the registry entry is orphaned. - # Skipping the check would re-open the same authorization gap. - if mcp_server is None: - raise HTTPException( - status_code=503, - detail=( - f"MCP server for tool '{name}' is not available; " - "refusing to dispatch without authorization checks. " - "Retry once the server is registered." - ), - ) - - # `pre_call_tool_check` calls into `proxy_logging_obj` for the - # pre-call guardrail hooks, so source it from the canonical - # `proxy_server` module the same way `_handle_managed_mcp_tool` - # does. `kwargs.get("proxy_logging_obj")` is None on the MCP - # entry path and would crash with AttributeError after the - # security checks pass. - from litellm.proxy.proxy_server import proxy_logging_obj - - hook_result = await global_mcp_server_manager.pre_call_tool_check( - name=original_tool_name, - arguments=arguments or {}, - server_name=server_name or mcp_server.name, - user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, - server=mcp_server, - raw_headers=raw_headers, - litellm_logging_obj=litellm_logging_obj, - guardrail_context=guardrail_context, - ) - # `pre_call_tool_check` may return guardrail-modified - # arguments; honor them on the local path too. - if isinstance(hook_result, dict) and "arguments" in hook_result: - arguments = hook_result["arguments"] - - verbose_logger.debug("Executing local registry tool: %s", name) - # The credential rides ContextVars because the tool function has its - # headers baked into the closure at registration time. - auth_header_value, openapi_forwarded_headers, upstream_credential = _resolve_openapi_tool_auth( - mcp_server=mcp_server, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - raw_headers=raw_headers, - user_api_key_auth=user_api_key_auth, - ) - ( - resolved_auth_headers, - forwarded_headers, - ) = await global_mcp_server_manager.resolve_openapi_upstream_auth( - mcp_server=mcp_server, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - mcp_auth_header=upstream_credential, - user_api_key_auth=user_api_key_auth, - forwarded_headers=openapi_forwarded_headers, - ) - - _auth_token: Final = _request_auth_header.set(auth_header_value) - _extra_token: Final = _request_extra_headers.set(forwarded_headers) - _resolved_token: Final = _request_resolved_auth_headers.set(resolved_auth_headers) - try: - response = await _handle_local_mcp_tool(name, arguments) - finally: - _request_auth_header.reset(_auth_token) - _request_extra_headers.reset(_extra_token) - _request_resolved_auth_headers.reset(_resolved_token) - - # Try managed MCP server tool (the name is bare; the prefix boundary was - # already resolved above against this server's registered prefixes) - # Primary and recommended way to use external MCP servers - ######################################################### - elif mcp_server: - response = await _handle_managed_mcp_tool( - server_name=server_name, - name=original_tool_name, - arguments=arguments, - user_api_key_auth=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, - litellm_logging_obj=litellm_logging_obj, - guardrail_context=guardrail_context, - host_progress_callback=host_progress_callback, - ) - - # Fall back to local tool registry with original name (legacy support) - ######################################################### - # Deprecated: Local MCP Server Tool - ######################################################### - else: - # Gate only what can actually dispatch. When the unprefixed name is - # not in the registry either, `_handle_local_mcp_tool` below reports - # 404 and nothing runs, so demanding a server here would turn every - # unknown tool name into a misleading 503. - if global_mcp_tool_registry.get_tool(original_tool_name) is not None: - # `mcp_server` is None here because the tool name is not in the - # tool -> server mapping, but the name still carries a prefix - # that the server-level check above compared against the - # caller's `allowed_mcp_servers` by exact `name`. So the named - # server is in that list and can carry the tool-level checks, - # even with the mapping cold. Resolve it from - # `allowed_mcp_servers` rather than the registry: the registry - # would happily return a server the caller holds no grant for, - # and matching anything other than `name` would accept a server - # the check never validated. - prefix_server: Final = next( - (candidate for candidate in allowed_mcp_servers if candidate.name == server_name), - None, - ) - if prefix_server is None: - # A non-empty prefix that passed the server-level check - # always matches here, so this arm only fires when the - # prefix was empty, which is exactly the case that check - # skips. Fail closed rather than dispatch with no server to - # evaluate a tool ceiling against. - raise HTTPException( - status_code=503, - detail=( - f"MCP server for tool '{original_tool_name}' is not available; " - "refusing to dispatch without authorization checks. " - "Retry once the server is registered." - ), - ) - - from litellm.proxy.proxy_server import proxy_logging_obj - - hook_result = await global_mcp_server_manager.pre_call_tool_check( - name=original_tool_name, - arguments=arguments, - server_name=server_name, - user_api_key_auth=user_api_key_auth, - proxy_logging_obj=proxy_logging_obj, - server=prefix_server, - raw_headers=raw_headers, - litellm_logging_obj=litellm_logging_obj, - guardrail_context=guardrail_context, - ) - if "arguments" in hook_result: - arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args - - response = await _handle_local_mcp_tool(original_tool_name, arguments) - - return await _run_post_mcp_call_guardrails( - result=response, - litellm_logging_obj=litellm_logging_obj, - user_api_key_auth=user_api_key_auth, - request_data=kwargs, - ) - - async def _run_post_mcp_call_guardrails( - result: CallToolResult, - litellm_logging_obj: LiteLLMLoggingObj | None, - user_api_key_auth: UserAPIKeyAuth | None, - request_data: Mapping[str, object], - ) -> CallToolResult: - """Run ``post_mcp_call`` guardrails over an executed tool result. - - Lives on ``execute_mcp_tool``'s return path rather than inside - ``_fire_mcp_tool_call_logging`` so enforcement never depends on logging - being configured, and so every dispatch route gets it: the MCP protocol - handler, the REST endpoint, and tool search all funnel through here. - A guardrail that rejects the result raises, matching ``pre_mcp_call``. - """ - from litellm.proxy.proxy_server import proxy_logging_obj - - if proxy_logging_obj is None: - return result - return await proxy_logging_obj.post_mcp_call_hook( - response=result, - request_data=( - litellm_logging_obj.model_call_details if litellm_logging_obj is not None else dict(request_data) - ), - user_api_key_dict=user_api_key_auth, - ) - - _MCP_CREDENTIAL_REQUEST_FIELDS: Final = frozenset( - { - "raw_headers", - "mcp_auth_header", - "mcp_server_auth_headers", - "oauth2_headers", - "user_api_key_auth", - } + from litellm.proxy._experimental.mcp_server.operations import ( + _MCP_CREDENTIAL_REQUEST_FIELDS, + _aggregate_server_key, + _check_byok_credential, + _fire_mcp_tool_call_logging, + _get_byok_credential, + _get_prompts_from_mcp_servers, + _get_resource_templates_from_mcp_servers, + _get_resources_from_mcp_servers, + _get_standard_logging_mcp_tool_call, + _get_tools_from_mcp_servers, + _handle_local_mcp_tool, + _handle_managed_mcp_tool, + _list_mcp_prompts, + _list_mcp_resource_templates, + _list_mcp_resources, + _list_mcp_tools, + _list_tools_before_first_call, + _resolve_display_name_to_original, + _run_post_mcp_call_guardrails, + call_mcp_tool, + execute_mcp_tool, + filter_tools_by_key_team_permissions, + fire_mcp_tool_call_failure_logging, + mcp_get_prompt, + mcp_read_resource, ) - async def _fire_mcp_tool_call_logging( - logging_obj: LiteLLMLoggingObj, - result: CallToolResult, - start_time: datetime, - end_time: datetime, - user_api_key_auth: UserAPIKeyAuth | None = None, - request_data: Mapping[str, object] | None = None, - ) -> CallToolResult: - """Fire post-call logging for an executed MCP tool call, returning the result to send. - - The returned result is what the caller must forward to the client: a - ``post_mcp_call`` guardrail may rewrite the tool output (e.g. mask - sensitive values) or reject it, in which case its exception propagates. - Guardrails run before the success/failure logging so the masked text, not - the raw one, is what gets logged. - - A result with ``is_error=True`` is logged as a failure (``status="failure"`` - payload, so OTel marks the span ERROR) while the HTTP wire behavior stays - 200 + ``isError: true`` per the MCP spec. The error check runs after - ``async_post_mcp_tool_call_hook`` because guardrails may flip the result - to ``is_error=True`` in that hook. Raised exceptions never reach here (the - ``@client`` wrapper and ``call_mcp_tool``'s except path log those), so - this cannot double-log a failure. - - ``request_data`` may carry credential-bearing fields (the REST path puts - ``raw_headers``, ``mcp_auth_header``, ``mcp_server_auth_headers``, and - ``oauth2_headers`` at the top level of its data dict), so those are - stripped before the dict is handed to ``post_call_failure_hook`` - callbacks. - """ - from litellm.proxy.proxy_server import proxy_logging_obj - - logging_obj.post_call(original_response=result) - await logging_obj.async_post_mcp_tool_call_hook( - kwargs=logging_obj.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - logging_obj.call_type = CallTypes.call_mcp_tool.value - error_message: Final = extract_mcp_tool_result_error_message(result) - if error_message is None: - await logging_obj.async_success_handler(result=result, start_time=start_time, end_time=end_time) - return result - - logging_obj.has_run_logging(event_type="sync_success") - logging_obj.has_run_logging(event_type="async_success") - tool_error: Final = MCPToolResultError(error_message) - logging_obj.failure_handler(tool_error, "", start_time, end_time) - await logging_obj.async_failure_handler(tool_error, "", start_time, end_time) - - if user_api_key_auth is None: - return result - - if proxy_logging_obj: - sanitized_request_data: Final = { - key: value for key, value in (request_data or {}).items() if key not in _MCP_CREDENTIAL_REQUEST_FIELDS - } - await proxy_logging_obj.post_call_failure_hook( - request_data=sanitized_request_data, - original_exception=tool_error, - user_api_key_dict=user_api_key_auth, - route="/mcp/call_tool", - ) - return result - - async def fire_mcp_tool_call_failure_logging( - logging_obj: LiteLLMLoggingObj | None, - exception: Exception, - start_time: datetime, - user_api_key_auth: UserAPIKeyAuth | None, - request_data: Mapping[str, object], - ) -> None: - """Failure logging shared by the ``/mcp`` path and the REST endpoint. Call from - inside the ``except`` block so the traceback is still available. - - The failure handlers run first because ``_ProxyDBLogger.async_post_call_failure_hook`` - builds the failure spend-log row from the ``standard_logging_object`` they produce; - both gate on ``should_run_logging``, so the ``@client`` wrapper does not log twice. - A relayed upstream 401 (``MCPUpstreamAuthError``) is an expected caller-must-reauth - signal and skips ``post_call_failure_hook``, which fires the ``llm_exceptions`` alert. - """ - from litellm.proxy.proxy_server import proxy_logging_obj - - traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) - if logging_obj is not None: - end_time: Final = datetime.now() # noqa: DTZ005 # naive to match `start_time`, which it is subtracted from - logging_obj.failure_handler(exception, traceback_str, start_time, end_time) - await logging_obj.async_failure_handler(exception, traceback_str, start_time, end_time) - - if isinstance(exception, MCPUpstreamAuthError) or not proxy_logging_obj or user_api_key_auth is None: - return - sanitized_request_data: Final = { - key: value for key, value in request_data.items() if key not in _MCP_CREDENTIAL_REQUEST_FIELDS - } - await proxy_logging_obj.post_call_failure_hook( - request_data=sanitized_request_data, - original_exception=exception, - user_api_key_dict=user_api_key_auth, - route="/mcp/call_tool", - traceback_str=traceback_str, - ) - - @client - async def call_mcp_tool( - name: str, - arguments: dict[str, object] | None = None, - user_api_key_auth: UserAPIKeyAuth | None = None, - mcp_auth_header: str | None = None, - mcp_servers: list[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, - client_ip: str | None = None, - **kwargs: Any, - ) -> CallToolResult: - """ - Call a specific tool with the provided arguments (handles prefixed tool names). - """ - start_time: Final = datetime.now() - litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None) - - try: - if arguments is None: - raise HTTPException(status_code=400, detail="Request arguments are required") - - ## CHECK IF USER IS ALLOWED TO CALL THIS TOOL - allowed_mcp_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth, - ) - - allowed_mcp_servers: list[MCPServer] = [] - for allowed_mcp_server_id in allowed_mcp_server_ids: - allowed_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id) - if allowed_server is not None: - # Same request-time oauth2_flow backstop the listing path applies, - # so a null-flow M2M-shape row is treated as M2M on tool calls too. - allowed_server = MCPServerManager.resolve_oauth2_flow_for_request(allowed_server) - allowed_mcp_servers.append(allowed_server) - - allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( - mcp_servers=mcp_servers, - allowed_mcp_servers=allowed_mcp_servers, - ) - if mcp_servers and not allowed_mcp_servers: - await raise_denied_scoped_mcp_access( - requested_names=mcp_servers, - user_api_key_auth=user_api_key_auth, - client_ip=client_ip, - ) - if not allowed_mcp_servers: - raise HTTPException( - status_code=403, - detail="User not allowed to call this tool.", - ) - - # Delegate to execute_mcp_tool for execution - response = await execute_mcp_tool( - name=name, - arguments=arguments, - allowed_mcp_servers=allowed_mcp_servers, - start_time=start_time, - user_api_key_auth=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, - **kwargs, - ) - except Exception as e: - await fire_mcp_tool_call_failure_logging(litellm_logging_obj, e, start_time, user_api_key_auth, kwargs) - raise - - if litellm_logging_obj: - response = await _fire_mcp_tool_call_logging( - logging_obj=litellm_logging_obj, - result=response, - start_time=start_time, - end_time=datetime.now(), - user_api_key_auth=user_api_key_auth, - request_data=kwargs, - ) - return response - - async def mcp_get_prompt( - name: str, - arguments: dict[str, object] | None = None, - user_api_key_auth: UserAPIKeyAuth | None = None, - mcp_auth_header: str | None = None, - mcp_servers: list[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, - ) -> GetPromptResult: - """ - Fetch a specific MCP prompt, handling both prefixed and unprefixed names. - """ - allowed_mcp_servers: Final = await _get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_servers=mcp_servers, - ) - - if not allowed_mcp_servers: - raise HTTPException( - status_code=403, - detail="User not allowed to get this prompt.", - ) - - # Extract server name from prefixed prompt name - original_prompt_name, server_name = split_server_prefix_from_name(name) - - server: Final = next((s for s in allowed_mcp_servers if s.name == server_name), None) - if server is None: - raise HTTPException( - status_code=403, - detail="User not allowed to get this prompt.", - ) - - server_auth_header, extra_headers = _prepare_mcp_server_headers( - server=server, - mcp_server_auth_headers=mcp_server_auth_headers, - mcp_auth_header=mcp_auth_header, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - user_api_key_auth=user_api_key_auth, - ) - - return await global_mcp_server_manager.get_prompt_from_server( - server=server, - user_api_key_auth=user_api_key_auth, - prompt_name=original_prompt_name, - arguments=arguments, - mcp_auth_header=server_auth_header, - extra_headers=extra_headers, - raw_headers=raw_headers, - ) - - async def mcp_read_resource( - url: AnyUrl, - user_api_key_auth: UserAPIKeyAuth | None = None, - mcp_auth_header: str | None = None, - mcp_servers: list[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, - ) -> ReadResourceResult: - """Read resource contents from upstream MCP servers.""" - - allowed_mcp_servers: Final = await _get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_servers=mcp_servers, - ) - - if not allowed_mcp_servers: - raise HTTPException( - status_code=403, - detail="User not allowed to read this resource.", - ) - - if len(allowed_mcp_servers) != 1: - raise HTTPException( - status_code=400, - detail=( - "Multiple MCP servers configured; read_resource currently supports exactly one allowed server." - ), - ) - - server: Final = allowed_mcp_servers[0] - - server_auth_header, extra_headers = _prepare_mcp_server_headers( - server=server, - mcp_server_auth_headers=mcp_server_auth_headers, - mcp_auth_header=mcp_auth_header, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - user_api_key_auth=user_api_key_auth, - ) - - return await global_mcp_server_manager.read_resource_from_server( - server=server, - user_api_key_auth=user_api_key_auth, - url=url, - mcp_auth_header=server_auth_header, - extra_headers=extra_headers, - raw_headers=raw_headers, - ) - - def _get_standard_logging_mcp_tool_call( - name: str, - arguments: dict[str, object], - server_name: str | None, - session_id: str | None = None, - ) -> StandardLoggingMCPToolCall: - mcp_server: Final = global_mcp_server_manager._get_mcp_server_from_tool_name( - add_server_prefix_to_name(name, server_name) if server_name else name - ) - namespaced_tool_name: Final = f"{server_name}/{name}" if server_name else name - if mcp_server: - mcp_info: Final = mcp_server.mcp_info or {} - return StandardLoggingMCPToolCall( - name=name, - arguments=arguments, - mcp_server_name=mcp_info.get("server_name"), - mcp_server_logo_url=mcp_info.get("logo_url"), - namespaced_tool_name=namespaced_tool_name, - mcp_session_id=session_id, - mcp_auth_mode=mcp_server.auth_type, - mcp_server_resource=_redact_mcp_resource_url(mcp_server.url), - ) - else: - return StandardLoggingMCPToolCall( - name=name, - arguments=arguments, - namespaced_tool_name=namespaced_tool_name, - mcp_session_id=session_id, - ) - - async def _handle_managed_mcp_tool( - server_name: str, - name: str, - arguments: dict[str, object], - user_api_key_auth: UserAPIKeyAuth | None = 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_logging_obj: LiteLLMLoggingObj | None = None, - host_progress_callback: Callable | None = None, - guardrail_context: Mapping[str, object] | None = None, - ) -> CallToolResult: - """Handle tool execution for managed server tools""" - # Import here to avoid circular import - from litellm.proxy.proxy_server import proxy_logging_obj - - call_tool_result: Final = await global_mcp_server_manager.call_tool( - server_name=server_name, - name=name, - arguments=arguments, - user_api_key_auth=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, - host_progress_callback=host_progress_callback, - litellm_logging_obj=litellm_logging_obj, - guardrail_context=guardrail_context, - ) - verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) - return call_tool_result - - async def _handle_local_mcp_tool(name: str, arguments: dict[str, object]) -> CallToolResult: - """Execute a local-registry tool and report whether it succeeded. - - Returns the result rather than bare content because the verdict is part of it: the content - alone cannot say whether the handler failed, so callers used to stamp is_error=False on every - outcome and an upstream rejection was served as tool output. - - A failure is reported as ``is_error=True`` here rather than raised, because the REST surface - turns an unrecognized exception into a 500 and an upstream 403 or 429 is not a gateway crash. - ``MCPUpstreamAuthError`` is the exception: it propagates so the caller is told to - re-authenticate, which both renderers already know how to say. - - Note: Local tools don't use prefixes, so we use the original name - """ - import inspect - - tool: Final = global_mcp_tool_registry.get_tool(name) - if not tool: - raise HTTPException(status_code=404, detail=f"Tool '{name}' not found") - - try: - if inspect.iscoroutinefunction(tool.handler): - result = await tool.handler(**arguments) - else: - result = tool.handler(**arguments) - except MCPUpstreamAuthError: - raise - except Exception as e: - verbose_logger.exception("Error executing local tool %s: %s", name, e) - return CallToolResult( - content=[TextContent(text=f"Error: {e}", type="text")], # mutable-ok: MCP result content - is_error=True, - ) - return CallToolResult( - content=[TextContent(text=str(result), type="text")], # mutable-ok: MCP result content - is_error=False, - ) - def _get_mcp_servers_in_path(path: str) -> list[str] | None: """ Get the MCP servers from the path @@ -4178,7 +1536,9 @@ if MCP_AVAILABLE: detail=f"API key does not have access to toolset '{toolset_id}'.", ) - tool_permissions = await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=[toolset_id]) + tool_permissions = await operations.global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=[toolset_id] + ) server_ids: Final = list(tool_permissions.keys()) existing_op: Final = user_api_key_auth.object_permission if existing_op is not None: @@ -4197,7 +1557,7 @@ if MCP_AVAILABLE: mcp_servers=server_ids, mcp_tool_permissions=tool_permissions, ) - return user_api_key_auth.model_copy(update={"object_permission": updated_op}) + return user_api_key_auth.model_copy(update={"object_permission": updated_op, "mcp_toolset_id": toolset_id}) async def _raise_preemptive_401_for_unauthenticated_servers( scope: Scope, @@ -4221,7 +1581,7 @@ if MCP_AVAILABLE: a server it will be 403'd on immediately after authentication. """ for server_name in mcp_servers or []: - server = global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip) + server = operations.global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip) if server is not None and allowed_server_ids is not None and server.server_id not in allowed_server_ids: # Caller's narrowed scope excludes this server — skip the # preemptive challenge and let downstream authorization @@ -4234,7 +1594,7 @@ if MCP_AVAILABLE: # authorization_url/token_url can change their inferred flow. continue if server is not None: - server = await global_mcp_server_manager.ensure_oauth_metadata_discovered(server) + server = await operations.global_mcp_server_manager.ensure_oauth_metadata_discovered(server) if server and server.auth_type == MCPAuth.oauth2: # The challenge decision is per oauth2 sub-mode, not per header: # gateway-managed modes (M2M and interactive authorization_code) @@ -4262,7 +1622,7 @@ if MCP_AVAILABLE: # authorization server is the gateway itself, vaulting via the # authorize interlude); the per-server relay advertised below # cannot vault without a litellm key on its token request. - if await global_mcp_server_manager.has_user_oauth_token(server, user_api_key_auth): + if await operations.global_mcp_server_manager.has_user_oauth_token(server, user_api_key_auth): continue if _is_mcp_admitted_user_subject(user_api_key_auth): @@ -4345,12 +1705,12 @@ if MCP_AVAILABLE: and server.server_id in frozenset( allowed.server_id - for allowed in await _get_allowed_mcp_servers( + for allowed in await operations._get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip ) ) ): - await global_mcp_server_manager.preflight_token_exchange( + await operations.global_mcp_server_manager.preflight_token_exchange( server=server, oauth2_headers=oauth2_headers, user_api_key_auth=user_api_key_auth, @@ -4366,7 +1726,9 @@ if MCP_AVAILABLE: if ( server and server.is_oauth_passthrough - and not _client_has_passthrough_authorization(server, oauth2_headers, mcp_server_auth_headers) + and not operations._client_has_passthrough_authorization( + server, oauth2_headers, mcp_server_auth_headers + ) ): www_authenticate = get_passthrough_www_authenticate( scope=scope, @@ -4383,7 +1745,7 @@ if MCP_AVAILABLE: and server.is_oauth_delegate and len(mcp_servers or []) == 1 and _get_forwarded_auth_from_scope(scope) is None - and not _client_has_per_server_auth_header(server, mcp_server_auth_headers) + and not operations._client_has_per_server_auth_header(server, mcp_server_auth_headers) ): www_authenticate = get_passthrough_www_authenticate( scope=scope, @@ -4400,7 +1762,7 @@ if MCP_AVAILABLE: and server.is_true_passthrough and len(mcp_servers or []) == 1 and not _scope_has_authorization_header(scope) - and not _client_has_per_server_auth_header(server, mcp_server_auth_headers) + and not operations._client_has_per_server_auth_header(server, mcp_server_auth_headers) ): if server.is_dcr_bridge: raise HTTPException( @@ -4528,7 +1890,7 @@ if MCP_AVAILABLE: # Use the authorized server set, not the raw user-supplied names, so that # a caller cannot force a probe to a server their key is not allowed to use. - allowed_servers: Final = await _get_allowed_mcp_servers( + allowed_servers: Final = await operations._get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip, diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index a482d02c31d..3650c722103 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -463,8 +463,8 @@ async def handle_mcp_tool_search( oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, ) -> CallToolResult: - from litellm.proxy._experimental.mcp_server.server import ( - _list_mcp_tools, # pyright: ignore[reportPrivateUsage] # shared catalog owner + from litellm.proxy._experimental.mcp_server.operations import ( + _list_mcp_tools, ) from litellm.proxy.proxy_server import llm_router, proxy_logging_obj @@ -519,8 +519,8 @@ async def handle_mcp_proxy_tool( from jsonschema import validate from litellm.proxy import proxy_server - from litellm.proxy._experimental.mcp_server.server import ( # pyright: ignore[reportPrivateUsage] # shared catalog owner - _list_mcp_tools, # pyright: ignore[reportPrivateUsage] # shared catalog owner + from litellm.proxy._experimental.mcp_server.operations import ( + _list_mcp_tools, ) listing: Final = await _list_mcp_tools( @@ -607,7 +607,7 @@ async def handle_mcp_tool_call( requested_server_id: str | None = None, guardrail_context: Mapping[str, object] | None = None, ) -> CallToolResult: - from litellm.proxy._experimental.mcp_server.server import ( + from litellm.proxy._experimental.mcp_server.operations import ( _get_allowed_mcp_servers, execute_mcp_tool, raise_denied_scoped_mcp_access, @@ -643,6 +643,7 @@ async def handle_mcp_tool_call( mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, + client_ip=client_ip, litellm_logging_obj=litellm_logging_obj, requested_server_id=requested_server_id, guardrail_context=guardrail_context, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3f6ef89fbc6..0592616f06a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3239,6 +3239,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob # above; a forged value could at most narrow, but the stripping keeps the field's provenance # single-owner so its meaning stays trustworthy. mcp_session_resource_server_id: str | None = Field(default=None, exclude=True) + mcp_toolset_id: str | None = Field(default=None, exclude=True) via_virtual_key: bool = Field( default=False, exclude=True, @@ -3280,6 +3281,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob values.pop("mcp_admitted_user_subject", None) values.pop("mcp_source_team_rpm_limits", None) values.pop("mcp_session_resource_server_id", None) + values.pop("mcp_toolset_id", None) values.pop("via_virtual_key", None) if values.get("api_key") is not None: values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))}) diff --git a/scripts/check_mcp_operation_boundary.py b/scripts/check_mcp_operation_boundary.py new file mode 100644 index 00000000000..b6c9dcefabf --- /dev/null +++ b/scripts/check_mcp_operation_boundary.py @@ -0,0 +1,65 @@ +import ast +import sys +from pathlib import Path +from typing import Final + +PACKAGE: Final = Path("litellm/proxy/_experimental/mcp_server") +LEGACY_ADAPTERS: Final = frozenset({"server.py", "legacy_callbacks.py", "mcp_context.py", "mcp_debug.py"}) +CONFINED_NAMES: Final = frozenset( + { + "auth_context_var", + "active_mcp_session_var", + "active_mcp_request_ctx_var", + "get_active_auth_context", + "get_active_mcp_session", + "get_active_mcp_request_ctx", + "get_or_extract_auth_context", + "_session_obj_auth_storage", + "WeakKeyDictionary", + "_mcp_active_toolset_id", + "_mcp_gateway_initialize_instructions", + "_mcp_gateway_server_name", + "_mcp_proxy_mode", + } +) + + +def is_confined(name: str) -> bool: + return name in CONFINED_NAMES or name.startswith("_stateful_session_") + + +def violations(path: Path, source: str) -> tuple[str, ...]: + if path.name in LEGACY_ADAPTERS: + return () + tree: Final = ast.parse(source, filename=str(path)) + return tuple( + f"{path}:{node.lineno}: MCP request/session state belongs in a legacy adapter" + for node in ast.walk(tree) + if ( + isinstance(node, ast.ImportFrom) + and ( + (node.module or "").endswith(".mcp_context") + or any(is_confined(alias.name) for alias in node.names) + or (path.name in {"operations.py", "contracts.py"} and (node.module or "").endswith(".server")) + ) + or isinstance(node, ast.Name) + and is_confined(node.id) + or isinstance(node, ast.Attribute) + and is_confined(node.attr) + ) + ) + + +def main() -> int: + findings: Final = tuple( + finding for path in sorted(PACKAGE.rglob("*.py")) for finding in violations(path, path.read_text()) + ) + if findings: + print("\n".join(findings), file=sys.stderr) + return 1 + print("MCP operation boundary: passed") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index 1abd415d237..22cc38f841c 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -102,6 +102,9 @@ ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|s ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$' litellm_py_files=$(scope_match "$litellm_py_pattern") +if [ -n "$(scope_match '^(litellm/proxy/_experimental/mcp_server/|scripts/check_mcp_operation_boundary\.py)')" ]; then + uv run --no-sync python scripts/check_mcp_operation_boundary.py || exit 1 +fi e2e_py_files=$(scope_match "$e2e_py_pattern") test_tree_files=$(scope_match "$test_tree_pattern") # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it. diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index ed8829945e5..41d0e2cb59b 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -142,7 +142,7 @@ async def test_mcp_cost_tracking(): local_mcp_server_manager, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", local_mcp_server_manager, ), ): @@ -293,7 +293,7 @@ async def test_mcp_cost_tracking_per_tool(): local_mcp_server_manager, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", local_mcp_server_manager, ), ): @@ -451,7 +451,7 @@ async def test_mcp_tool_call_hook(): local_mcp_server_manager, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", local_mcp_server_manager, ), ): diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 94cf35b675d..2b92367f186 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -922,7 +922,7 @@ async def test_get_tools_from_mcp_servers(): ) with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): # Test with specific servers @@ -950,6 +950,7 @@ async def test_get_tools_from_mcp_servers(): extra_headers=None, add_prefix=False, raw_headers=None, + client_ip=None, user_api_key_auth=None, oauth2_headers=None, ): @@ -966,7 +967,7 @@ async def test_get_tools_from_mcp_servers(): ) with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager_2, ): result = await _get_tools_from_mcp_servers( @@ -998,7 +999,7 @@ async def test_get_tools_from_mcp_servers(): ) with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): with patch( @@ -1981,6 +1982,7 @@ async def test_get_tools_for_single_server(): extra_headers=None, add_prefix=False, raw_headers=None, + client_ip=None, user_api_key_auth=None, ) @@ -2076,7 +2078,7 @@ async def test_rest_listing_hides_key_grants_dispatch_would_refuse(): with patch( "litellm.proxy._experimental.mcp_server.rest_endpoints.global_mcp_server_manager" ) as mock_manager, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager" + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager" ) as mock_server_manager, patch.object( MCPRequestHandler, "get_allowed_tools_for_server", @@ -2473,7 +2475,7 @@ async def test_filter_tools_by_allowed_tools_integration(): # Mock the global MCP server manager with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager" + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager" ) as mock_manager: # Mock manager methods mock_manager.get_allowed_mcp_servers = AsyncMock( @@ -2588,7 +2590,7 @@ async def test_filter_tools_by_disallowed_tools_integration(): # Mock the global MCP server manager with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager" + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager" ) as mock_manager: # Mock manager methods mock_manager.get_allowed_mcp_servers = AsyncMock( @@ -2689,7 +2691,7 @@ async def test_filter_tools_no_restrictions_integration(): # Mock the global MCP server manager with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager" + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager" ) as mock_manager: # Mock manager methods mock_manager.get_allowed_mcp_servers = AsyncMock( @@ -2970,10 +2972,10 @@ async def test_call_mcp_tool_uses_manager_permission_lookup(): return_value=mock_server, ) as mock_get_server, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_tool_registry" + "litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry" ) as mock_tool_registry, patch( - "litellm.proxy._experimental.mcp_server.server._handle_managed_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations._handle_managed_mcp_tool", new_callable=AsyncMock, ) as mock_handle_managed, patch( @@ -3046,10 +3048,10 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission return_value=mock_server, ) as mock_get_server, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_tool_registry" + "litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry" ) as mock_tool_registry, patch( - "litellm.proxy._experimental.mcp_server.server._handle_managed_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations._handle_managed_mcp_tool", new_callable=AsyncMock, ) as mock_handle_managed, patch( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index 87e23893616..a77b4c8d565 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -1,3 +1,4 @@ +from litellm.proxy._experimental.mcp_server import operations as mcp_operations """ Unit tests for the BYOK OAuth 2.1 authorization server endpoints. @@ -592,7 +593,7 @@ async def test_check_byok_credential_missing_credential(monkeypatch): monkeypatch.delenv("PROXY_BASE_URL", raising=False) monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) - server_module.byok_credential_cache.flush_cache() + mcp_operations.byok_credential_cache.flush_cache() mock_prisma = MagicMock() with ( @@ -628,13 +629,13 @@ async def test_execute_byok_tool_missing_credential_advertises_api_key_flow(monk from litellm.types.mcp_server.mcp_server_manager import MCPServer monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy") - mcp_module.byok_credential_cache.flush_cache() + mcp_operations.byok_credential_cache.flush_cache() server = MCPServer(server_id="byok-discovery", name="byok-discovery", transport=MCPTransport.http, is_byok=True) prisma = MagicMock() prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None) monkeypatch.setattr(proxy_server, "prisma_client", prisma) with pytest.raises(HTTPException) as exc_info: - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="list_regions", arguments={}, allowed_mcp_servers=[server], @@ -687,7 +688,7 @@ async def test_invalidate_byok_cred_cache_evicts_locally_and_broadcasts_the_same server = MCPServer(server_id="byok-revoke", name="byok-server", transport=MCPTransport.http, is_byok=True) user_auth = UserAPIKeyAuth(user_id="mallory", api_key="sk-test") - server_module.byok_credential_cache.flush_cache() + mcp_operations.byok_credential_cache.flush_cache() db_lookup = AsyncMock(side_effect=["sk-before-revoke", None]) publish = AsyncMock() @@ -699,13 +700,13 @@ async def test_invalidate_byok_cred_cache_evicts_locally_and_broadcasts_the_same "litellm.proxy.proxy_server.prisma_client", MagicMock() ), patch.object( # test-quality-ok: the redis publisher is module-level; asserting the broadcast without a redis - server_module, "publish_auth_cache_invalidation", new=publish + mcp_operations, "publish_auth_cache_invalidation", new=publish ), ): - assert await server_module._get_byok_credential(server, user_auth) == "sk-before-revoke" - assert await server_module._get_byok_credential(server, user_auth) == "sk-before-revoke" - await server_module._invalidate_byok_cred_cache("mallory", "byok-revoke") - assert await server_module._get_byok_credential(server, user_auth) is None + assert await mcp_operations._get_byok_credential(server, user_auth) == "sk-before-revoke" + assert await mcp_operations._get_byok_credential(server, user_auth) == "sk-before-revoke" + await mcp_operations._invalidate_byok_cred_cache("mallory", "byok-revoke") + assert await mcp_operations._get_byok_credential(server, user_auth) is None assert db_lookup.await_count == 2 publish.assert_awaited_once_with(cache_key=byok_credential_cache_key("mallory", "byok-revoke")) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_contracts.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_contracts.py new file mode 100644 index 00000000000..e13ecdfcce9 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_contracts.py @@ -0,0 +1,60 @@ +from dataclasses import FrozenInstanceError + +import pytest + +from litellm.proxy._experimental.mcp_server.operations import prepare_context +from litellm.proxy._types import UserAPIKeyAuth + + +def test_operation_context_isolates_nested_headers_and_caller_permissions(): + caller = UserAPIKeyAuth(user_id="alpha", models=["allowed"]) + caller.mcp_admitted_user_subject = True + caller.mcp_session_resource_server_id = "alpha-server" + caller.mcp_toolset_id = "toolset-alpha" + caller.mcp_source_team_rpm_limits = {"team": {"alpha-server": 2}} + headers = {"x-caller": "alpha"} + server_headers = {"alpha-server": {"authorization": "alpha-token"}} + context = prepare_context(caller, raw_headers=headers, mcp_server_auth_headers=server_headers) + + caller.models.append("forbidden") + caller.mcp_source_team_rpm_limits["team"]["alpha-server"] = 999 + headers["x-caller"] = "bravo" + server_headers["alpha-server"]["authorization"] = "bravo-token" + captured = context.user_api_key_auth + assert captured is not None + assert captured.models == ["allowed"] + assert captured.mcp_admitted_user_subject is True + assert captured.mcp_session_resource_server_id == "alpha-server" + assert captured.mcp_toolset_id == "toolset-alpha" + assert captured.mcp_source_team_rpm_limits == {"team": {"alpha-server": 2}} + captured.models.append("also-forbidden") + assert context.user_api_key_auth.models == ["allowed"] + assert context.raw_headers == {"x-caller": "alpha"} + assert context.mcp_server_auth_headers == {"alpha-server": {"authorization": "alpha-token"}} + with pytest.raises(TypeError): + context.raw_headers["x-caller"] = "changed" + with pytest.raises(TypeError): + context.mcp_server_auth_headers["alpha-server"]["authorization"] = "changed" + with pytest.raises(FrozenInstanceError): + context.client_ip = "untrusted" + + +def test_operation_context_preserves_missing_and_empty_inputs(): + missing = prepare_context() + empty = prepare_context(mcp_servers=[], raw_headers={}, oauth2_headers={}, mcp_server_auth_headers={}) + assert missing.user_api_key_auth is None + assert missing.mcp_servers is None + assert missing.raw_headers is None + assert missing.oauth2_headers is None + assert missing.mcp_server_auth_headers is None + assert empty.mcp_servers == () + assert empty.raw_headers == {} + assert empty.oauth2_headers == {} + assert empty.mcp_server_auth_headers == {} + + +def test_toolset_request_marker_cannot_be_supplied_by_caller_or_serialized(): + auth = UserAPIKeyAuth.model_validate({"user_id": "alpha", "mcp_toolset_id": "forged"}) + assert auth.mcp_toolset_id is None + auth.mcp_toolset_id = "server-resolved" + assert "mcp_toolset_id" not in auth.model_dump() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py index 64d926bc5e3..b8aadef430f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py @@ -1,5 +1,6 @@ +from litellm.proxy._experimental.mcp_server import operations as mcp_operations """Tests for guardrail-block recording in -``litellm.proxy._experimental.mcp_server.server.call_mcp_tool``. +``litellm.proxy._experimental.mcp_server.operations.call_mcp_tool``. A pre-call MCP guardrail block *raises* into ``call_mcp_tool``'s ``except Exception``. The failure spend-log row that the Guardrails Monitor's @@ -70,7 +71,7 @@ async def _call_block(logging_obj, order: list, *, user_api_key_auth=mock.sentin with mock.patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}): with contextlib.suppress(HTTPException): - await server.call_mcp_tool.__wrapped__( + await mcp_operations.call_mcp_tool.__wrapped__( name="t", arguments=None, user_api_key_auth=user_api_key_auth, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 28faf375ab8..9659eb1cbc2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1229,7 +1229,7 @@ class TestResolveByokMcpAuthHeader: user_auth = UserAPIKeyAuth(user_id="user-1", api_key="sk-dashboard") with patch( - "litellm.proxy._experimental.mcp_server.server._get_byok_credential", + "litellm.proxy._experimental.mcp_server.operations._get_byok_credential", new=AsyncMock(return_value="stored-cred"), ): result = await _resolve_byok_mcp_auth_header(server, user_auth, None) @@ -1249,7 +1249,7 @@ class TestResolveByokMcpAuthHeader: user_auth = UserAPIKeyAuth(user_id="user-1", api_key="sk-dashboard") with patch( - "litellm.proxy._experimental.mcp_server.server._get_byok_credential", + "litellm.proxy._experimental.mcp_server.operations._get_byok_credential", new=AsyncMock(return_value=None), ): with pytest.raises(HTTPException) as exc_info: @@ -1272,7 +1272,7 @@ class TestResolveByokMcpAuthHeader: check_mock = AsyncMock(return_value=None) with patch( - "litellm.proxy._experimental.mcp_server.server._check_byok_credential", + "litellm.proxy._experimental.mcp_server.operations._check_byok_credential", new=check_mock, ): result = await _resolve_byok_mcp_auth_header(server, user_auth, "caller-header") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index 3f5d4ad83ea..1909e3306a2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -1,3 +1,4 @@ +from litellm.proxy._experimental.mcp_server import operations as mcp_operations """Unit tests for MCP OAuth passthrough tool-fetch behavior.""" import logging @@ -339,16 +340,16 @@ async def test_aggregate_list_tools_absorbs_one_unauthenticated_server(): raise MCPUpstreamAuthError(status_code=401, www_authenticate=None, server_name=server.name) return [good_tool] - with patch.object(mcp_server, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate, working])), patch.object( - mcp_server, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={}) - ), patch.object(mcp_server, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object( - mcp_server, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None) + with patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate, working])), patch.object( + mcp_operations, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={}) + ), patch.object(mcp_operations, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object( + mcp_operations, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None) ), patch.object( - mcp_server, "filter_tools_by_key_team_permissions", AsyncMock(side_effect=lambda tools, **k: tools) + mcp_operations, "filter_tools_by_key_team_permissions", AsyncMock(side_effect=lambda tools, **k: tools) ), patch.object( - mcp_server.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools) + mcp_operations.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools) ): - listing = await mcp_server._get_tools_from_mcp_servers( + listing = await mcp_operations._get_tools_from_mcp_servers( user_api_key_auth=UserAPIKeyAuth(token="h", user_id="u1"), mcp_auth_header=None, mcp_servers=None, @@ -382,14 +383,14 @@ async def test_single_server_route_also_absorbs_upstream_auth_error(): # //mcp sets the path-derived single-server scope; absorption must hold even then. token = _mcp_gateway_server_name.set("delegate_docs") try: - with patch.object(mcp_server, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate])), patch.object( - mcp_server, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={}) - ), patch.object(mcp_server, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object( - mcp_server, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None) + with patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate])), patch.object( + mcp_operations, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={}) + ), patch.object(mcp_operations, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object( + mcp_operations, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None) ), patch.object( - mcp_server.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools) + mcp_operations.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools) ): - listing = await mcp_server._get_tools_from_mcp_servers( + listing = await mcp_operations._get_tools_from_mcp_servers( user_api_key_auth=UserAPIKeyAuth(token="h", user_id="u1"), mcp_auth_header=None, mcp_servers=["delegate_docs"], @@ -419,15 +420,15 @@ async def test_aggregate_with_single_accessible_server_still_absorbs(): async def fake_get_tools(server, **kwargs): raise MCPUpstreamAuthError(status_code=401, www_authenticate=None, server_name=server.name) - with patch.object(mcp_server, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate])), patch.object( - mcp_server, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={}) - ), patch.object(mcp_server, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object( - mcp_server, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None) + with patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[delegate])), patch.object( + mcp_operations, "_prefetch_oauth_creds_for_user", AsyncMock(return_value={}) + ), patch.object(mcp_operations, "_prepare_mcp_server_headers", MagicMock(return_value=(None, None))), patch.object( + mcp_operations, "_get_user_oauth_extra_headers_from_db", AsyncMock(return_value=None) ), patch.object( - mcp_server.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools) + mcp_operations.global_mcp_server_manager, "_get_tools_from_server", AsyncMock(side_effect=fake_get_tools) ): # Aggregate route: no explicit server filter, even though only one server is accessible. - listing = await mcp_server._get_tools_from_mcp_servers( + listing = await mcp_operations._get_tools_from_mcp_servers( user_api_key_auth=UserAPIKeyAuth(token="h", user_id="u1"), mcp_auth_header=None, mcp_servers=None, @@ -475,3 +476,25 @@ async def test_client_creation_failure_logs_sanitized_exchange(monkeypatch, capl await manager._get_tools_from_server(server) assert "POST https://upstream/ -> HTTP 500" in caplog.text assert "missing_scope" in caplog.text and "query-secret" not in caplog.text + + +@pytest.mark.parametrize( + "oauth_headers,server_headers,authorized", + [ + ({"Authorization": "Bearer upstream"}, None, True), + ({"AUTHORIZATION": "Bearer upstream"}, None, True), + ({"x-unrelated": "present"}, None, False), + (None, {"catalog": {"Authorization": "Bearer scoped"}}, True), + (None, {"other-server": {"Authorization": "Bearer unrelated"}}, False), + (None, {"catalog": {"x-unrelated": "present"}}, False), + (None, {"catalog": "Bearer legacy"}, True), + (None, {"catalog": " "}, False), + ], +) +def test_passthrough_admission_recognizes_only_matching_authorization(oauth_headers, server_headers, authorized): + from litellm.proxy._experimental.mcp_server.operations import _client_has_passthrough_authorization + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer(server_id="catalog", name="catalog", alias="catalog", transport=MCPTransport.http) + assert _client_has_passthrough_authorization(server, oauth_headers, server_headers) is authorized diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py index 84d4f1fd083..ed5d67164bd 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py @@ -1,3 +1,4 @@ +from litellm.proxy._experimental.mcp_server import operations as mcp_operations import json from datetime import datetime @@ -27,8 +28,8 @@ def proxy_mode(): @pytest.mark.asyncio @pytest.mark.usefixtures("proxy_mode") async def test_proxy_call_rejects_non_proxy_tool_names() -> None: - result = await server._dispatch_virtual_mcp_tool( - name="math_stdio-add", arguments={"a": 1, "b": 2}, user_api_key_auth=AUTH, client_ip=None + result = await mcp_operations._dispatch_virtual_mcp_tool( + name="math_stdio-add", arguments={"a": 1, "b": 2}, user_api_key_auth=AUTH, client_ip=None, mcp_proxy_mode=True ) assert result is not None @@ -105,12 +106,13 @@ async def test_proxy_scope_exception_emits_failure_log(monkeypatch: pytest.Monke arguments = {"tool_id": "denied-scope", "arguments": {}} with pytest.raises(HTTPException) as denied: - await server._dispatch_virtual_mcp_tool( + await mcp_operations._dispatch_virtual_mcp_tool( name="call_tool", arguments=arguments, user_api_key_auth=auth, client_ip=None, mcp_servers=["ungranted"], + mcp_proxy_mode=True, raw_headers={"authorization": "Bearer raw-scope-secret", "x-litellm-call-id": "scope-denial"}, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index b2eded67430..b715fe67e20 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,3 +1,4 @@ +from litellm.proxy._experimental.mcp_server import operations as mcp_operations import asyncio import contextlib import contextvars @@ -138,7 +139,7 @@ async def test_mcp_server_tool_call_body_contains_request_data(_mcp_request_ctx) mock_add_litellm_data_to_request, ): with patch( - "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.call_mcp_tool", mock_call_mcp_tool, ): with patch( @@ -194,7 +195,7 @@ async def test_mcp_server_tool_call_forwards_client_headers_to_logging(_mcp_requ mock_add_litellm_data_to_request, ): with patch( - "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.call_mcp_tool", mock_call_mcp_tool, ): with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): @@ -241,7 +242,7 @@ async def test_mcp_server_tool_call_strips_custom_litellm_key_header(_mcp_reques capturing_add_litellm_data_to_request, ): with patch( - "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.call_mcp_tool", mock_call_mcp_tool, ): with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): @@ -287,11 +288,11 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror(_mcp_r mock_add_litellm_data_to_request, ): with patch( - "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.call_mcp_tool", mock_call_mcp_tool, ): with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()): - with patch("litellm.proxy._experimental.mcp_server.server.verbose_logger", mock_logger): + with patch("litellm.proxy._experimental.mcp_server.operations.verbose_logger", mock_logger): result = await mcp_server_tool_call(_mcp_request_ctx(), _call_tool_params("test_tool", {"param": "value"})) assert result.is_error is True @@ -867,15 +868,15 @@ async def test_get_prompts_from_mcp_servers_success(): with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", AsyncMock(return_value=[server_a, server_b]), ) as mock_allowed, patch( - "litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers", + "litellm.proxy._experimental.mcp_server.operations._prepare_mcp_server_headers", return_value=(None, None), ) as mock_headers, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", ) as mock_manager, ): mock_manager.get_prompts_from_server = AsyncMock( @@ -927,15 +928,15 @@ async def test_get_resources_from_mcp_servers_success(): with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", AsyncMock(return_value=[server_a, server_b]), ) as mock_allowed, patch( - "litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers", + "litellm.proxy._experimental.mcp_server.operations._prepare_mcp_server_headers", return_value=(None, None), ) as mock_headers, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", ) as mock_manager, ): mock_manager.get_resources_from_server = AsyncMock( @@ -992,15 +993,15 @@ async def test_get_resource_templates_from_mcp_servers_success(): with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", AsyncMock(return_value=[server]), ) as mock_allowed, patch( - "litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers", + "litellm.proxy._experimental.mcp_server.operations._prepare_mcp_server_headers", return_value=(None, None), ) as mock_headers, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", ) as mock_manager, ): mock_manager.get_resource_templates_from_server = AsyncMock( @@ -1042,15 +1043,15 @@ async def test_mcp_get_prompt_success(): with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", AsyncMock(return_value=[server]), ) as mock_allowed, patch( - "litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers", + "litellm.proxy._experimental.mcp_server.operations._prepare_mcp_server_headers", return_value=({"Authorization": "token"}, {"X-Test": "1"}), ) as mock_headers, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", ) as mock_manager, ): mock_manager.get_prompt_from_server = AsyncMock(return_value=prompt_result) @@ -1078,6 +1079,7 @@ async def test_mcp_get_prompt_success(): mcp_auth_header={"Authorization": "token"}, extra_headers={"X-Test": "1"}, raw_headers=None, + client_ip=None, ) assert result is prompt_result @@ -1106,15 +1108,15 @@ async def test_mcp_read_resource_success(): with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", AsyncMock(return_value=[server]), ) as mock_allowed, patch( - "litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers", + "litellm.proxy._experimental.mcp_server.operations._prepare_mcp_server_headers", return_value=({"Authorization": "token"}, {"X-Test": "1"}), ) as mock_headers, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", ) as mock_manager, ): mock_manager.read_resource_from_server = AsyncMock(return_value=read_result) @@ -1140,6 +1142,7 @@ async def test_mcp_read_resource_success(): mcp_auth_header={"Authorization": "token"}, extra_headers={"X-Test": "1"}, raw_headers=None, + client_ip=None, ) assert result is read_result @@ -1264,7 +1267,7 @@ async def test_mcp_read_resource_multiple_servers_error(): server_b.name = "server_b" with patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", AsyncMock(return_value=[server_a, server_b]), ) as mock_allowed: with pytest.raises(HTTPException) as exc_info: @@ -1354,11 +1357,11 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): mock_manager._get_tools_from_server = mock_get_tools_from_server with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): with patch( - "litellm.proxy._experimental.mcp_server.server.verbose_logger", + "litellm.proxy._experimental.mcp_server.operations.verbose_logger", ) as mock_logger: # Test with server-specific auth headers mcp_server_auth_headers = { @@ -1450,11 +1453,11 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): mock_manager._get_tools_from_server = mock_get_tools_from_server with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): with patch( - "litellm.proxy._experimental.mcp_server.server.verbose_logger", + "litellm.proxy._experimental.mcp_server.operations.verbose_logger", ) as mock_logger: # Test with server-specific auth headers mcp_server_auth_headers = { @@ -1524,11 +1527,11 @@ async def _denied_scoped_list( with ( patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", resolver, ), patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ), ): @@ -1575,11 +1578,11 @@ async def test_empty_scope_lists_nothing_instead_of_raising_a_nameless_denial(): with ( patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", resolver, ), patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", _denied_scope_manager({"github": "srv-github"}), ), ): @@ -1721,7 +1724,8 @@ async def test_scoped_list_agent_veto_attributed_for_differently_cased_server_na @pytest.mark.asyncio -async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error(_mcp_request_ctx): +@pytest.mark.parametrize("denial_at_auth", [False, True]) +async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error(_mcp_request_ctx, denial_at_auth): """The MCP protocol handler surfaces a permission HTTPException as a clean JSON-RPC error (MCPError, INVALID_REQUEST) carrying the denial message, instead of a raw 500.""" try: @@ -1738,10 +1742,10 @@ async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error( with ( patch( # test-quality-ok: the protocol handler reads auth from module context; no injection seam "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", - new=AsyncMock(return_value=(None, None, None, None, None, None, None)), + new=AsyncMock(return_value=(None, None, None, None, None, None, None), side_effect=denial if denial_at_auth else None), ), patch( # test-quality-ok: the listing helper is the handler's only collaborator; the suite's seam - "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + "litellm.proxy._experimental.mcp_server.operations._list_mcp_tools", new=AsyncMock(side_effect=denial), ), ): @@ -1768,7 +1772,7 @@ async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict(_mcp_ new=AsyncMock(return_value=(None, None, None, None, None, None, None)), ), patch( # test-quality-ok: the tool-call helper is the handler's only collaborator; the suite's seam - "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.call_mcp_tool", new=AsyncMock(side_effect=denial), ), ): @@ -1819,7 +1823,7 @@ async def test_mcp_server_tool_call_body_with_none_arguments(_mcp_request_ctx): mock_add_litellm_data_to_request, ): with patch( - "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.call_mcp_tool", mock_call_mcp_tool, ): with patch( @@ -1893,7 +1897,7 @@ async def test_concurrent_initialize_session_managers(): "run", return_value=mock_cm_sse, ) as mock_sse_run, - patch("litellm.proxy._experimental.mcp_server.server.verbose_logger"), + patch("litellm.proxy._experimental.mcp_server.operations.verbose_logger"), ): # Create multiple concurrent tasks that call initialize_session_managers async def init_task(): @@ -2334,7 +2338,7 @@ async def test_mcp_routing_chunked_initialize_to_stateful(): "litellm.proxy._experimental.mcp_server.server.set_auth_context", ), patch( # test-quality-ok: registry is empty in unit tests; key owns one server - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[MagicMock()], ), @@ -2886,7 +2890,7 @@ async def test_initialize_request_tracks_active_session_after_response_header(): return_value=(owner_auth, None, None, None, None, None), ), patch( # test-quality-ok: registry is empty in unit tests; key owns one server - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[MagicMock()], ), @@ -3039,7 +3043,7 @@ async def test_initialize_request_records_client_name_in_gateway_sessions_report return_value=(owner_auth, None, None, None, None, None), ), patch( # test-quality-ok: registry is empty in unit tests; key owns one server - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[MagicMock()], ), @@ -3514,7 +3518,7 @@ async def test_initialize_request_with_existing_session_tracks_new_session(): ), ), patch( # test-quality-ok: registry is empty in unit tests; key owns one server - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[MagicMock()], ), @@ -4248,7 +4252,7 @@ async def test_mcp_routing_with_conflicting_alias_and_group_name(): with ( patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_allowed_mcp_servers", mock_get_allowed, ), patch( @@ -4256,7 +4260,7 @@ async def test_mcp_routing_with_conflicting_alias_and_group_name(): mock_db_lookup, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager._get_tools_from_server", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager._get_tools_from_server", mock_get_tools_spy, ), ): @@ -4365,16 +4369,16 @@ async def test_oauth2_caller_headers_not_forwarded_for_migrated_server(): side_effect=mock_fetch_tools_with_timeout, ), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", AsyncMock(return_value=[oauth2_server]), ), patch( - "litellm.proxy._experimental.mcp_server.server._prefetch_oauth_creds_for_user", + "litellm.proxy._experimental.mcp_server.operations._prefetch_oauth_creds_for_user", new_callable=AsyncMock, return_value={}, ), patch( - "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + "litellm.proxy._experimental.mcp_server.operations._get_user_oauth_extra_headers_from_db", new_callable=AsyncMock, return_value=None, ), @@ -4456,7 +4460,7 @@ async def test_list_tools_single_server_unprefixed_names(): mock_manager._get_tools_from_server = mock_get_tools_from_server with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): listing = await _get_tools_from_mcp_servers( @@ -4535,7 +4539,7 @@ async def test_list_tools_multiple_servers_prefixed_names(): mock_manager._get_tools_from_server = mock_get_tools_from_server with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): listing = await _get_tools_from_mcp_servers( @@ -4715,7 +4719,7 @@ async def test_call_mcp_tool_user_unauthorized_access(): AsyncMock(return_value=["allowed_server", "another_server"]), ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_id", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_id", side_effect=mock_get_server_by_id, ), ): @@ -4745,11 +4749,11 @@ async def test_call_mcp_tool_scoped_denial_names_the_binding_agent(): with ( patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_allowed_mcp_servers", AsyncMock(return_value=[]), ), patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", _scope_resolver({"github": "srv-github"}), ), ): @@ -4821,7 +4825,7 @@ async def test_call_mcp_tool_unauthorized_403_does_not_leak_server_credentials() AsyncMock(return_value=["allowed_server"]), ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_id", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_id", side_effect=mock_get_server_by_id, ), ): @@ -4964,7 +4968,7 @@ async def test_list_tools_filters_by_key_team_permissions(): mock_manager._get_tools_from_server = mock_get_tools_from_server with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): listing = await _get_tools_from_mcp_servers( @@ -5075,7 +5079,7 @@ async def test_list_tools_with_team_tool_permissions_inheritance(): mock_manager._get_tools_from_server = mock_get_tools_from_server with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): # Mock the team object permission retrieval @@ -5167,7 +5171,7 @@ async def test_list_tools_with_no_tool_permissions_shows_all(): mock_manager._get_tools_from_server = mock_get_tools_from_server with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): listing = await _get_tools_from_mcp_servers( @@ -5273,7 +5277,7 @@ async def test_list_tools_strips_prefix_when_matching_permissions(): mock_manager._get_tools_from_server = mock_get_tools_from_server with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): listing = await _get_tools_from_mcp_servers( @@ -5715,12 +5719,12 @@ async def test_call_mcp_tool_logs_failure_via_post_call_failure_hook(): return_value=mock_server, ), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers_from_mcp_server_names", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers_from_mcp_server_names", new_callable=AsyncMock, return_value=[mock_server], ), patch( - "litellm.proxy._experimental.mcp_server.server.execute_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.execute_mcp_tool", new_callable=AsyncMock, side_effect=Exception("boom"), ), @@ -5784,26 +5788,26 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=[server_a]), ), patch( - "litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers", + "litellm.proxy._experimental.mcp_server.operations._prepare_mcp_server_headers", return_value=(None, None), ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", ) as mock_manager, patch( - "litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools", + "litellm.proxy._experimental.mcp_server.operations.filter_tools_by_allowed_tools", side_effect=lambda tools, _server: tools, ), patch( - "litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions", + "litellm.proxy._experimental.mcp_server.operations.filter_tools_by_key_team_permissions", new=AsyncMock(side_effect=lambda tools, **_: tools), ), patch( - "litellm.proxy._experimental.mcp_server.server.function_setup", + "litellm.proxy._experimental.mcp_server.operations.function_setup", side_effect=_capture_function_setup, ), ): @@ -5866,26 +5870,26 @@ async def test_get_tools_from_mcp_servers_returns_tools_when_success_logging_fai with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=[server_a]), ), patch( - "litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers", + "litellm.proxy._experimental.mcp_server.operations._prepare_mcp_server_headers", return_value=(None, None), ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", ) as mock_manager, patch( - "litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools", + "litellm.proxy._experimental.mcp_server.operations.filter_tools_by_allowed_tools", side_effect=lambda tools, _server: tools, ), patch( - "litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions", + "litellm.proxy._experimental.mcp_server.operations.filter_tools_by_key_team_permissions", new=AsyncMock(side_effect=lambda tools, **_: tools), ), patch( - "litellm.proxy._experimental.mcp_server.server.function_setup", + "litellm.proxy._experimental.mcp_server.operations.function_setup", return_value=(dummy_logging_obj, None), ), ): @@ -6186,23 +6190,23 @@ async def test_get_tools_from_mcp_servers_injects_stored_oauth2_token(): with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=[oauth2_server]), ), patch( # Patch the bulk prefetch so no real DB connection is needed - "litellm.proxy._experimental.mcp_server.server._prefetch_oauth_creds_for_user", + "litellm.proxy._experimental.mcp_server.operations._prefetch_oauth_creds_for_user", new=AsyncMock(return_value=prefetched_creds), ) as mock_prefetch, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", ) as mock_manager, patch( - "litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools", + "litellm.proxy._experimental.mcp_server.operations.filter_tools_by_allowed_tools", side_effect=lambda tools, _server: tools, ), patch( - "litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions", + "litellm.proxy._experimental.mcp_server.operations.filter_tools_by_key_team_permissions", new=AsyncMock(side_effect=lambda tools, **_: tools), ), ): @@ -6534,7 +6538,7 @@ class TestGatewayCreateInitializationOptions: with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[scoped_server], ), @@ -6564,7 +6568,7 @@ class TestGatewayCreateInitializationOptions: from litellm.proxy._types import UserAPIKeyAuth with patch( # test-quality-ok: grant resolution is the input under test - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[], ): @@ -6590,7 +6594,7 @@ class TestGatewayCreateInitializationOptions: from litellm.proxy._types import UserAPIKeyAuth with patch( # test-quality-ok: grant resolution is the input under test - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[], ): @@ -6615,7 +6619,7 @@ class TestGatewayCreateInitializationOptions: from litellm.proxy._types import UserAPIKeyAuth with patch( # test-quality-ok: grant resolution is the input under test - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[], ): @@ -6671,7 +6675,7 @@ class TestGatewayCreateInitializationOptions: ), ), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[scoped_server], ), @@ -6806,14 +6810,14 @@ async def test_list_tools_with_legacy_db_m2m_server_resolves_oauth2_flow(): with ( patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", ) as mock_manager, patch( - "litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools", + "litellm.proxy._experimental.mcp_server.operations.filter_tools_by_allowed_tools", side_effect=lambda tools, _server: tools, ), patch( - "litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions", + "litellm.proxy._experimental.mcp_server.operations.filter_tools_by_key_team_permissions", new=AsyncMock(side_effect=lambda tools, **_: tools), ), ): @@ -7076,7 +7080,7 @@ def _patch_delegate_resolver(server: MCPServer, *resolvable_names: str): return server if name in resolvable_names else None return patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", side_effect=_resolve, ) @@ -7095,7 +7099,7 @@ async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-qual with ( _patch_delegate_resolver(server, "delegate_test"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=[server]), ), patch( @@ -7131,7 +7135,7 @@ async def test_legacy_delegate_dual_credentials_are_not_probed_upstream(): # te with ( patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=[server]), ), patch( # test-quality-ok: the removed probe call is the security regression under test @@ -7178,7 +7182,7 @@ async def test_oauth_passthrough_preflight_preserves_status_contract(probe_statu with ( patch( # test-quality-ok: isolate authorized-server resolution so this test exercises the preflight contract - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=[server]), ), patch( # test-quality-ok: the upstream transport boundary is the behavior being mapped to an HTTP response @@ -7224,7 +7228,7 @@ async def test_delegate_tokenless_request_not_probed(): with ( _patch_delegate_resolver(server, "delegate_test"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=[server]), ), patch( @@ -7257,7 +7261,7 @@ async def test_delegate_preflight_skipped_on_multi_server_routes(): with ( _patch_delegate_resolver(servers[0], "delegate_test", "other_server"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=servers), ), patch( @@ -7300,7 +7304,7 @@ async def test_bare_authorization_never_probes_passthrough_servers(): with ( _patch_delegate_resolver(passthrough_server, "pt_server"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=[passthrough_server]), ), patch( @@ -7346,7 +7350,7 @@ async def test_delegate_not_probed_when_named_only_via_server_id(): with ( _patch_delegate_resolver(server, "delegate_test"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=[server]), ), patch( @@ -7379,7 +7383,7 @@ async def test_delegate_probe_not_fanned_out_to_access_group_members(): with ( _patch_delegate_resolver(group_member, "delegate_test"), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(return_value=[group_member]), ), patch( @@ -7475,11 +7479,11 @@ async def test_execute_mcp_tool_rest_server_id_authoritative_for_unprefixed_tool with ( patch.dict( - mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_name_mapping, + mcp_operations.global_mcp_server_manager.tool_name_to_mcp_server_name_mapping, {"echo": oauth_server.name}, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "get_registry", return_value={ api_key_server.server_id: api_key_server, @@ -7487,13 +7491,12 @@ async def test_execute_mcp_tool_rest_server_id_authoritative_for_unprefixed_tool }, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=oauth_server, ), patch.object( - mcp_module, - "_handle_managed_mcp_tool", + mcp_operations, "_handle_managed_mcp_tool", new=fake_handle_managed_mcp_tool, ), patch.object( @@ -7502,12 +7505,12 @@ async def test_execute_mcp_tool_rest_server_id_authoritative_for_unprefixed_tool return_value=True, ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=None, ), ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="echo", arguments={"message": "hello"}, allowed_mcp_servers=[api_key_server, oauth_server], @@ -7540,7 +7543,7 @@ def _worker_that_never_listed(server: MCPServer, upstream_tools: tuple[str, ...] from litellm.proxy._experimental.mcp_server import server as mcp_module - mcp_module.global_mcp_server_manager.registry[server.server_id] = server + mcp_operations.global_mcp_server_manager.registry[server.server_id] = server dispatched: dict[str, object] = {} async def fake_handle_managed_mcp_tool(**kwargs): @@ -7552,17 +7555,17 @@ def _worker_that_never_listed(server: MCPServer, upstream_tools: tuple[str, ...] with ( patch.object( # test-quality-ok: the upstream MCP session is the boundary; a real one needs an initialize handshake over a live server - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_create_mcp_client", new=AsyncMock(return_value=MagicMock()), ) as create_client, patch.object( # test-quality-ok: same boundary, this is the tools/list answer the upstream would give - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_fetch_tools_with_timeout", side_effect=fake_fetch_tools, ) as fetch_tools, patch.object( # test-quality-ok: records the resolved server and bare name the managed call would forward upstream - mcp_module, "_handle_managed_mcp_tool", new=fake_handle_managed_mcp_tool + mcp_operations, "_handle_managed_mcp_tool", new=fake_handle_managed_mcp_tool ), ): yield SimpleNamespace(create_client=create_client, fetch_tools=fetch_tools, dispatched=dispatched) @@ -7576,7 +7579,7 @@ async def test_execute_mcp_tool_lists_never_listed_passthrough_server_with_calle server = _never_listed_passthrough_server() with _worker_that_never_listed(server, upstream_tools=("add",)) as worker: - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="lazy_map-add", arguments={"a": 1, "b": 2}, allowed_mcp_servers=[server], @@ -7597,7 +7600,7 @@ async def test_execute_mcp_tool_rest_server_id_lists_never_listed_server_first() server = _never_listed_passthrough_server() with _worker_that_never_listed(server, upstream_tools=("add",)) as worker: - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="add", arguments={"a": 1, "b": 2}, allowed_mcp_servers=[server], @@ -7620,7 +7623,7 @@ async def test_execute_mcp_tool_unknown_tool_on_never_listed_server_lists_once_t _worker_that_never_listed(server, upstream_tools=("add",)) as worker, pytest.raises(HTTPException) as exc_info, ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="lazy_map-nope", arguments={}, allowed_mcp_servers=[server], @@ -7641,8 +7644,8 @@ async def test_execute_mcp_tool_does_not_relist_a_server_this_worker_already_lis server = _never_listed_passthrough_server() with _worker_that_never_listed(server, upstream_tools=("add",)) as worker: - mcp_module.global_mcp_server_manager._create_prefixed_tools([MCPTool(name="add", inputSchema={})], server) - await mcp_module.execute_mcp_tool( + mcp_operations.global_mcp_server_manager._create_prefixed_tools([MCPTool(name="add", inputSchema={})], server) + await mcp_operations.execute_mcp_tool( name="lazy_map-add", arguments={"a": 1, "b": 2}, allowed_mcp_servers=[server], @@ -7664,8 +7667,8 @@ async def test_execute_mcp_tool_lists_a_tool_this_worker_has_not_yet_seen_on_a_l server = _never_listed_passthrough_server() with _worker_that_never_listed(server, upstream_tools=("add", "multiply")) as worker: - mcp_module.global_mcp_server_manager._create_prefixed_tools([MCPTool(name="add", inputSchema={})], server) - await mcp_module.execute_mcp_tool( + mcp_operations.global_mcp_server_manager._create_prefixed_tools([MCPTool(name="add", inputSchema={})], server) + await mcp_operations.execute_mcp_tool( name="lazy_map-multiply", arguments={"a": 1, "b": 2}, allowed_mcp_servers=[server], @@ -7688,7 +7691,7 @@ async def test_execute_mcp_tool_never_lists_a_server_the_caller_cannot_access(): _worker_that_never_listed(server, upstream_tools=("add",)) as worker, pytest.raises(HTTPException) as exc_info, ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="lazy_map-add", arguments={"a": 1, "b": 2}, allowed_mcp_servers=[other_server], @@ -7734,13 +7737,12 @@ async def test_execute_mcp_tool_strips_a_prefix_that_contains_the_separator(): with ( patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=alias_less_server, ), patch.object( - mcp_module, - "_handle_managed_mcp_tool", + mcp_operations, "_handle_managed_mcp_tool", new=fake_handle_managed_mcp_tool, ), patch.object( @@ -7749,12 +7751,12 @@ async def test_execute_mcp_tool_strips_a_prefix_that_contains_the_separator(): return_value=True, ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=None, ), ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name=f"{server_id}-read_wiki_contents", arguments={"repoName": "acme/wiki"}, allowed_mcp_servers=[alias_less_server], @@ -7808,11 +7810,11 @@ async def test_execute_mcp_tool_rest_server_id_injects_requested_server_credenti with ( patch.dict( - mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_name_mapping, + mcp_operations.global_mcp_server_manager.tool_name_to_mcp_server_name_mapping, {"echo": collision_server.name, "echo_requested-echo": requested_server.name}, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "get_registry", return_value={ requested_server.server_id: requested_server, @@ -7820,7 +7822,7 @@ async def test_execute_mcp_tool_rest_server_id_injects_requested_server_credenti }, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_create_mcp_client", new=fake_create_mcp_client, ), @@ -7830,13 +7832,13 @@ async def test_execute_mcp_tool_rest_server_id_injects_requested_server_credenti return_value=True, ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=None, ), patch("litellm.proxy.proxy_server.proxy_logging_obj", None), ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="echo", arguments={"message": "hello"}, allowed_mcp_servers=[requested_server, collision_server], @@ -7879,7 +7881,7 @@ async def test_execute_mcp_tool_rest_prefixed_tool_still_validates_server_id(): with ( patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "get_registry", return_value={ api_key_server.server_id: api_key_server, @@ -7887,7 +7889,7 @@ async def test_execute_mcp_tool_rest_prefixed_tool_still_validates_server_id(): }, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=oauth_server, ), @@ -7897,13 +7899,13 @@ async def test_execute_mcp_tool_rest_prefixed_tool_still_validates_server_id(): return_value=True, ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=None, ), pytest.raises(HTTPException) as exc_info, ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="echo_oauth_m2m-echo", arguments={"message": "hello"}, allowed_mcp_servers=[api_key_server, oauth_server], @@ -7941,7 +7943,7 @@ async def test_execute_mcp_tool_rest_unauthorized_prefix_still_mismatches(): with ( patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "get_registry", return_value={ api_key_server.server_id: api_key_server, @@ -7949,7 +7951,7 @@ async def test_execute_mcp_tool_rest_unauthorized_prefix_still_mismatches(): }, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=restricted_server, ), @@ -7959,13 +7961,13 @@ async def test_execute_mcp_tool_rest_unauthorized_prefix_still_mismatches(): return_value=True, ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=None, ), pytest.raises(HTTPException) as exc_info, ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="restricted_server-echo", arguments={"message": "hello"}, allowed_mcp_servers=[api_key_server], @@ -8005,18 +8007,17 @@ async def test_execute_mcp_tool_rest_hyphenated_upstream_tool_name_routes_to_req with ( patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "get_registry", return_value={api_key_server.server_id: api_key_server}, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=None, ), patch.object( - mcp_module, - "_handle_managed_mcp_tool", + mcp_operations, "_handle_managed_mcp_tool", new=fake_handle_managed_mcp_tool, ), patch.object( @@ -8025,12 +8026,12 @@ async def test_execute_mcp_tool_rest_hyphenated_upstream_tool_name_routes_to_req return_value=True, ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=None, ), ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="text-to-speech", arguments={"message": "hello"}, allowed_mcp_servers=[api_key_server], @@ -8089,22 +8090,22 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details(): with ( patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=fake_server, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "pre_call_tool_check", new=AsyncMock(return_value={}), ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=fake_tool, ), patch( - "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool", new=AsyncMock(return_value=[]), ), patch( @@ -8112,7 +8113,7 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details(): return_value=True, ), ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="list_pets", arguments={"limit": 10}, allowed_mcp_servers=[fake_server], @@ -8168,7 +8169,7 @@ async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requeste with ( patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "get_registry", return_value={ requested_server.server_id: requested_server, @@ -8176,13 +8177,12 @@ async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requeste }, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=None, ), patch.object( - mcp_module, - "_handle_managed_mcp_tool", + mcp_operations, "_handle_managed_mcp_tool", new=fake_handle_managed_mcp_tool, ), patch.object( @@ -8191,12 +8191,12 @@ async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requeste return_value=True, ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=None, ), ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="known_prefix-list_things", arguments={"message": "hello"}, allowed_mcp_servers=[requested_server, prefix_owner], @@ -8248,7 +8248,7 @@ async def test_execute_mcp_tool_rest_prefix_retry_resolution_still_enforces_serv with ( patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "get_registry", return_value={ requested_server.server_id: requested_server, @@ -8256,7 +8256,7 @@ async def test_execute_mcp_tool_rest_prefix_retry_resolution_still_enforces_serv }, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", side_effect=resolve_only_when_requested_prefix_added, ), @@ -8266,13 +8266,13 @@ async def test_execute_mcp_tool_rest_prefix_retry_resolution_still_enforces_serv return_value=True, ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=None, ), pytest.raises(HTTPException) as exc_info, ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="known_prefix-echo", arguments={"message": "hello"}, allowed_mcp_servers=[requested_server, prefix_owner], @@ -9175,14 +9175,14 @@ async def test_call_tool_with_legacy_db_m2m_server_resolves_oauth2_flow(): with ( patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", ) as mock_manager, patch( - "litellm.proxy._experimental.mcp_server.server.execute_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.execute_mcp_tool", side_effect=capture_execute, ), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers_from_mcp_server_names", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers_from_mcp_server_names", new=AsyncMock(side_effect=lambda mcp_servers, allowed_mcp_servers: allowed_mcp_servers), ), ): @@ -9260,12 +9260,12 @@ async def test_call_mcp_tool_skips_failure_hook_for_upstream_auth_error(): ), patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=mock_server), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers_from_mcp_server_names", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers_from_mcp_server_names", new_callable=AsyncMock, return_value=[mock_server], ), patch( - "litellm.proxy._experimental.mcp_server.server.execute_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.execute_mcp_tool", new_callable=AsyncMock, side_effect=MCPUpstreamAuthError(status_code=401, www_authenticate="Bearer", server_name="test_server"), ), @@ -9345,7 +9345,7 @@ async def test_aggregate_listing_reports_per_server_outcomes(): mock_manager._get_tools_from_server = mock_get_tools_from_server with patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ): listing = await _get_tools_from_mcp_servers( @@ -9419,7 +9419,7 @@ async def test_handle_list_tools_attaches_outcome_meta(_mcp_request_ctx): new=AsyncMock(return_value=(None, None, None, None, None, None, None)), ), patch( - "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + "litellm.proxy._experimental.mcp_server.operations._list_mcp_tools", new=AsyncMock(return_value=listing), ), ): @@ -9485,12 +9485,12 @@ class TestPreemptive401ModeAware: with ( patch.object( - server_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "get_mcp_server_by_name", return_value=server, ), patch.object( - server_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "has_user_oauth_token", new_callable=AsyncMock, return_value=has_stored_token, @@ -9509,7 +9509,7 @@ class TestPreemptive401ModeAware: async def test_deferred_discovery_runs_before_delegate_challenge(self): from litellm.proxy._experimental.mcp_server import server as server_module - manager = server_module.global_mcp_server_manager + manager = mcp_operations.global_mcp_server_manager server = _make_oauth2_server( "lazy_delegate", oauth2_flow="authorization_code", @@ -9541,7 +9541,7 @@ class TestPreemptive401ModeAware: async def test_stamped_m2m_challenge_skips_deferred_discovery(self): from litellm.proxy._experimental.mcp_server import server as server_module - manager = server_module.global_mcp_server_manager + manager = mcp_operations.global_mcp_server_manager server = _make_oauth2_server("stamped_m2m", oauth2_flow="client_credentials") with patch.object( @@ -9584,12 +9584,12 @@ class TestPreemptive401ModeAware: with ( patch.dict(os.environ, {"SERVER_ROOT_PATH": "/litellm"}), patch.object( - server_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "get_mcp_server_by_name", return_value=server, ), patch.object( - server_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "has_user_oauth_token", new_callable=AsyncMock, return_value=False, @@ -9691,17 +9691,17 @@ class TestSingleServerPreflightReachesIdJag: with ( patch.object( # test-quality-ok: route wiring must use the manager's configured server - server_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "get_mcp_server_by_name", return_value=server, ), patch.object( # test-quality-ok: route wiring must invoke the manager preflight - server_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "preflight_token_exchange", preflight, ), patch.object( # test-quality-ok: allowed-set resolution needs the DB; the test controls its answer - server_module, "_get_allowed_mcp_servers", AsyncMock(return_value=[server]) + mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server]) ), ): await server_module._raise_preemptive_401_for_unauthenticated_servers( @@ -9751,12 +9751,12 @@ class TestSingleServerPreflightReachesIdJag: with ( patch.object( # test-quality-ok: route wiring must use the manager's configured server - server_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "get_mcp_server_by_name", return_value=token_exchange, ), patch.object( # test-quality-ok: route wiring must invoke the manager preflight - server_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "preflight_token_exchange", preflight, ), @@ -9817,13 +9817,13 @@ class TestOboPreflightScopedToAllowedServers: preflight = AsyncMock() with ( patch.object( # test-quality-ok: route handler reads the module-level manager, no injection seam - server_module.global_mcp_server_manager, "get_mcp_server_by_name", return_value=requested + mcp_operations.global_mcp_server_manager, "get_mcp_server_by_name", return_value=requested ), patch.object( # test-quality-ok: the exchanger is the observable; a real one would call an IdP - server_module.global_mcp_server_manager, "preflight_token_exchange", preflight + mcp_operations.global_mcp_server_manager, "preflight_token_exchange", preflight ), patch.object( # test-quality-ok: allowed-set resolution needs the DB; the test controls its answer - server_module, "_get_allowed_mcp_servers", allowed_lookup + mcp_operations, "_get_allowed_mcp_servers", allowed_lookup ), ): await server_module._raise_preemptive_401_for_unauthenticated_servers( @@ -10132,7 +10132,7 @@ class TestListFiltersHonorThePrefixBoundary: with ( patch.object(MCPRequestHandler, "get_allowed_tools_for_server", AsyncMock(return_value=grants)), - patch("litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager") as mock_manager, + patch("litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager") as mock_manager, ): mock_manager.get_mcp_server_by_id.return_value = server @@ -10195,11 +10195,11 @@ async def test_list_tools_injects_byok_credential_for_non_oauth2_auth_types(auth with ( patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", mock_manager, ), patch( - "litellm.proxy._experimental.mcp_server.server._get_byok_credential", + "litellm.proxy._experimental.mcp_server.operations._get_byok_credential", AsyncMock(return_value="personal-api-key"), ), ): @@ -10292,3 +10292,44 @@ async def test_streamable_http_rejects_modern_protocol_version(header_value: str assert header_value in body["error"]["message"] for version in body["error"]["message"].split("supported: ")[1].split(", "): assert version in HANDSHAKE_PROTOCOL_VERSIONS + + +@pytest.mark.asyncio +@pytest.mark.parametrize("handler_name,field", [ + ("handle_list_tools", "tools"), + ("list_prompts", "prompts"), + ("list_resources", "resources"), + ("list_resource_templates", "resource_templates"), +]) +async def test_native_listing_preserves_empty_result_on_auth_failure(_mcp_request_ctx, handler_name, field): + from litellm.proxy._experimental.mcp_server import server + + with patch.object(server, "get_or_extract_auth_context", AsyncMock(side_effect=RuntimeError("auth failure"))): + result = await getattr(server, handler_name)(_mcp_request_ctx(), _paged_params()) + assert getattr(result, field) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_hook_raises", [False, True]) +async def test_tool_listing_preserves_permission_denial_when_failure_logging_fails(failure_hook_raises): + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy import proxy_server + + auth = UserAPIKeyAuth(user_id="denied-caller") + denial = HTTPException(status_code=403, detail="scope denied") + logger = MagicMock() + logger.post_call_failure_hook = AsyncMock(side_effect=RuntimeError("log unavailable") if failure_hook_raises else None) + upstream = AsyncMock() + with ( + patch.object(operations, "_get_allowed_mcp_servers", AsyncMock(side_effect=denial)), + patch.object(operations, "function_setup", return_value=(None, None)), + patch.object(proxy_server, "proxy_logging_obj", logger), + patch.object(operations.global_mcp_server_manager, "_get_tools_from_server", upstream), + ): + with pytest.raises(HTTPException) as rejected: + await operations._get_tools_from_mcp_servers(user_api_key_auth=auth, mcp_auth_header=None, mcp_servers=["catalog"], log_list_tools_to_spendlogs=True) + assert rejected.value is denial + upstream.assert_not_awaited() + logger.post_call_failure_hook.assert_awaited_once() + assert logger.post_call_failure_hook.await_args.kwargs["original_exception"] is denial + assert logger.post_call_failure_hook.await_args.kwargs["user_api_key_dict"] == auth diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ac23831b5f4..9f42a523350 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -73,6 +73,135 @@ from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +@pytest.mark.asyncio +async def test_manager_sampling_preserves_explicit_headers_without_ambient_context(): + from litellm.proxy._experimental.mcp_server import server as legacy_server + + caller = UserAPIKeyAuth(user_id="sampling-caller") + upstream = MCPServer( + server_id="sampling-context", + name="sampling_context", + url="https://example.invalid/mcp", + transport=MCPTransport.http, + allow_sampling=True, + ) + sampling = AsyncMock() + client = MagicMock() + client.call_tool = AsyncMock(return_value=CallToolResult(content=[])) + assert legacy_server.get_active_auth_context() is None + with ( + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient", return_value=client) as factory, + patch("litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", sampling), + ): + await MCPServerManager()._call_regular_mcp_tool( + mcp_server=upstream, + original_tool_name="probe", + arguments={}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=None, + raw_headers={"x-test-caller": "sampling-caller"}, + proxy_logging_obj=None, + user_api_key_auth=caller, + ) + callback = factory.call_args.kwargs["sampling_callback"] + await callback(None, None) + assert sampling.await_args.kwargs["user_api_key_auth"].user_id == "sampling-caller" + assert sampling.await_args.kwargs["raw_headers"] == {"x-test-caller": "sampling-caller"} + + + +@pytest.mark.asyncio +async def test_sampling_callback_keeps_creation_context_after_caller_switch(): + from mcp.server.auth.middleware.auth_context import auth_context_var + + from litellm.proxy._experimental.mcp_server import server as legacy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _create_sampling_callback + + token = auth_context_var.set(None) + recorder = AsyncMock() + try: + original = UserAPIKeyAuth(user_id="alpha", models=["alpha-model"]) + original.mcp_admitted_user_subject = True + headers = {"x-caller": "alpha"} + legacy_server.set_auth_context(original, raw_headers=headers, client_ip="192.0.2.1") + callback = _create_sampling_callback() + original.models.append("bravo-model") + headers["x-caller"] = "bravo" + legacy_server.set_auth_context(UserAPIKeyAuth(user_id="bravo"), raw_headers={"x-caller": "bravo"}) + with patch("litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", recorder): + await callback(None, None) + observed = recorder.await_args.kwargs + assert observed["user_api_key_auth"].user_id == "alpha" + assert observed["user_api_key_auth"].models == ["alpha-model"] + assert observed["user_api_key_auth"].mcp_admitted_user_subject is True + assert observed["raw_headers"] == {"x-caller": "alpha"} + assert observed["client_ip"] == "192.0.2.1" + finally: + auth_context_var.reset(token) + + +@pytest.mark.asyncio +async def test_elicitation_callback_keeps_initiating_session(): + from litellm.proxy._experimental.mcp_server import server as legacy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _create_elicitation_callback + + initiating = MagicMock() + replacement = MagicMock() + recorder = AsyncMock() + token = legacy_server.active_mcp_session_var.set(initiating) + try: + callback = _create_elicitation_callback() + legacy_server.active_mcp_session_var.set(replacement) + with patch("litellm.proxy._experimental.mcp_server.elicitation_handler.handle_elicitation_request", recorder): + await callback(None, None) + assert recorder.await_args.kwargs["downstream_session"] is initiating + assert recorder.await_args.kwargs["downstream_capabilities"] is initiating.capabilities + finally: + legacy_server.active_mcp_session_var.reset(token) + + +@pytest.mark.asyncio +async def test_sampling_callbacks_isolate_callers_and_cancellation(): + from mcp.types import ErrorData + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _create_sampling_callback + + started = asyncio.Event() + cancelled = asyncio.Event() + observed = {} + + async def record_sampling(*, user_api_key_auth, raw_headers, **kwargs): + label = user_api_key_auth.user_id + if label == "cancelled": + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + raise + await asyncio.sleep(0) + observed[label] = raw_headers["x-caller"] + return ErrorData(code=-1, message=label) + + callbacks = tuple( + _create_sampling_callback(UserAPIKeyAuth(user_id=label), raw_headers={"x-caller": label}) + for label in ("alpha", "bravo", "cancelled") + ) + with patch( + "litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", record_sampling + ): + tasks = tuple(asyncio.create_task(callback(None, None)) for callback in callbacks) + await asyncio.wait_for(started.wait(), timeout=2) + tasks[2].cancel() + results = await asyncio.gather(*tasks, return_exceptions=True) + assert observed == {"alpha": "alpha", "bravo": "bravo"} + assert [result.message for result in results[:2]] == ["alpha", "bravo"] + assert isinstance(results[2], asyncio.CancelledError) + assert cancelled.is_set() + + def _reload_mcp_manager_module(): utils_module = sys.modules["litellm.proxy._experimental.mcp_server.utils"] manager_module = sys.modules["litellm.proxy._experimental.mcp_server.mcp_server_manager"] @@ -84,6 +213,9 @@ def _reload_mcp_manager_module(): server_module = sys.modules.get("litellm.proxy._experimental.mcp_server.server") if server_module is not None and hasattr(server_module, "global_mcp_server_manager"): server_module.global_mcp_server_manager = reloaded.global_mcp_server_manager + operations_module = sys.modules.get("litellm.proxy._experimental.mcp_server.operations") + if operations_module is not None: + operations_module.global_mcp_server_manager = reloaded.global_mcp_server_manager return reloaded @@ -3923,6 +4055,7 @@ class TestMCPServerManager: result = await manager.get_resource_templates_from_server( server=server, user_api_key_auth=None, + raw_headers=None, mcp_auth_header="auth", extra_headers=None, add_prefix=False, @@ -3935,6 +4068,8 @@ class TestMCPServerManager: stdio_env=None, subject_token=None, user_api_key_auth=None, + raw_headers=None, + client_ip=None, ) mock_client.list_resource_templates.assert_awaited_once() assert result == expected_templates @@ -5849,7 +5984,7 @@ class TestMCPServerManager: stored = {"Authorization": "Bearer stored-user-token"} with patch( - "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + "litellm.proxy._experimental.mcp_server.operations._get_user_oauth_extra_headers_from_db", new=AsyncMock(return_value=stored), ) as mock_lookup: result = await manager._resolve_oauth2_headers_for_tool_call( @@ -5876,7 +6011,7 @@ class TestMCPServerManager: user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") with patch( - "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + "litellm.proxy._experimental.mcp_server.operations._get_user_oauth_extra_headers_from_db", new=AsyncMock(return_value={"Authorization": "Bearer should-not-be-used"}), ) as mock_lookup: result = await manager._resolve_oauth2_headers_for_tool_call( @@ -5902,7 +6037,7 @@ class TestMCPServerManager: user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") with patch( - "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + "litellm.proxy._experimental.mcp_server.operations._get_user_oauth_extra_headers_from_db", new=AsyncMock(side_effect=RuntimeError("redis down")), ): result = await manager._resolve_oauth2_headers_for_tool_call( @@ -6058,7 +6193,7 @@ class TestMCPServerManager: user_auth = UserAPIKeyAuth(api_key="sk-test") with patch( - "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + "litellm.proxy._experimental.mcp_server.operations._get_user_oauth_extra_headers_from_db", new=AsyncMock(return_value={"Authorization": "Bearer x"}), ) as mock_lookup: result = await manager._resolve_oauth2_headers_for_tool_call( @@ -6862,7 +6997,8 @@ class TestMCPServerManager: } user_api_key_auth = UserAPIKeyAuth(api_key="sk-test", user_id="user-123") - token = _mcp_active_toolset_id.set("toolset-abc") + user_api_key_auth.mcp_toolset_id = "toolset-abc" + token = _mcp_active_toolset_id.set("unrelated-ambient-toolset") try: with ( patch.object(proxy_server_module, "user_api_key_cache", cache), @@ -14332,3 +14468,36 @@ async def test_request_selected_during_guardrail_runs_concurrently_with_tool(mon assert guardrail_started.is_set() is selected assert result.is_error is False assert result.content[0].text == "executed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("with_caller,legacy_factory", [(True, False), (False, False), (True, True)]) +async def test_client_sampling_does_not_fill_explicit_context_from_another_ambient_caller(with_caller, legacy_factory): + from mcp.server.auth.middleware.auth_context import auth_context_var + from litellm.proxy._experimental.mcp_server import server as legacy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _create_sampling_callback + + upstream = MCPServer(server_id="explicit-empty", name="explicit_empty", url="https://example.invalid/mcp", transport=MCPTransport.http, allow_sampling=True) + token = auth_context_var.set(None) + sampling = AsyncMock() + try: + legacy_server.set_auth_context(UserAPIKeyAuth(user_id="unrelated"), raw_headers={"authorization": "unrelated-credential"}, client_ip="192.0.2.99") + with ( + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient") as factory, + patch("litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", sampling), + ): + if legacy_factory: + callback = _create_sampling_callback(user_api_key_auth=UserAPIKeyAuth(user_id="explicit")) + else: + await MCPServerManager()._create_mcp_client(upstream, user_api_key_auth=UserAPIKeyAuth(user_id="explicit") if with_caller else None) + callback = factory.call_args.kwargs["sampling_callback"] + await callback(None, None) + captured = sampling.await_args.kwargs + if with_caller: + assert captured["user_api_key_auth"].user_id == "explicit" + else: + assert captured["user_api_key_auth"] is None + assert captured["raw_headers"] is None + assert captured["client_ip"] is None + finally: + auth_context_var.reset(token) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index 9420eecd222..ec6fdef69ee 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -639,12 +639,12 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): return_value=False, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.has_user_oauth_token", new_callable=AsyncMock, return_value=False, ) as mock_has_token, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=oauth_server, ), patch.object( @@ -727,12 +727,12 @@ async def test_admitted_subject_missing_stored_token_challenged_with_resource_me return_value=False, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.has_user_oauth_token", new_callable=AsyncMock, return_value=False, ) as mock_has_token, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=oauth_server, ), patch.object( @@ -833,11 +833,11 @@ async def test_client_credentials_server_is_not_preemptively_challenged(m2m_fiel return_value=False, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.has_user_oauth_token", new_callable=AsyncMock, ) as mock_has_token, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=m2m_server, ), patch.object(session_manager_stateless, "handle_request", new_callable=AsyncMock) as mock_handle_request, @@ -929,16 +929,16 @@ async def test_handle_streamable_http_mcp_delegated_server_surfaces_upstream_cha return_value=False, ), patch( - "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + "litellm.proxy._experimental.mcp_server.operations._get_user_oauth_extra_headers_from_db", new_callable=AsyncMock, return_value=None, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=delegated_server, ), patch( # test-quality-ok: registry is empty in unit tests; key owns the delegated server - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[delegated_server], ), @@ -1022,12 +1022,12 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): return_value=False, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.has_user_oauth_token", new_callable=AsyncMock, return_value=True, ) as mock_has_token, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=oauth_server, ), patch.object( @@ -1126,11 +1126,11 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns return_value=False, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.has_user_oauth_token", new_callable=AsyncMock, ) as mock_has_token, patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=delegated_server, ), patch.object( @@ -1218,7 +1218,7 @@ async def test_handle_streamable_http_mcp_token_exchange_without_subject_returns return_value=False, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=obo_server, ), patch.object( @@ -1317,7 +1317,7 @@ async def test_handle_streamable_http_mcp_oauth_delegate_without_token_returns_g True, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=od_server, ), patch.object( @@ -1391,7 +1391,7 @@ async def test_handle_streamable_http_mcp_oauth_delegate_with_forwarded_token_sk new_callable=AsyncMock, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=od_server, ), patch.object( @@ -1453,7 +1453,7 @@ async def _run_passthrough_connect( new_callable=AsyncMock, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=server, ), patch.object(session_manager_stateless, "handle_request", new_callable=AsyncMock) as mock_handle_request, @@ -1574,7 +1574,7 @@ async def test_handle_streamable_http_mcp_true_passthrough_without_token_surface return_value=probe_client, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=tp_server, ), patch.object( @@ -1642,7 +1642,7 @@ async def test_handle_streamable_http_mcp_true_passthrough_dcr_bridge_challenges return_value=probe_client, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=bridge_server, ), patch.object( @@ -1720,7 +1720,7 @@ async def test_handle_streamable_http_mcp_true_passthrough_with_token_skips_prob return_value=probe_client, ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=tp_server, ), patch.object( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index cb43d2c2592..4575741aa8b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -1,3 +1,4 @@ +from litellm.proxy._experimental.mcp_server import operations as mcp_operations """ Tests for MCP tool search feature. @@ -572,7 +573,7 @@ class TestCallToolRestApiVirtualTools: mock_tool.input_schema = {"type": "object", "properties": {}} with patch( - "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + "litellm.proxy._experimental.mcp_server.operations._list_mcp_tools", new_callable=AsyncMock, return_value=AggregateToolListing(tools=[mock_tool], outcomes={}), ): @@ -616,12 +617,12 @@ class TestCallToolRestApiVirtualTools: with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[MagicMock()], ), patch( - "litellm.proxy._experimental.mcp_server.server.execute_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.execute_mcp_tool", new_callable=AsyncMock, return_value=fake_result, ) as mock_execute, @@ -669,12 +670,12 @@ class TestCallToolRestApiVirtualTools: return_value="203.0.113.7", ), patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[MagicMock()], ) as mock_allowed, patch( - "litellm.proxy._experimental.mcp_server.server.execute_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.execute_mcp_tool", new_callable=AsyncMock, return_value=fake_result, ), @@ -699,7 +700,7 @@ class TestCallToolRestApiVirtualTools: return_value="203.0.113.7", ), patch( - "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + "litellm.proxy._experimental.mcp_server.operations._list_mcp_tools", new_callable=AsyncMock, return_value=AggregateToolListing(tools=[], outcomes={}), ) as mock_list, @@ -832,7 +833,7 @@ class TestCallToolRestApiVirtualTools: "litellm.proxy.proxy_server.proxy_logging_obj", key_limits ), patch( # test-quality-ok: the authorized catalog is the seam every virtual tool shares; the ranking under test stays real - "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + "litellm.proxy._experimental.mcp_server.operations._list_mcp_tools", new_callable=AsyncMock, return_value=AggregateToolListing(tools=list(CATALOG), outcomes={}), ) as mock_list, @@ -939,7 +940,7 @@ class TestDispatchVirtualMcpTool: new_callable=AsyncMock, return_value="SEARCH_RESULT", ) as mock_search: - result = await srv._dispatch_virtual_mcp_tool( + result = await mcp_operations._dispatch_virtual_mcp_tool( name=MCP_TOOL_SEARCH_TOOL_NAME, arguments={"query": "q", "top_k": 3}, user_api_key_auth=uak, @@ -961,7 +962,7 @@ class TestDispatchVirtualMcpTool: new_callable=AsyncMock, return_value="AGENT_RESULT", ) as mock_agent_search: - result = await srv._dispatch_virtual_mcp_tool( + result = await mcp_operations._dispatch_virtual_mcp_tool( name=AGENT_SEARCH_TOOL_NAME, arguments={"query": "translate a document", "top_k": "2"}, user_api_key_auth=uak, @@ -996,7 +997,7 @@ class TestDispatchVirtualMcpTool: new_callable=AsyncMock, return_value="CALL_RESULT", ) as mock_call: - result = await srv._dispatch_virtual_mcp_tool( + result = await mcp_operations._dispatch_virtual_mcp_tool( name=MCP_TOOL_CALL_TOOL_NAME, arguments={"tool_name": "math-add", "arguments": {"a": 1, "b": 2}}, user_api_key_auth=uak, @@ -1027,8 +1028,7 @@ class TestDispatchVirtualMcpTool: sentinel_logging_obj = object() with ( patch.object( - srv, - "_build_virtual_call_logging_obj", + mcp_operations, "_build_virtual_call_logging_obj", new_callable=AsyncMock, return_value=sentinel_logging_obj, ) as mock_build, @@ -1038,7 +1038,7 @@ class TestDispatchVirtualMcpTool: return_value="CALL_RESULT", ) as mock_call, ): - await srv._dispatch_virtual_mcp_tool( + await mcp_operations._dispatch_virtual_mcp_tool( name=MCP_TOOL_CALL_TOOL_NAME, arguments={"tool_name": "math-add", "arguments": {"a": 1}}, user_api_key_auth=uak, @@ -1060,7 +1060,7 @@ class TestDispatchVirtualMcpTool: new_callable=AsyncMock, return_value="SEARCH_RESULT", ) as mock_search: - await srv._dispatch_virtual_mcp_tool( + await mcp_operations._dispatch_virtual_mcp_tool( name=MCP_TOOL_SEARCH_TOOL_NAME, arguments={"query": "issue", "top_k": "not-a-number"}, user_api_key_auth=uak, @@ -1083,12 +1083,12 @@ class TestDispatchVirtualMcpTool: fake = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[MagicMock()], ) as mock_allowed, patch( - "litellm.proxy._experimental.mcp_server.server.execute_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.execute_mcp_tool", new_callable=AsyncMock, return_value=fake, ) as mock_exec, @@ -1130,12 +1130,12 @@ class TestDispatchVirtualMcpTool: uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) with ( patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new_callable=AsyncMock, return_value=[], ), patch( - "litellm.proxy._experimental.mcp_server.server.execute_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations.execute_mcp_tool", new_callable=AsyncMock, ) as mock_exec, ): @@ -1217,7 +1217,7 @@ class TestMcpServerToolCallErrorHandling: return_value=(uak, None, None, None, None, None, None), ), patch( - "litellm.proxy._experimental.mcp_server.server._dispatch_virtual_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations._dispatch_virtual_mcp_tool", new_callable=AsyncMock, side_effect=HTTPException(status_code=403, detail="User not allowed to call this tool"), ), @@ -1254,7 +1254,7 @@ async def test_handle_mcp_tool_call_scoped_denial_names_the_binding_agent() -> N ] with patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", new=AsyncMock(side_effect=resolve), ): with pytest.raises(HTTPException) as exc_info: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py index 519acc241c6..1398884783e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py @@ -58,6 +58,22 @@ class TestApplyToolsetScope: assert set(op.mcp_servers or []) == {"server-a", "server-b"} assert op.mcp_tool_permissions == toolset_perms + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy._experimental.mcp_server.operations import prepare_context + + manager = MCPServerManager() + unscoped_open = await manager.operator_open_server_ids( + auth, allow_all_server_ids=["operator-open-outside-toolset"], submitted_server_ids=[] + ) + scoped_open = await manager.operator_open_server_ids( + prepare_context(result).user_api_key_auth, + allow_all_server_ids=["operator-open-outside-toolset"], + submitted_server_ids=[], + ) + assert unscoped_open == {"operator-open-outside-toolset"} + assert scoped_open == set() + assert auth.mcp_toolset_id is None + @pytest.mark.asyncio async def test_admin_creates_object_permission_when_none(self): """Admin key with object_permission=None can access any toolset.""" @@ -564,7 +580,7 @@ class TestMCPActiveToolsetContextVar: MagicMock(get_mcp_client_ip=MagicMock(return_value="127.0.0.1")), ), patch( - "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", MagicMock(get_mcp_server_by_name=MagicMock(return_value=None)), ), patch( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index ac716bace3c..bb70f38285c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -1,3 +1,4 @@ +from litellm.proxy._experimental.mcp_server import operations as mcp_operations """ VERIA-7 regression: OpenAPI-backed (local-registry) MCP tools must run through `pre_call_tool_check` before dispatch, the same as managed @@ -49,22 +50,22 @@ async def test_openapi_local_tool_runs_pre_call_tool_check(): with ( patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=fake_server, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "pre_call_tool_check", new=pre_call, ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=fake_tool, ), patch( - "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool", new=handle_local, ), patch( @@ -72,7 +73,7 @@ async def test_openapi_local_tool_runs_pre_call_tool_check(): return_value=True, ), ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="list_pets", arguments={"limit": 10}, allowed_mcp_servers=[fake_server], @@ -92,7 +93,7 @@ async def test_openapi_local_tool_runs_pre_call_tool_check(): assert pre_call_kwargs["guardrail_context"] == {"metadata": {"guardrails": ("block-all",)}} assert pre_call_kwargs["name"] == "list_pets" assert pre_call_kwargs["server"] is fake_server - assert pre_call_kwargs["user_api_key_auth"] is user + assert pre_call_kwargs["user_api_key_auth"] == user # `proxy_logging_obj` must be sourced from the canonical proxy_server # module (same as the managed path) — passing None would crash the # downstream `_create_mcp_request_object_from_kwargs` call with @@ -134,22 +135,22 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises(): with ( patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=fake_server, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "pre_call_tool_check", new=pre_call, ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=fake_tool, ), patch( - "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool", new=handle_local, ), patch( @@ -158,7 +159,7 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises(): ), ): with pytest.raises(HTTPException) as exc: - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="delete_pet", arguments={}, allowed_mcp_servers=[fake_server], @@ -195,24 +196,24 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable(): # `_get_mcp_server_from_tool_name` returns None — no server context. with ( - patch.object(mcp_module, "_resolve_openapi_tool_auth", new=resolve_auth), + patch.object(mcp_operations, "_resolve_openapi_tool_auth", new=resolve_auth), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=None, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "pre_call_tool_check", new=pre_call, ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=fake_tool, ), patch( - "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool", new=handle_local, ), patch( @@ -221,7 +222,7 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable(): ), ): with pytest.raises(HTTPException) as exc: - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="list_pets", arguments={}, allowed_mcp_servers=[], @@ -280,27 +281,27 @@ async def test_openapi_local_tool_injects_resolved_oauth_token(): with ( patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=oauth_server, ), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "pre_call_tool_check", new=AsyncMock(return_value={}), ), patch.object( - mcp_module.global_mcp_tool_registry, + mcp_operations.global_mcp_tool_registry, "get_tool", return_value=fake_tool, ), patch.object( - mcp_module.global_mcp_server_manager._cred_provider, + mcp_operations.global_mcp_server_manager._cred_provider, "resolve_credentials", new=AsyncMock(return_value=Ok(StaticHeaderAuth("Bearer stored-user-token"))), ), patch( - "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool", new=handle_local, ), patch( @@ -308,7 +309,7 @@ async def test_openapi_local_tool_injects_resolved_oauth_token(): return_value=True, ), ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="get_values", arguments={}, allowed_mcp_servers=[oauth_server], @@ -417,7 +418,7 @@ async def test_legacy_local_tool_fallback_refuses_unentitled_caller(legacy_local ) with pytest.raises(HTTPException) as exc: - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name=f"{LEGACY_SERVER_NAME}-{LEGACY_TOOL}", arguments={}, allowed_mcp_servers=[server], @@ -451,7 +452,7 @@ async def test_legacy_local_tool_fallback_still_dispatches_entitled_caller( server, executed = legacy_local_tool user = _caller_entitled_to([LEGACY_TOOL]) - result = await mcp_module.execute_mcp_tool( + result = await mcp_operations.execute_mcp_tool( name=f"{LEGACY_SERVER_NAME}-{LEGACY_TOOL}", arguments={}, allowed_mcp_servers=[server], @@ -481,7 +482,7 @@ async def test_legacy_local_tool_fallback_fails_closed_on_empty_prefix( _server, executed = legacy_local_tool with pytest.raises(HTTPException) as exc: - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name=f"-{LEGACY_TOOL}", arguments={}, allowed_mcp_servers=[], @@ -523,7 +524,7 @@ async def test_legacy_local_tool_fallback_fails_closed_when_prefix_names_no_serv return_value=True, ): with pytest.raises(HTTPException) as exc: - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name=f"{LEGACY_SERVER_NAME}-{LEGACY_TOOL}", arguments={}, allowed_mcp_servers=[other_server], @@ -546,7 +547,7 @@ async def test_unknown_tool_name_still_reports_not_found(): from litellm.proxy._experimental.mcp_server import server as mcp_module with pytest.raises(HTTPException) as exc: - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="tool_no_registry_knows", arguments={}, allowed_mcp_servers=[], @@ -610,7 +611,7 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc captured["injected"] = _request_auth_header.get() return [] - manager = mcp_module.global_mcp_server_manager + manager = mcp_operations.global_mcp_server_manager with ( patch.object(manager, "resolve_openapi_upstream_auth", new=fake_resolver), patch.object(manager, "pre_call_tool_check", new=AsyncMock(return_value={})), @@ -620,9 +621,9 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc fake_tool.name = "list_reports" with ( patch.object(manager, "_get_mcp_server_from_tool_name", return_value=server), - patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=fake_tool), + patch.object(mcp_operations.global_mcp_tool_registry, "get_tool", return_value=fake_tool), patch( - "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + "litellm.proxy._experimental.mcp_server.operations._handle_local_mcp_tool", new=capture_local, ), patch( @@ -630,7 +631,7 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc return_value=True, ), ): - await mcp_module.execute_mcp_tool( + await mcp_operations.execute_mcp_tool( name="list_reports", arguments={}, allowed_mcp_servers=[server], @@ -702,11 +703,11 @@ async def test_local_dispatch_reports_the_outcome_instead_of_success(failure: st user = UserAPIKeyAuth(api_key="sk-user", user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) with ( - patch.object(mcp_module.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=server), - patch.object(mcp_module.global_mcp_server_manager, "pre_call_tool_check", new=AsyncMock(return_value={})), - patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=fake_tool), + patch.object(mcp_operations.global_mcp_server_manager, "_get_mcp_server_from_tool_name", return_value=server), + patch.object(mcp_operations.global_mcp_server_manager, "pre_call_tool_check", new=AsyncMock(return_value={})), + patch.object(mcp_operations.global_mcp_tool_registry, "get_tool", return_value=fake_tool), patch.object( - mcp_module.global_mcp_server_manager, + mcp_operations.global_mcp_server_manager, "resolve_openapi_upstream_auth", new=AsyncMock(return_value=(None, None)), ), @@ -715,7 +716,7 @@ async def test_local_dispatch_reports_the_outcome_instead_of_success(failure: st return_value=True, ), ): - call = mcp_module.execute_mcp_tool( + call = mcp_operations.execute_mcp_tool( name="list_reports", arguments={}, allowed_mcp_servers=[server], diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py new file mode 100644 index 00000000000..abb925ddc77 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py @@ -0,0 +1,365 @@ +import asyncio +from unittest.mock import AsyncMock, patch + +import pytest +from mcp.types import GetPromptRequest, GetPromptRequestParams, GetPromptResult + +from litellm.proxy._experimental.mcp_server.operations import GatewayOperations, prepare_context +from litellm.proxy._types import UserAPIKeyAuth + + +@pytest.mark.asyncio +async def test_oauth_prefetch_failure_does_not_log_caller_or_exception_text(caplog): + from litellm.proxy._experimental.mcp_server.operations import _prefetch_oauth_creds_for_user + + user_id = "caller\nFORGED-USER-LINE" + fetch = AsyncMock(side_effect=RuntimeError("database\nFORGED-ERROR-LINE")) + database = object() + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=database), + patch("litellm.proxy._experimental.mcp_server.db.list_user_oauth_credentials", fetch), + caplog.at_level("WARNING", logger="LiteLLM"), + ): + result = await _prefetch_oauth_creds_for_user(UserAPIKeyAuth(user_id=user_id)) + assert result == {} + fetch.assert_awaited_once_with(database, user_id) + warnings = [record.getMessage() for record in caplog.records if "prefetch" in record.getMessage()] + assert len(warnings) == 1 + assert "failed" in warnings[0] + assert "\n" not in warnings[0] + assert "FORGED" not in warnings[0] + + +@pytest.mark.asyncio +async def test_dispatch_uses_explicit_context_when_ambient_caller_differs(): + from mcp.server.auth.middleware.auth_context import auth_context_var + from litellm.proxy._experimental.mcp_server.server import set_auth_context + + context = prepare_context( + UserAPIKeyAuth(user_id="alpha"), + raw_headers={"x-caller": "alpha"}, + mcp_servers=["alpha-server"], + client_ip="192.0.2.1", + ) + token = auth_context_var.set(None) + handler = AsyncMock(return_value=GetPromptResult(messages=[])) + try: + set_auth_context(UserAPIKeyAuth(user_id="bravo"), raw_headers={"x-caller": "bravo"}) + with patch("litellm.proxy._experimental.mcp_server.operations.mcp_get_prompt", handler): + result = await GatewayOperations().execute( + GetPromptRequest(params=GetPromptRequestParams(name="alpha-prompt")), context + ) + assert result.messages == [] + assert handler.await_args.kwargs["name"] == "alpha-prompt" + assert handler.await_args.kwargs["user_api_key_auth"].user_id == "alpha" + assert handler.await_args.kwargs["raw_headers"] == {"x-caller": "alpha"} + assert handler.await_args.kwargs["mcp_servers"] == ["alpha-server"] + assert handler.await_args.kwargs["client_ip"] == "192.0.2.1" + finally: + auth_context_var.reset(token) + + +@pytest.mark.asyncio +async def test_legacy_adapter_cleans_context_after_cancelled_operation(): + from types import SimpleNamespace + from litellm.proxy._experimental.mcp_server import server + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var + + previous_session = server.active_mcp_session_var.get() + previous_request = active_mcp_request_ctx_var.get() + request = SimpleNamespace(session=object()) + auth = (None, None, None, None, None, None, None) + + async def cancelled_operation(): + async with server._legacy_operation_context(request, trace=False): + assert server.active_mcp_session_var.get() is request.session + assert active_mcp_request_ctx_var.get() is request + raise asyncio.CancelledError + + with patch( + "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", AsyncMock(return_value=auth) + ): + with pytest.raises(asyncio.CancelledError): + await cancelled_operation() + assert server.active_mcp_session_var.get() is previous_session + assert active_mcp_request_ctx_var.get() is previous_request + + +@pytest.mark.asyncio +async def test_legacy_adapter_cleans_context_when_trace_setup_fails(): + from types import SimpleNamespace + from litellm.proxy._experimental.mcp_server import server + from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var + + previous_session = server.active_mcp_session_var.get() + previous_request = active_mcp_request_ctx_var.get() + request = SimpleNamespace(session=object()) + + async def enter_operation(): + async with server._legacy_operation_context(request, trace=True): + pytest.fail("Trace setup failure must prevent dispatch") + + with patch.object(server, "_otel_set_mcp_transport_span", side_effect=RuntimeError("trace failure")): + with pytest.raises(RuntimeError, match="trace failure"): + await enter_operation() + assert server.active_mcp_session_var.get() is previous_session + assert active_mcp_request_ctx_var.get() is previous_request + + +@pytest.mark.asyncio +async def test_prompt_sampling_receives_explicit_operation_caller_headers_and_ip(): + from unittest.mock import MagicMock + from litellm.proxy._experimental.mcp_server import operations + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + upstream = MCPServer( + server_id="explicit-prompt", + name="explicit_prompt", + url="https://example.invalid/mcp", + transport=MCPTransport.http, + allow_sampling=True, + ) + context = prepare_context( + UserAPIKeyAuth(user_id="prompt-caller"), + raw_headers={"x-caller": "prompt-caller"}, + client_ip="192.0.2.41", + ) + client = MagicMock() + client.get_prompt = AsyncMock(return_value=GetPromptResult(messages=[])) + sampling = AsyncMock() + with ( + patch.object(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[upstream])), + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient", return_value=client) as factory, + patch("litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", sampling), + ): + result = await GatewayOperations().execute( + GetPromptRequest(params=GetPromptRequestParams(name="explicit_prompt-prompt")), context + ) + assert result.messages == [] + await factory.call_args.kwargs["sampling_callback"](None, None) + captured = sampling.await_args.kwargs + assert captured["user_api_key_auth"] is not None + assert captured["user_api_key_auth"].user_id == "prompt-caller" + assert captured["raw_headers"] == {"x-caller": "prompt-caller"} + assert captured["client_ip"] == "192.0.2.41" + + +def _catalog_case(method): + from mcp import types + + cases = { + "prompts/list": ( + types.ListPromptsRequest(), + "list_prompts", + "get_prompts_from_server", + [types.Prompt(name="catalog-prompt")], + "prompts", + ), + "prompts/get": ( + types.GetPromptRequest( + params=types.GetPromptRequestParams(name="catalog-prompt", arguments={"topic": "test"}) + ), + "get_prompt", + "get_prompt_from_server", + types.GetPromptResult(messages=[]), + None, + ), + "resources/list": ( + types.ListResourcesRequest(), + "list_resources", + "get_resources_from_server", + [types.Resource(name="document", uri="https://example.com/document")], + "resources", + ), + "resources/templates/list": ( + types.ListResourceTemplatesRequest(), + "list_resource_templates", + "get_resource_templates_from_server", + [types.ResourceTemplate(name="document", uri_template="https://example.com/{name}")], + "resource_templates", + ), + "resources/read": ( + types.ReadResourceRequest(params=types.ReadResourceRequestParams(uri="https://example.com/document")), + "read_resource", + "read_resource_from_server", + types.ReadResourceResult( + contents=[types.TextResourceContents(uri="https://example.com/document", text="document body")] + ), + None, + ), + } + return cases[method] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method", ["prompts/list", "prompts/get", "resources/list", "resources/templates/list", "resources/read"] +) +@pytest.mark.parametrize("state", ["success", "denied", "upstream_failure", "scope_failure"]) +async def test_native_catalog_operations_preserve_context_results_and_failure_policy(method, state): + from types import SimpleNamespace + from fastapi import HTTPException + from mcp.server.context import ServerRequestContext + from mcp.types import PaginatedRequestParams + from litellm.proxy._experimental.mcp_server import operations, server + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + operation, handler_name, manager_method, payload, collection = _catalog_case(method) + caller = UserAPIKeyAuth(user_id="catalog-caller") + headers = {"x-caller": "catalog-caller"} + upstream_server = MCPServer(server_id="catalog", name="catalog", transport=MCPTransport.http) + allowed = AsyncMock( + return_value=[] if state == "denied" else [upstream_server], + side_effect=HTTPException(status_code=403, detail="scope denied") if state == "scope_failure" else None, + ) + upstream = AsyncMock( + return_value=payload, side_effect=RuntimeError("upstream unavailable") if state == "upstream_failure" else None + ) + ctx = ServerRequestContext( + session=SimpleNamespace(), lifespan_context={}, protocol_version="2025-06-18", method=method + ) + auth = (caller, None, ["catalog"], None, None, headers, "192.0.2.41") + with ( + patch.object(server, "get_or_extract_auth_context", AsyncMock(return_value=auth)), + patch.object(operations, "_get_allowed_mcp_servers", allowed), + patch.object(operations.global_mcp_server_manager, manager_method, upstream), + ): + if collection is None and state != "success": + expected_error = RuntimeError if state == "upstream_failure" else HTTPException + with pytest.raises(expected_error): + await getattr(server, handler_name)(ctx, operation.params) + else: + result = await getattr(server, handler_name)(ctx, operation.params or PaginatedRequestParams()) + if collection: + assert getattr(result, collection) == (payload if state == "success" else []) + else: + assert result == payload + assert allowed.await_args.kwargs == { + "user_api_key_auth": caller, + "mcp_servers": ["catalog"], + "client_ip": "192.0.2.41", + } + if state in ("denied", "scope_failure"): + upstream.assert_not_awaited() + else: + upstream.assert_awaited_once() + forwarded = upstream.await_args.kwargs + assert forwarded["user_api_key_auth"] == caller + assert forwarded["raw_headers"] == headers + assert forwarded["client_ip"] == "192.0.2.41" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method", ["prompts/list", "prompts/get", "resources/list", "resources/templates/list", "resources/read"] +) +async def test_explicit_proxy_context_rejects_catalog_operations_before_upstream_access(method): + from mcp.shared.exceptions import MCPError + from mcp.types import METHOD_NOT_FOUND + from litellm.proxy._experimental.mcp_server import operations + + operation, _, manager_method, _, _ = _catalog_case(method) + upstream = AsyncMock() + with patch.object(operations.global_mcp_server_manager, manager_method, upstream): + with pytest.raises(MCPError) as rejected: + await GatewayOperations().execute(operation, prepare_context(mcp_proxy_mode=True)) + assert rejected.value.error.code == METHOD_NOT_FOUND + assert rejected.value.error.message == "Operation unavailable on /mcp/proxy" + upstream.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["missing_env", "pii", "guardrail", "unexpected"]) +async def test_tool_operation_preserves_failure_messages_and_request_trace(failure): + from mcp.types import CallToolRequest, CallToolRequestParams + from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.utils import MCPMissingUserEnvVarsError + + failures = { + "missing_env": ( + MCPMissingUserEnvVarsError( + server_id="server", server_name="server", missing=["TOKEN"], setup_url="https://example.com/setup" + ), + "https://example.com/setup", + ), + "pii": ( + BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="test"), + "Blocked PII entity detected", + ), + "guardrail": (GuardrailRaisedException(message="request denied"), "Guardrail violation"), + "unexpected": (RuntimeError("upstream unavailable"), "Error: upstream unavailable"), + } + error, expected = failures[failure] + dispatch = AsyncMock(side_effect=error) + context = prepare_context( + raw_headers={"x-litellm-trace-id": "operation-trace", "authorization": "private-test-header"} + ) + with patch.object(operations, "call_mcp_tool", dispatch): + result = await GatewayOperations().execute( + CallToolRequest(params=CallToolRequestParams(name="catalog-tool", arguments={})), context + ) + assert result.is_error is True + assert expected in result.content[0].text + assert "private-test-header" not in result.content[0].text + dispatch.assert_awaited_once() + assert dispatch.await_args.kwargs["litellm_trace_id"] == "operation-trace" + assert dispatch.await_args.kwargs["litellm_session_id"] == "operation-trace" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method,helper", + [ + ("prompts/list", "_list_mcp_prompts"), + ("resources/list", "_list_mcp_resources"), + ("resources/templates/list", "_list_mcp_resource_templates"), + ], +) +async def test_catalog_operation_preserves_empty_result_for_malformed_upstream_items(method, helper): + from litellm.proxy._experimental.mcp_server import operations + + operation, _, _, _, collection = _catalog_case(method) + with patch.object(operations, helper, AsyncMock(return_value=[{"unexpected": "item"}])): + result = await GatewayOperations().execute(operation, prepare_context()) + assert getattr(result, collection) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("catalog_unavailable", [False, True]) +async def test_tool_listing_returns_empty_result_without_dispatch_for_unavailable_catalog(catalog_unavailable): + from mcp.types import ListToolsRequest + from litellm.proxy._experimental.mcp_server import operations + + allowed = AsyncMock( + return_value=[], side_effect=RuntimeError("catalog unavailable") if catalog_unavailable else None + ) + upstream = AsyncMock() + with ( + patch.object(operations, "_get_allowed_mcp_servers", allowed), + patch.object(operations.global_mcp_server_manager, "_get_tools_from_server", upstream), + ): + result = await GatewayOperations().execute(ListToolsRequest(), prepare_context()) + assert result.tools == [] + allowed.assert_awaited_once() + upstream.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_explicit_proxy_context_lists_builtin_tools_and_blocks_direct_tool_dispatch(): + from mcp.types import CallToolRequest, CallToolRequestParams, ListToolsRequest + from litellm.proxy._experimental.mcp_server import operations + + context = prepare_context(mcp_proxy_mode=True) + allowed = AsyncMock() + with patch.object(operations, "_get_allowed_mcp_servers", allowed): + listing = await GatewayOperations().execute(ListToolsRequest(), context) + denied = await GatewayOperations().execute( + CallToolRequest(params=CallToolRequestParams(name="catalog-tool", arguments={})), context + ) + assert {tool.name for tool in listing.tools} == {"search_tools", "get_tool_schema", "call_tool"} + assert denied.is_error is True + assert "unavailable on /mcp/proxy" in denied.content[0].text + allowed.assert_not_awaited() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 13af58c15c0..233a8cc96ba 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1,3 +1,4 @@ +from litellm.proxy._experimental.mcp_server import operations as mcp_operations import asyncio import inspect import json @@ -1253,6 +1254,7 @@ class TestListToolsRestAPI: user_api_key_auth=None, extra_headers=None, apply_tool_filters=True, + client_ip=None, ): captured["called"] = True captured["server"] = server @@ -1338,6 +1340,7 @@ class TestListToolsRestAPI: user_api_key_auth=None, extra_headers=None, apply_tool_filters=True, + client_ip=None, ): captured["user_api_key_auth"] = user_api_key_auth return ["tool-1"] @@ -1891,6 +1894,7 @@ class TestListToolsRestAPI: user_api_key_auth=None, extra_headers=None, apply_tool_filters=True, + client_ip=None, ): captured["called"] = True captured["server_arg"] = server @@ -2027,6 +2031,7 @@ class TestListToolsRestAPI: user_api_key_auth=None, extra_headers=None, apply_tool_filters=True, + client_ip=None, ): captured["called"] = True captured["server_arg"] = server @@ -2112,6 +2117,7 @@ class TestListToolsRestAPI: user_api_key_auth=None, extra_headers=None, apply_tool_filters=True, + client_ip=None, ): return ["scoped-tool"] @@ -2319,6 +2325,7 @@ class TestListToolsRestAPI: user_api_key_auth=None, extra_headers=None, apply_tool_filters=True, + client_ip=None, ): captured["server"] = server captured["auth_header"] = server_auth_header @@ -3145,10 +3152,10 @@ async def test_request_selected_tool_specific_guardrail_applies_to_virtual_execu monkeypatch.setattr(litellm, "callbacks", [guardrail]) monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) - monkeypatch.setattr(server, "global_mcp_tool_registry", registry) - monkeypatch.setattr(server, "global_mcp_server_manager", manager) + monkeypatch.setattr(mcp_operations, "global_mcp_tool_registry", registry) + monkeypatch.setattr(mcp_operations, "global_mcp_server_manager", manager) monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) - monkeypatch.setattr(server, "_get_allowed_mcp_servers", AsyncMock(return_value=[managed_server])) + monkeypatch.setattr(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[managed_server])) monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache())) monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", passthrough_request_data) monkeypatch.setattr(proxy_server, "proxy_config", {}) diff --git a/tests/test_litellm/test_check_mcp_operation_boundary.py b/tests/test_litellm/test_check_mcp_operation_boundary.py new file mode 100644 index 00000000000..d7ac72de9f0 --- /dev/null +++ b/tests/test_litellm/test_check_mcp_operation_boundary.py @@ -0,0 +1,52 @@ +from pathlib import Path + +import pytest + +from scripts.check_mcp_operation_boundary import main, violations + + +@pytest.mark.parametrize( + "source", + ( + "from mcp.server.auth.middleware.auth_context import auth_context_var as hidden", + "from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_proxy_mode as mode", + "caller = legacy.get_active_auth_context()", + "owners = transport._stateful_session_owners", + "from weakref import WeakKeyDictionary", + "from litellm.proxy._experimental.mcp_server.server import get_auth_context", + ), +) +def test_shared_operation_boundary_rejects_ambient_state(source): + assert violations(Path("operations.py"), source) + + +def test_legacy_adapter_may_resolve_context_but_policy_must_receive_it(): + source = "from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_proxy_mode" + assert violations(Path("server.py"), source) == () + assert violations(Path("legacy_callbacks.py"), source) == () + assert violations(Path("operations.py"), "def execute(context):\n return context.client_ip") == () + assert violations(Path("mcp_server_manager.py"), "def _mcp_registry_key(server):\n return server.name") == () + + +def test_boundary_command_rejects_shared_state_and_accepts_explicit_context(tmp_path, monkeypatch, capsys): + import subprocess + import sys + + package = tmp_path / "litellm/proxy/_experimental/mcp_server" + package.mkdir(parents=True) + module = package / "operations.py" + module.write_text("from mcp.server.auth.middleware.auth_context import auth_context_var as hidden\n") + command = [sys.executable, str(Path(__file__).resolve().parents[2] / "scripts/check_mcp_operation_boundary.py")] + monkeypatch.chdir(tmp_path) + assert main() == 1 + assert "operations.py:1:" in capsys.readouterr().err + rejected = subprocess.run(command, cwd=tmp_path, capture_output=True, text=True, check=False) + assert rejected.returncode == 1 + assert "operations.py:1: MCP request/session state belongs in a legacy adapter" in rejected.stderr + + module.write_text("def execute(context):\n return context.client_ip\n") + assert main() == 0 + assert "MCP operation boundary: passed" in capsys.readouterr().out + accepted = subprocess.run(command, cwd=tmp_path, capture_output=True, text=True, check=False) + assert accepted.returncode == 0 + assert "MCP operation boundary: passed" in accepted.stdout