From 067514bc76794fc3f9a66b9083f1cf52e6728427 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Wed, 16 Sep 2026 21:33:26 -0700 Subject: [PATCH 1/8] fix(responses): patch custom_tool_call_output in place on guardrail write-back --- .../guardrail_translation/handler.py | 2 +- ...test_openai_responses_guardrail_handler.py | 50 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 1ab4811b8a0..9739ca95089 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -209,7 +209,7 @@ _TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | f ) _OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) _PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType( - {"function_call_output": "output", "message": "content"} + {"function_call_output": "output", "custom_tool_call_output": "output", "message": "content"} ) _EMPTY_RESPONSES_REQUEST: Final[ResponsesAPIOptionalRequestParams] = {} diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index c714b5d378a..ea58a48a3d5 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -2223,6 +2223,56 @@ class TestStructuredMessagesWriteBack: } assert result["input"][3] == {"role": "user", "content": "What is the codename?"} + @pytest.mark.asyncio + async def test_codex_custom_tool_items_survive_tool_output_compression(self): + handler = OpenAIResponsesHandler() + additional_tools_item = { + "type": "additional_tools", + "tools": [{"type": "custom", "name": "exec", "description": "Run a JavaScript snippet"}], + } + reasoning_item = { + "id": "rs_456", + "type": "reasoning", + "summary": [], + "encrypted_content": "gAAAAA-signed-reasoning", + } + custom_tool_call_item = { + "id": "ctc_456", + "type": "custom_tool_call", + "call_id": "call_exec", + "name": "exec", + "input": 'const r = await tools.exec_command({"cmd": "cat memo.txt"});\ntext(r.output);', + "status": "completed", + } + data = { + "model": "gpt-5.6", + "input": [ + additional_tools_item, + {"role": "user", "content": "What is the codename?"}, + reasoning_item, + custom_tool_call_item, + { + "type": "custom_tool_call_output", + "call_id": "call_exec", + "output": [ + {"type": "input_text", "text": "Script completed\nOutput:\n"}, + {"type": "input_text", "text": "memo " * 400}, + ], + }, + ], + } + + result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail()) + + assert result["input"][0] is additional_tools_item + assert result["input"][1] == {"role": "user", "content": "What is the codename?"} + assert result["input"][2] is reasoning_item + assert result["input"][3] is custom_tool_call_item + assert result["input"][4]["type"] == "custom_tool_call_output" + assert result["input"][4]["call_id"] == "call_exec" + assert COMPRESSED_MARKER in str(result["input"][4]["output"]) + assert len(result["input"]) == 5 + @pytest.mark.asyncio async def test_web_search_call_item_preserved_verbatim(self): handler = OpenAIResponsesHandler() From 1098604ed661ee54b66ec7a378a6163f589b10df Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:24:15 -0700 Subject: [PATCH 2/8] refactor(mcp): extract explicit operation context and dispatch --- .github/workflows/test-linting.yml | 4 + Makefile | 1 + .../_experimental/mcp_server/contracts.py | 95 + .../mcp_server/legacy_callbacks.py | 83 + .../mcp_server/mcp_server_manager.py | 129 +- .../_experimental/mcp_server/operations.py | 3102 +++++++++++++++++ .../mcp_server/rest_endpoints.py | 13 +- .../proxy/_experimental/mcp_server/server.py | 3059 ++-------------- .../_experimental/mcp_server/tool_search.py | 11 +- litellm/proxy/_types.py | 2 + scripts/check_mcp_operation_boundary.py | 65 + scripts/pre_commit_lint.sh | 3 + .../mcp_server/test_byok_oauth_endpoints.py | 19 +- .../mcp_server/test_contracts.py | 60 + .../mcp_server/test_mcp_block_recording.py | 5 +- .../mcp_server/test_mcp_hook_extra_headers.py | 6 +- .../test_mcp_oauth_passthrough_tools.py | 61 +- .../mcp_server/test_mcp_proxy_mode.py | 8 +- .../mcp_server/test_mcp_server.py | 385 +- .../mcp_server/test_mcp_server_manager.py | 174 +- .../mcp_server/test_mcp_stale_session.py | 40 +- .../mcp_server/test_mcp_tool_search.py | 40 +- .../mcp_server/test_mcp_toolset_scope.py | 2 +- .../mcp_server/test_openapi_tool_auth.py | 75 +- .../mcp_server/test_operations.py | 343 ++ .../mcp_server/test_rest_endpoints.py | 13 +- .../test_check_mcp_operation_boundary.py | 52 + 27 files changed, 4624 insertions(+), 3226 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/contracts.py create mode 100644 litellm/proxy/_experimental/mcp_server/legacy_callbacks.py create mode 100644 litellm/proxy/_experimental/mcp_server/operations.py create mode 100644 scripts/check_mcp_operation_boundary.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_contracts.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py create mode 100644 tests/test_litellm/test_check_mcp_operation_boundary.py 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..4424907c28a --- /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 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 b293ab5a206..264d7cf19aa 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( @@ -3373,17 +3319,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() @@ -4151,6 +4093,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. @@ -4199,7 +4143,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 @@ -4344,6 +4294,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. @@ -4433,6 +4384,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 @@ -4543,6 +4496,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 = ( @@ -4563,6 +4517,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( @@ -4586,6 +4542,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 = ( @@ -4606,6 +4563,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( @@ -4629,6 +4588,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 = ( @@ -4649,6 +4609,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( @@ -4672,6 +4634,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.""" @@ -4692,6 +4655,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) @@ -4705,6 +4671,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.""" @@ -4725,6 +4692,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( @@ -5805,6 +5775,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. @@ -5830,6 +5802,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) @@ -5847,6 +5821,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. @@ -5991,6 +5966,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( @@ -6014,6 +5991,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() @@ -6189,7 +6168,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, ) @@ -6295,6 +6274,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 @@ -6421,6 +6401,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..3ee4f37add7 --- /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, assert_never, 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 + +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 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) + + +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 _: + 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 3a9bca926b0..33978bb9182 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,56 @@ def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: if MCP_AVAILABLE: + __all__ = ( + "_MCP_CREDENTIAL_REQUEST_FIELDS", + "ListMCPToolsRestAPIResponseObject", + "_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 +474,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 +508,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 +798,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 +873,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 +954,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 +981,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 +1017,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 +1533,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 +1554,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 +1578,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 +1591,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 +1619,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 +1702,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 +1723,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 +1742,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 +1759,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 +1887,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 3b34440d0bf..07338bcff52 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3238,6 +3238,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, @@ -3279,6 +3280,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/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 3668a06203c..169de9095cf 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(): @@ -2333,7 +2337,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()], ), @@ -2823,7 +2827,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()], ), @@ -2976,7 +2980,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()], ), @@ -3451,7 +3455,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()], ), @@ -4164,7 +4168,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( @@ -4172,7 +4176,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, ), ): @@ -4281,16 +4285,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, ), @@ -4372,7 +4376,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( @@ -4451,7 +4455,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( @@ -4631,7 +4635,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, ), ): @@ -4661,11 +4665,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"}), ), ): @@ -4737,7 +4741,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, ), ): @@ -4880,7 +4884,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( @@ -4991,7 +4995,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 @@ -5083,7 +5087,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( @@ -5189,7 +5193,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( @@ -5631,12 +5635,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"), ), @@ -5700,26 +5704,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, ), ): @@ -5782,26 +5786,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), ), ): @@ -6102,23 +6106,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), ), ): @@ -6450,7 +6454,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], ), @@ -6480,7 +6484,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=[], ): @@ -6506,7 +6510,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=[], ): @@ -6531,7 +6535,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=[], ): @@ -6587,7 +6591,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], ), @@ -6722,14 +6726,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), ), ): @@ -6992,7 +6996,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, ) @@ -7011,7 +7015,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( @@ -7047,7 +7051,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 @@ -7094,7 +7098,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 @@ -7140,7 +7144,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( @@ -7173,7 +7177,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( @@ -7216,7 +7220,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( @@ -7262,7 +7266,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( @@ -7295,7 +7299,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( @@ -7391,11 +7395,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, @@ -7403,13 +7407,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( @@ -7418,12 +7421,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], @@ -7456,7 +7459,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): @@ -7468,17 +7471,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) @@ -7492,7 +7495,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], @@ -7513,7 +7516,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], @@ -7536,7 +7539,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], @@ -7557,8 +7560,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], @@ -7580,8 +7583,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], @@ -7604,7 +7607,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], @@ -7650,13 +7653,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( @@ -7665,12 +7667,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], @@ -7724,11 +7726,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, @@ -7736,7 +7738,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, ), @@ -7746,13 +7748,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], @@ -7795,7 +7797,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, @@ -7803,7 +7805,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, ), @@ -7813,13 +7815,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], @@ -7857,7 +7859,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, @@ -7865,7 +7867,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, ), @@ -7875,13 +7877,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], @@ -7921,18 +7923,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( @@ -7941,12 +7942,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], @@ -8005,22 +8006,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( @@ -8028,7 +8029,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], @@ -8084,7 +8085,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, @@ -8092,13 +8093,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( @@ -8107,12 +8107,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], @@ -8164,7 +8164,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, @@ -8172,7 +8172,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, ), @@ -8182,13 +8182,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], @@ -9091,14 +9091,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), ), ): @@ -9176,12 +9176,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"), ), @@ -9261,7 +9261,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( @@ -9335,7 +9335,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), ), ): @@ -9401,12 +9401,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, @@ -9425,7 +9425,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", @@ -9457,7 +9457,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( @@ -9500,12 +9500,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, @@ -9607,17 +9607,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( @@ -9667,12 +9667,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, ), @@ -9733,13 +9733,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( @@ -10048,7 +10048,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 @@ -10111,11 +10111,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"), ), ): @@ -10208,3 +10208,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 9140ac61f1a..7be79f9b514 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 @@ -71,6 +71,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"] @@ -82,6 +211,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 @@ -3921,6 +4053,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, @@ -3933,6 +4066,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 @@ -5847,7 +5982,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( @@ -5874,7 +6009,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( @@ -5900,7 +6035,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( @@ -6056,7 +6191,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( @@ -6860,7 +6995,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), @@ -14075,3 +14211,31 @@ 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", [True, False]) +async def test_client_sampling_does_not_fill_explicit_context_from_another_ambient_caller(with_caller): + from mcp.server.auth.middleware.auth_context import auth_context_var + from litellm.proxy._experimental.mcp_server import server as legacy_server + + 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), + ): + await MCPServerManager()._create_mcp_client(upstream, user_api_key_auth=UserAPIKeyAuth(user_id="explicit") if with_caller else None) + await factory.call_args.kwargs["sampling_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..60098b1f656 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 @@ -564,7 +564,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..870f2cd5f8f --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py @@ -0,0 +1,343 @@ +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_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 From d2f30a77fded3a39c32b5dfc2f7f464bcd8f0409 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:31:48 -0700 Subject: [PATCH 3/8] test(mcp): preserve toolset scope across explicit context --- .../mcp_server/test_mcp_toolset_scope.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) 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 60098b1f656..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.""" From e208b4e89ed41c6e4239c638f8e51a1239560959 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:40:12 -0700 Subject: [PATCH 4/8] fix(mcp): keep explicit legacy sampling callers isolated --- .../_experimental/mcp_server/legacy_callbacks.py | 2 +- .../mcp_server/test_mcp_server_manager.py | 13 +++++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py b/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py index 4424907c28a..9e321062643 100644 --- a/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py +++ b/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py @@ -33,7 +33,7 @@ def create_sampling_callback( ) -> SamplingCallback: from litellm.proxy._experimental.mcp_server.server import get_active_auth_context - auth: Final = get_active_auth_context() if operation_context is None else None + 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 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 7be79f9b514..2dc7e01966a 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 @@ -14214,10 +14214,11 @@ async def test_request_selected_during_guardrail_runs_concurrently_with_tool(mon @pytest.mark.asyncio -@pytest.mark.parametrize("with_caller", [True, False]) -async def test_client_sampling_does_not_fill_explicit_context_from_another_ambient_caller(with_caller): +@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) @@ -14228,8 +14229,12 @@ async def test_client_sampling_does_not_fill_explicit_context_from_another_ambie 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), ): - await MCPServerManager()._create_mcp_client(upstream, user_api_key_auth=UserAPIKeyAuth(user_id="explicit") if with_caller else None) - await factory.call_args.kwargs["sampling_callback"](None, None) + 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" From 9002749e29bcff3de2f8991ee988f923ffc5a3c5 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:58:06 -0700 Subject: [PATCH 5/8] fix(mcp): preserve Python 3.10 imports and integration test seams --- .../_experimental/mcp_server/operations.py | 4 ++-- tests/mcp_tests/test_mcp_logging.py | 6 ++--- tests/mcp_tests/test_mcp_server.py | 24 ++++++++++--------- 3 files changed, 18 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 3ee4f37add7..a579c33e702 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -6,7 +6,7 @@ import types import uuid from collections.abc import Mapping, Sequence from datetime import datetime -from typing import Any, Final, NoReturn, TypeAlias, assert_never, overload +from typing import Any, Final, NoReturn, TypeAlias, overload from fastapi import HTTPException from mcp import ReadResourceResult, Resource @@ -34,7 +34,7 @@ from mcp.types import ( ) from mcp.types import Tool as MCPTool from pydantic import AnyUrl, ConfigDict, Field, TypeAdapter -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_logger from litellm.constants import ( 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( From ef67412e50c9f3250eb473a0a14fbf00617e36ea Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:13:20 -0700 Subject: [PATCH 6/8] fix(mcp): keep OAuth prefetch failure logs free of caller data --- .../_experimental/mcp_server/operations.py | 6 ++--- .../proxy/_experimental/mcp_server/server.py | 3 +++ .../mcp_server/test_operations.py | 22 +++++++++++++++++++ 3 files changed, 28 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index a579c33e702..fcee3483e15 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -763,8 +763,8 @@ async def _prefetch_oauth_creds_for_user( ) 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) + except Exception: + verbose_logger.warning("_prefetch_oauth_creds_for_user: failed to prefetch OAuth credentials") return {} @@ -3099,4 +3099,4 @@ class GatewayOperations: case ReadResourceRequest(params=params): return await _execute_read_resource(context, params, self._host_progress_callback) case _: - assert_never(operation) + return assert_never(operation) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 33978bb9182..a57eed04419 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -416,7 +416,10 @@ 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", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py index 870f2cd5f8f..abb925ddc77 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_operations.py @@ -8,6 +8,28 @@ from litellm.proxy._experimental.mcp_server.operations import GatewayOperations, 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 5c6854510400ca53a8718401adf0ca911f65e62a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 21 Sep 2026 15:00:54 -0700 Subject: [PATCH 7/8] test: fix stale budget-status and bad-database-url assertions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two CI checks were asserting behaviour the proxy no longer has. Neither was catching anything; both now fail for the right reason. budget_exceeded (tests/otel_tests/test_e2e_budgeting.py) bf804f5188 made 422 the default for budget refusals and added budget_exceeded_status_code to restore 429 for callers that need it. The e2e budget tests still asserted 429, so all six have been failing on a status change that was deliberate. Assert 422, the documented default, rather than reading litellm.budget_exceeded_status_code back — a test that asks the code what it does would have passed straight through this change and through the next one. The helpers also caught bare Exception, so a connection error reached `e.body` and surfaced as an AttributeError instead of a failed assertion. Narrow both to openai.APIStatusError, which is what a refusal actually raises (UnprocessableEntityError for 422, RateLimitError for 429), and let anything else propagate as itself. test_bad_database_url (.circleci/config.yml) The check required "Database setup failed after multiple retries", which only the v1 resolver emits, OR uvicorn's "Application startup failed. Exiting.". With v2 the default, the first branch is dead and the whole assertion rests on an incidental uvicorn line that HEAD's run did not emit at all. Assert the behaviour instead of the wording: the container exits non-zero, the log names the unreachable server (P1001), it never reaches "Application startup complete", and it is not left running. The exit code was previously discarded by `|| true`, so the one thing the job most needed to check was never checked. Verified by running the bad-DATABASE_URL container: exit 3, P1001 present, no startup-complete line, container stopped — the new check passes and the old one passed only by the uvicorn line's accident. --- .circleci/config.yml | 23 ++++++++++++----------- tests/otel_tests/test_e2e_budgeting.py | 13 +++++++------ 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index cc9aa7fe1c4..c69ed6a3510 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3033,28 +3033,29 @@ jobs: - run: name: Run Docker container with bad DATABASE_URL command: | + set +e docker run --name my-app \ -p 4000:4000 \ -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \ -e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \ -e DATABASE_URL="postgresql://wrong:wrong@wrong:5432/wrong" \ myapp:latest \ - --port 4000 > docker_output.log 2>&1 || true + --port 4000 > docker_output.log 2>&1 + echo "$?" > docker_exit_code + set -e - run: name: Display Docker logs command: cat docker_output.log - run: - name: Check for expected error + name: Proxy must refuse to serve on an unreachable database command: | - if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \ - (grep -q "Database setup failed after multiple retries" docker_output.log || \ - grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then - echo "Expected error found. Test passed." - else - echo "Expected error not found. Test failed." - cat docker_output.log - exit 1 - fi + fail() { echo "FAILED: $1"; cat docker_output.log; exit 1; } + exit_code="$(cat docker_exit_code)" + [ "$exit_code" -ne 0 ] || fail "proxy exited 0 with an unreachable database" + grep -q "P1001" docker_output.log || fail "log does not name the unreachable database server" + ! grep -q "Application startup complete" docker_output.log || fail "proxy reached serving state" + ! docker exec my-app true 2>/dev/null || fail "container is still running" + echo "Proxy refused to serve (exit $exit_code) and never reached startup. Test passed." provider_replay_harness: docker: diff --git a/tests/otel_tests/test_e2e_budgeting.py b/tests/otel_tests/test_e2e_budgeting.py index ca5058818e4..ae8f0ddc3ec 100644 --- a/tests/otel_tests/test_e2e_budgeting.py +++ b/tests/otel_tests/test_e2e_budgeting.py @@ -5,6 +5,7 @@ import uuid from typing import Any, Optional import aiohttp +import openai import pytest from httpx import AsyncClient @@ -23,7 +24,7 @@ async def make_calls_until_budget_exceeded(session, key: str, call_function, **k call_count += 1 await asyncio.sleep(0.1) # allow spend tracking to catch up pytest.fail(f"Budget was not exceeded after {MAX_CALLS} calls") - except Exception as e: + except openai.APIStatusError as e: print("vars: ", vars(e)) print("e.body: ", e.body) @@ -32,8 +33,8 @@ async def make_calls_until_budget_exceeded(session, key: str, call_function, **k # Check error structure and values that should be consistent assert ( - error_dict["code"] == "429" - ), f"Expected error code 429, got: {error_dict['code']}" + error_dict["code"] == "422" + ), f"Expected error code 422, got: {error_dict['code']}" assert ( error_dict["type"] == "budget_exceeded" ), f"Expected error type budget_exceeded, got: {error_dict['type']}" @@ -506,9 +507,9 @@ async def make_calls_until_team_budget_exceeded_cli_sso( call_count += 1 await asyncio.sleep(0.1) pytest.fail(f"Budget was not exceeded after {MAX_CALLS} calls") - except Exception as e: + except openai.APIStatusError as e: error_dict = e.body - assert error_dict["code"] == "429" + assert error_dict["code"] == "422" assert error_dict["type"] == "budget_exceeded" message = error_dict["message"] assert "Budget has been exceeded!" in message @@ -556,7 +557,7 @@ async def test_team_budget_enforcement_cli_sso_token(): 1. Create team with a tiny max_budget and a user on that team 2. Obtain a CLI SSO JWT (HTTP poll flow when Redis is shared, else mint) 3. Make chat completion calls until the team budget is exceeded - 4. Verify HTTP 429 budget_exceeded names the team + 4. Verify HTTP 422 budget_exceeded names the team """ user_id = f"cli-budget-user-{uuid.uuid4().hex[:8]}" user_email = f"{user_id}@example.com" From cf7988a2df65f6e775426bde8656cf4840f09f56 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Mon, 21 Sep 2026 22:31:01 +0000 Subject: [PATCH 8/8] chore(prices): sync OpenRouter prices: 3 models openrouter/~z-ai/glm-flash-latest: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/deepseek/deepseek-v4-pro: input_cost_per_token, output_cost_per_token, cache_read_input_token_cost openrouter/z-ai/glm-5.3-flash: max_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, cache_read_input_token_cost --- ...odel_prices_and_context_window_backup.json | 26 +++++++++---------- model_prices_and_context_window.json | 26 +++++++++---------- 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6b49b1d47a5..beb90c2fb8c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -43037,21 +43037,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.95578e-07, + "input_cost_per_token": 8.92272e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.791156e-06, + "output_cost_per_token": 1.784544e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.46315e-08, + "cache_read_input_token_cost": 7.4356e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -68212,13 +68212,13 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3-flash": { - "input_cost_per_token": 7.5e-08, - "output_cost_per_token": 2.5e-07, - "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 102400, - "max_tokens": 102400, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -73252,14 +73252,14 @@ "supports_web_search": true }, "openrouter/~z-ai/glm-flash-latest": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 7.5e-08, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 102400, - "max_tokens": 102400, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 2.5e-07, + "output_cost_per_token": 5e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6b49b1d47a5..beb90c2fb8c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -43037,21 +43037,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.95578e-07, + "input_cost_per_token": 8.92272e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.791156e-06, + "output_cost_per_token": 1.784544e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.46315e-08, + "cache_read_input_token_cost": 7.4356e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -68212,13 +68212,13 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3-flash": { - "input_cost_per_token": 7.5e-08, - "output_cost_per_token": 2.5e-07, - "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 5e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 102400, - "max_tokens": 102400, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -73252,14 +73252,14 @@ "supports_web_search": true }, "openrouter/~z-ai/glm-flash-latest": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 7.5e-08, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 102400, - "max_tokens": 102400, + "max_output_tokens": 943718, + "max_tokens": 943718, "mode": "chat", - "output_cost_per_token": 2.5e-07, + "output_cost_per_token": 5e-07, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true,