From 754a2afe12891f0fc9e8b20fbf17f7504d71215a Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 8 Sep 2026 17:30:20 -0700 Subject: [PATCH] feat(mcp): add schema discovery proxy mode (#40298) --- .../_experimental/mcp_server/mcp_context.py | 3 + .../proxy/_experimental/mcp_server/server.py | 100 ++++- .../_experimental/mcp_server/tool_search.py | 265 ++++++++++++- litellm/proxy/_lazy_openapi_snapshot.json | 128 +++++++ litellm/proxy/_types.py | 1 + litellm/proxy/proxy_server.py | 25 ++ tests/mcp_tests/test_proxy_mcp_e2e.py | 152 ++++++-- .../mcp_server/test_mcp_proxy_mode.py | 354 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 184 +++++++++ 9 files changed, 1154 insertions(+), 58 deletions(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_context.py b/litellm/proxy/_experimental/mcp_server/mcp_context.py index 7936edad753..74cc0c900d9 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_context.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_context.py @@ -21,3 +21,6 @@ _mcp_gateway_initialize_instructions: Final[ContextVar[str | None]] = ContextVar # Per-request scoped server name; set in MCP HTTP/SSE handlers when the path # identifies exactly one upstream server. Never populated from client-supplied headers. _mcp_gateway_server_name: Final[ContextVar[str | None]] = ContextVar("_mcp_gateway_server_name", default=None) + +# Set server-side by the /mcp/proxy route. Never populated from client-supplied headers. +_mcp_proxy_mode: Final[ContextVar[bool]] = ContextVar("_mcp_proxy_mode", default=False) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 1b1e77d27ef..44c00df996c 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -15,7 +15,7 @@ import types import uuid from collections.abc import AsyncIterator, Callable, Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol import httpx from fastapi import FastAPI, HTTPException @@ -47,6 +47,7 @@ from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_active_toolset_id, _mcp_gateway_initialize_instructions, _mcp_gateway_server_name, + _mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # server-owned request mode ) from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug from litellm.proxy._experimental.mcp_server.oauth_utils import ( @@ -537,11 +538,22 @@ if MCP_AVAILABLE: notification_options: NotificationOptions | None = None, experimental_capabilities: dict[str, dict[str, object]] | None = None, ) -> InitializationOptions: - opts: Final = Server.create_initialization_options( + base_options: Final = Server.create_initialization_options( self, notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, ) + opts: Final = ( + base_options.model_copy( + update={ # mutable-ok: Pydantic update payload + "capabilities": base_options.capabilities.model_copy( + update={"prompts": None, "resources": None} # mutable-ok: Pydantic update payload + ) + } + ) + if _mcp_proxy_mode.get() + else base_options + ) updates: Final[dict[str, str]] = {} merged: Final = _mcp_gateway_initialize_instructions.get() if merged is not None: @@ -822,17 +834,20 @@ if MCP_AVAILABLE: "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 _mcp_proxy_mode.get(): + return [Tool.model_validate(d) for d in get_mcp_proxy_tool_definitions()] # mutable-ok: MCP SDK list if getattr( getattr(user_api_key_auth, "object_permission", None), "mcp_tool_search_enabled", False, ): - from mcp.types import Tool - - from litellm.proxy._experimental.mcp_server.tool_search import ( - get_virtual_tool_definitions, - ) - return [Tool.model_validate(d) for d in get_virtual_tool_definitions()] # Get mcp_servers from context variable @@ -906,6 +921,12 @@ if MCP_AVAILABLE: verbose_logger.debug("Host progressToken captured: %s...", str(host_token)[:8]) return forward_progress + def _reject_mcp_proxy_operation() -> NoReturn: + from mcp.shared.exceptions import McpError + from mcp.types import METHOD_NOT_FOUND, ErrorData + + raise McpError(ErrorData(code=METHOD_NOT_FOUND, message="Operation unavailable on /mcp/proxy")) + async def _build_virtual_call_logging_obj( name: str, arguments: dict[str, object], @@ -961,16 +982,53 @@ if MCP_AVAILABLE: 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") + ], + isError=True, + ) + + if _mcp_proxy_mode.get() and name in MCP_PROXY_TOOL_NAMES: + assert user_api_key_auth is not None + 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 + ) + return 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, + ) + if name not in VIRTUAL_TOOL_NAMES: return None @@ -1216,6 +1274,8 @@ if MCP_AVAILABLE: """ List all available prompts """ + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) @@ -1273,8 +1333,8 @@ if MCP_AVAILABLE: Returns: GetPromptResult: Getting prompt execution results """ - - # Validate arguments + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) @@ -1311,6 +1371,8 @@ if MCP_AVAILABLE: @server.list_resources() async def list_resources() -> list[Resource]: """List all available resources.""" + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) @@ -1355,6 +1417,8 @@ if MCP_AVAILABLE: @server.list_resource_templates() async def list_resource_templates() -> list[ResourceTemplate]: """List all available resource templates.""" + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) @@ -1400,6 +1464,8 @@ if MCP_AVAILABLE: @server.read_resource() async def read_resource(url: AnyUrl) -> list[ReadResourceContents]: + if _mcp_proxy_mode.get(): + _reject_mcp_proxy_operation() from mcp.server.lowlevel.server import request_ctx req_ctx: Final = request_ctx.get(None) @@ -1998,6 +2064,7 @@ if MCP_AVAILABLE: 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. @@ -2177,9 +2244,14 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, ) - # Apply display-name/description overrides last so that - # permission filtering always works against original names. - filtered_tools = apply_tool_overrides(filtered_tools, server) + 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", @@ -2491,6 +2563,7 @@ if MCP_AVAILABLE: 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. @@ -2520,6 +2593,7 @@ if MCP_AVAILABLE: 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 diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 24e2f5ce64d..2c73f9b863b 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -1,5 +1,6 @@ from __future__ import annotations +import hashlib import json from collections.abc import Mapping, Sequence from dataclasses import dataclass @@ -30,6 +31,12 @@ if TYPE_CHECKING: MCP_TOOL_SEARCH_SETTINGS_KEY: Final[str] = "mcp_tool_search" MCP_TOOL_SEARCH_TOOL_NAME: Final[str] = "mcp_tool_search" MCP_TOOL_CALL_TOOL_NAME: Final[str] = "mcp_tool_call" +MCP_PROXY_SEARCH_TOOL_NAME: Final[str] = "search_tools" +MCP_PROXY_SCHEMA_TOOL_NAME: Final[str] = "get_tool_schema" +MCP_PROXY_CALL_TOOL_NAME: Final[str] = "call_tool" +MCP_PROXY_TOOL_NAMES: Final = frozenset( + (MCP_PROXY_SEARCH_TOOL_NAME, MCP_PROXY_SCHEMA_TOOL_NAME, MCP_PROXY_CALL_TOOL_NAME) +) AGENT_SEARCH_TOOL_NAME: Final[str] = "agent_search" SKILL_SEARCH_TOOL_NAME: Final[str] = "skill_search" VIRTUAL_TOOL_NAMES: Final = frozenset( @@ -51,6 +58,29 @@ class ToolSearchResult(TypedDict, total=False): score: ReadOnly[float] +class MCPProxySearchResult(TypedDict, total=False): + tool_id: Required[ReadOnly[str]] + name: Required[ReadOnly[str]] + description: Required[ReadOnly[str]] + score: ReadOnly[float] + + +class MCPProxySchemaResult(MCPProxySearchResult, total=False): + inputSchema: Required[ReadOnly[Mapping[str, object]]] + outputSchema: ReadOnly[Mapping[str, object]] + + +class MCPProxyToolIdentity(TypedDict): + server_id: ReadOnly[str] + tool_name: ReadOnly[str] + + +@dataclass(frozen=True, slots=True) +class MCPToolSearchHit: + tool: Tool + score: float | None = None + + @dataclass(frozen=True, slots=True) class SemanticToolRanker: embed: Embedder @@ -76,6 +106,55 @@ def _scored_result(tool: Tool, score: float) -> ToolSearchResult: return {"name": tool.name, "description": tool.description or "", "inputSchema": tool.inputSchema, "score": score} +_MCP_PROXY_IDENTITY_META_KEY: Final[str] = "litellm.ai/proxy_tool_identity" + + +def with_mcp_proxy_identity(tool: Tool, server_id: str) -> Tool: + identity: Final[MCPProxyToolIdentity] = {"server_id": server_id, "tool_name": tool.name} + return tool.model_copy( # mutable-ok: Pydantic requires mutable update and metadata mappings + update={ # mutable-ok: Pydantic update payload + "meta": {**(tool.meta or {}), _MCP_PROXY_IDENTITY_META_KEY: identity} # mutable-ok: metadata mapping + } + ) + + +def _mcp_proxy_identity(tool: Tool) -> MCPProxyToolIdentity: + identity: Final = (tool.meta or {}).get(_MCP_PROXY_IDENTITY_META_KEY) # mutable-ok: absent metadata default + if not isinstance(identity, Mapping): + raise TypeError("MCP proxy tool identity is missing") + server_id: Final = identity.get("server_id") + tool_name: Final = identity.get("tool_name") + if not isinstance(server_id, str) or not isinstance(tool_name, str): + raise TypeError("MCP proxy tool identity is invalid") + return {"server_id": server_id, "tool_name": tool_name} # mutable-ok: TypedDict identity payload + + +def mcp_proxy_tool_id(tool: Tool) -> str: + identity: Final = _mcp_proxy_identity(tool) + return hashlib.sha256(f"{identity['server_id']}\0{identity['tool_name']}".encode()).hexdigest()[:32] + + +def _proxy_search_result(hit: MCPToolSearchHit) -> MCPProxySearchResult: + base: Final[MCPProxySearchResult] = { + "tool_id": mcp_proxy_tool_id(hit.tool), + "name": hit.tool.name, + "description": hit.tool.description or "", + } + return {**base, "score": hit.score} if hit.score is not None else base # mutable-ok: wire result payload + + +def _proxy_schema_result(tool: Tool) -> MCPProxySchemaResult: + base: Final[MCPProxySchemaResult] = { + "tool_id": mcp_proxy_tool_id(tool), + "name": tool.name, + "description": tool.description or "", + "inputSchema": tool.inputSchema, + } + if tool.outputSchema is None: + return base + return {**base, "outputSchema": tool.outputSchema} # mutable-ok: wire schema payload + + def _tool_text(tool: Tool) -> str: return "\n".join(part for part in (tool.name, tool.description or "") if part) @@ -107,6 +186,38 @@ def search_tools(query: str, tools: Sequence[Tool], top_k: int = 5) -> tuple[Too return tuple(_tool_result(tool) for _, tool in _top_hits(tools, scores, minimum=1.0, limit=top_k)) +async def rank_mcp_tools( + query: str, + tools: Sequence[Tool], + top_k: int, + settings: MCPToolSearchSettings, + ranker: SemanticToolRanker | None, +) -> tuple[MCPToolSearchHit, ...] | EmbeddingFailed: + core, rest = _split_core_tools(tools, settings.core_tools) + core_hits: Final = tuple(MCPToolSearchHit(tool) for tool in core) + if not query: + return core_hits + limit: Final = min(top_k, settings.top_k) + if ranker is None: + scores: Final = tuple(_keyword_score(query, tool) for tool in rest) + return ( + *core_hits, + *(MCPToolSearchHit(tool) for _, tool in _top_hits(rest, scores, minimum=1.0, limit=limit)), + ) + semantic_scores: Final = await ranker.index.scores( + query, tuple(_tool_text(tool) for tool in rest), ranker.embed, ranker.embedding_model + ) + if isinstance(semantic_scores, EmbeddingFailed): + return semantic_scores + return ( + *core_hits, + *( + MCPToolSearchHit(tool, score) + for score, tool in _top_hits(rest, semantic_scores, settings.similarity_threshold, limit) + ), + ) + + async def search_mcp_tools( query: str, tools: Sequence[Tool], @@ -114,21 +225,12 @@ async def search_mcp_tools( settings: MCPToolSearchSettings, ranker: SemanticToolRanker | None, ) -> tuple[ToolSearchResult, ...] | EmbeddingFailed: - """Core tools the caller can access come first, then up to `top_k` ranked matches from the remaining tools.""" - core, rest = _split_core_tools(tools, settings.core_tools) - limit: Final = min(top_k, settings.top_k) - core_results: Final = tuple(_tool_result(tool) for tool in core) - if ranker is None: - return (*core_results, *search_tools(query, rest, limit)) - if not query: - return core_results - scores: Final = await ranker.index.scores( - query, tuple(_tool_text(tool) for tool in rest), ranker.embed, ranker.embedding_model + hits: Final = await rank_mcp_tools(query, tools, top_k, settings, ranker) + if isinstance(hits, EmbeddingFailed): + return hits + return tuple( + _scored_result(hit.tool, hit.score) if hit.score is not None else _tool_result(hit.tool) for hit in hits ) - if isinstance(scores, EmbeddingFailed): - return scores - hits: Final = _top_hits(rest, scores, minimum=settings.similarity_threshold, limit=limit) - return (*core_results, *(_scored_result(tool, score) for score, tool in hits)) class _ToolParamSchema(TypedDict, total=False): @@ -223,10 +325,48 @@ _SKILL_SEARCH_DEFINITION: Final[VirtualToolDefinition] = { } +_MCP_PROXY_SEARCH_DEFINITION: Final[VirtualToolDefinition] = { + "name": MCP_PROXY_SEARCH_TOOL_NAME, + "description": "Search accessible MCP tools by describing what you need. Returns opaque tool IDs.", + "inputSchema": { + "type": "object", + "properties": {"query": {"type": "string", "description": "What the tool should do."}}, + "required": _json_array("query"), + }, +} + +_MCP_PROXY_SCHEMA_DEFINITION: Final[VirtualToolDefinition] = { + "name": MCP_PROXY_SCHEMA_TOOL_NAME, + "description": "Return the complete schema for an accessible MCP tool ID.", + "inputSchema": { + "type": "object", + "properties": {"tool_id": {"type": "string", "description": "Opaque ID from search_tools."}}, + "required": _json_array("tool_id"), + }, +} + +_MCP_PROXY_CALL_DEFINITION: Final[VirtualToolDefinition] = { + "name": MCP_PROXY_CALL_TOOL_NAME, + "description": "Call an accessible MCP tool by opaque ID with schema-valid arguments.", + "inputSchema": { + "type": "object", + "properties": { + "tool_id": {"type": "string", "description": "Opaque ID from search_tools."}, + "arguments": {"type": "object", "description": "Arguments validated against the selected tool schema."}, + }, + "required": _json_array("tool_id"), + }, +} + + def get_virtual_tool_definitions() -> tuple[VirtualToolDefinition, ...]: return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION, _SKILL_SEARCH_DEFINITION) +def get_mcp_proxy_tool_definitions() -> tuple[VirtualToolDefinition, ...]: + return (_MCP_PROXY_SEARCH_DEFINITION, _MCP_PROXY_SCHEMA_DEFINITION, _MCP_PROXY_CALL_DEFINITION) + + def _text_tool_result(text: str, is_error: bool) -> CallToolResult: from mcp.types import CallToolResult, TextContent @@ -314,7 +454,9 @@ 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 + from litellm.proxy._experimental.mcp_server.server import ( + _list_mcp_tools, # pyright: ignore[reportPrivateUsage] # shared catalog owner + ) from litellm.proxy.proxy_server import llm_router, proxy_logging_obj settings: Final = mcp_tool_search_settings() @@ -351,6 +493,97 @@ async def handle_mcp_tool_search( return _text_tool_result(json.dumps(results), is_error=False) +async def handle_mcp_proxy_tool( + name: str, + arguments: dict[str, object], # mutable-ok: MCP dispatcher passes mutable call arguments + user_api_key_dict: UserAPIKeyAuth, + client_ip: str | None = None, + mcp_servers: list[str] | None = None, # mutable-ok: preserve MCP scope container for existing resolver + mcp_auth_header: str | None = None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, # mutable-ok: preserve forwarded headers + oauth2_headers: dict[str, str] | None = None, # mutable-ok: preserve forwarded headers + raw_headers: dict[str, str] | None = None, # mutable-ok: preserve request headers + litellm_logging_obj: LiteLLMLoggingObj | None = None, +) -> CallToolResult: + from fastapi import HTTPException + from jsonschema import ValidationError as JsonSchemaValidationError + 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 + ) + + listing: Final = await _list_mcp_tools( + user_api_key_auth=user_api_key_dict, + mcp_servers=mcp_servers, + client_ip=client_ip, + 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=True, + ) + tools_by_id: Final = {mcp_proxy_tool_id(tool): tool for tool in listing.tools} # mutable-ok: lookup index + + if name == MCP_PROXY_SEARCH_TOOL_NAME: + llm_router: Final = proxy_server.llm_router + proxy_logging_obj: Final = proxy_server.proxy_logging_obj + settings: Final = mcp_tool_search_settings() + if isinstance(settings, ValidationError): + return _text_tool_result(str(settings), is_error=True) + if settings.embedding_model is not None and llm_router is None: + return _text_tool_result( + f"litellm_settings.{MCP_TOOL_SEARCH_SETTINGS_KEY}.embedding_model needs a model_list so it can be called", + is_error=True, + ) + ranker: Final = ( + SemanticToolRanker( + embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict, proxy_logging_obj), + embedding_model=settings.embedding_model, + index=global_mcp_tool_search_index, + ) + if settings.embedding_model is not None and llm_router is not None + else None + ) + results: Final = await rank_mcp_tools(str(arguments.get("query", "")), listing.tools, 5, settings, ranker) + if isinstance(results, EmbeddingFailed): + return _text_tool_result(results.reason, is_error=True) + return _text_tool_result(json.dumps(tuple(_proxy_search_result(hit) for hit in results)), is_error=False) + + tool_id: Final = arguments.get("tool_id") + tool: Final = tools_by_id.get(tool_id) if isinstance(tool_id, str) else None + if tool is None: + return _text_tool_result("Unknown or unauthorized tool_id", is_error=True) + + if name == MCP_PROXY_SCHEMA_TOOL_NAME: + return _text_tool_result(json.dumps(_proxy_schema_result(tool)), is_error=False) + if name != MCP_PROXY_CALL_TOOL_NAME: + raise HTTPException(status_code=400, detail=f"Unknown MCP proxy tool: {name}") + + tool_arguments: Final = arguments.get("arguments", {}) # mutable-ok: JSON Schema validator consumes mapping + if not isinstance(tool_arguments, dict): + return _text_tool_result("arguments must be an object", is_error=True) + try: + validate(instance=tool_arguments, schema=tool.inputSchema) + except JsonSchemaValidationError as exc: + return _text_tool_result(f"Invalid arguments: {exc.message}", is_error=True) + + return await handle_mcp_tool_call( + tool_name=_mcp_proxy_identity(tool)["tool_name"], + arguments=tool_arguments, + user_api_key_dict=user_api_key_dict, + requested_server_id=_mcp_proxy_identity(tool)["server_id"], + 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=litellm_logging_obj, + ) + + async def handle_mcp_tool_call( tool_name: str, arguments: dict[str, Any], @@ -362,6 +595,7 @@ async def handle_mcp_tool_call( oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, litellm_logging_obj: LiteLLMLoggingObj | None = None, + requested_server_id: str | None = None, ) -> CallToolResult: from litellm.proxy._experimental.mcp_server.server import ( _get_allowed_mcp_servers, @@ -400,4 +634,5 @@ async def handle_mcp_tool_call( oauth2_headers=oauth2_headers, raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, + requested_server_id=requested_server_id, ) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 71475320c2c..4d5cdcc8003 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -17027,6 +17027,134 @@ "mcp_app" ] } + }, + "/mcp/proxy": { + "delete": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_delete", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "get": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "head": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_head", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "options": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_options", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "patch": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_patch", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "post": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + }, + "put": { + "description": "Serve the fixed three-tool MCP proxy surface.", + "operationId": "proxy_mcp_route_mcp_proxy_put", + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + } + }, + "summary": "Proxy Mcp Route", + "tags": [ + "mcp_app" + ] + } } } }, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1c8608305a2..d746cfd38d8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -502,6 +502,7 @@ class LiteLLMRoutes(enum.Enum): mcp_inference_routes = [ "/mcp", "/mcp/", + "/mcp/proxy", "/mcp/{subpath}", "/mcp/tools", "/mcp/tools/list", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a617aec9f5c..b22e9bba15b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -18484,6 +18484,31 @@ async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "Streami ######################################################## +@app.api_route( + "/mcp/proxy", + methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], # mutable-ok: FastAPI route methods +) +async def proxy_mcp_route(request: Request) -> Response: + """Serve the fixed three-tool MCP proxy surface.""" + from litellm.proxy._experimental.mcp_server.mcp_context import ( # pyright: ignore[reportPrivateUsage] # route-owned mode + _mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # route-owned mode + ) + from litellm.proxy._experimental.mcp_server.server import handle_streamable_http_mcp + from litellm.proxy._experimental.mcp_server.utils import is_mcp_available + + if not is_mcp_available(): + raise HTTPException(status_code=404, detail="Not Found") + + token: Final = _mcp_proxy_mode.set(True) + try: + scope: Final = dict(request.scope) # mutable-ok: ASGI scope rewrite + scope["_original_path"] = scope.get("path", "") + scope["path"] = BASE_MCP_ROUTE + return await _stream_mcp_asgi_response(handle_streamable_http_mcp, scope, request.receive) + finally: + _mcp_proxy_mode.reset(token) + + @app.api_route( BASE_MCP_ROUTE, methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], diff --git a/tests/mcp_tests/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py index 2dd57e13d3b..a97eed82e18 100644 --- a/tests/mcp_tests/test_proxy_mcp_e2e.py +++ b/tests/mcp_tests/test_proxy_mcp_e2e.py @@ -1,4 +1,5 @@ import asyncio +import json import os import socket import subprocess @@ -134,15 +135,11 @@ def math_streamable_http_server() -> str: @pytest.fixture(scope="session") -def proxy_server_url( - tmp_path_factory: pytest.TempPathFactory, math_streamable_http_server: str -): +def proxy_server_url(tmp_path_factory: pytest.TempPathFactory, math_streamable_http_server: str): config_dir = tmp_path_factory.mktemp("mcp_e2e") config_path = config_dir / "config.yaml" config = yaml.safe_load(CONFIG_TEMPLATE_PATH.read_text()) - config["mcp_servers"]["math_streamable_http"][ - "url" - ] = f"{math_streamable_http_server}/mcp" + config["mcp_servers"]["math_streamable_http"]["url"] = f"{math_streamable_http_server}/mcp" config_path.write_text(yaml.safe_dump(config)) server_url, server, thread, sock = _start_proxy_server(str(config_path)) @@ -177,9 +174,7 @@ class TestProxyMcpSimpleConnections: assert text == "7" @pytest.mark.asyncio - async def test_proxy_mcp_streamable_http_roundtrip( - self, proxy_server_url: str - ) -> None: + async def test_proxy_mcp_streamable_http_roundtrip(self, proxy_server_url: str) -> None: async with asyncio.timeout(20): async with streamablehttp_client( url=f"{proxy_server_url}/mcp", @@ -200,9 +195,7 @@ class TestProxyMcpSimpleConnections: assert text == "11" @pytest.mark.asyncio - async def test_proxy_mcp_lists_all_servers_without_header( - self, proxy_server_url: str - ) -> None: + async def test_proxy_mcp_lists_all_servers_without_header(self, proxy_server_url: str) -> None: async with asyncio.timeout(20): async with streamablehttp_client( url=f"{proxy_server_url}/mcp", @@ -220,20 +213,14 @@ class TestProxyMcpSimpleConnections: } assert expected_tool_names <= tool_names - async def _call_and_get_text( - tool_name: str, *, a: int, b: int - ) -> str | None: - result = await session.call_tool( - tool_name, arguments={"a": a, "b": b} - ) + async def _call_and_get_text(tool_name: str, *, a: int, b: int) -> str | None: + result = await session.call_tool(tool_name, arguments={"a": a, "b": b}) assert result.content first_content = result.content[0] return getattr(first_content, "text", None) stdio_result = await _call_and_get_text("math_stdio-add", a=2, b=3) - streamable_result = await _call_and_get_text( - "math_streamable_http-add", a=4, b=5 - ) + streamable_result = await _call_and_get_text("math_streamable_http-add", a=4, b=5) assert stdio_result == "5" assert streamable_result == "9" @@ -254,9 +241,7 @@ class TestProxyMcpStatelessBehavior: """ @pytest.mark.asyncio - async def test_independent_clients_no_shared_session( - self, proxy_server_url: str - ) -> None: + async def test_independent_clients_no_shared_session(self, proxy_server_url: str) -> None: """Two independent clients connect and operate without sharing session state.""" async with asyncio.timeout(30): # --- Client A: connect, initialize, call tool --- @@ -269,9 +254,7 @@ class TestProxyMcpStatelessBehavior: ) as (read_a, write_a, _get_sid_a): async with ClientSession(read_a, write_a) as session_a: await session_a.initialize() - result_a = await session_a.call_tool( - "add", arguments={"a": 10, "b": 20} - ) + result_a = await session_a.call_tool("add", arguments={"a": 10, "b": 20}) assert result_a.content text_a = getattr(result_a.content[0], "text", None) assert text_a == "30" @@ -293,9 +276,118 @@ class TestProxyMcpStatelessBehavior: await session_b.initialize() tools = await session_b.list_tools() assert any(t.name.endswith("add") for t in tools.tools) - result_b = await session_b.call_tool( - "add", arguments={"a": 100, "b": 200} - ) + result_b = await session_b.call_tool("add", arguments={"a": 100, "b": 200}) assert result_b.content text_b = getattr(result_b.content[0], "text", None) assert text_b == "300" + + +PROXY_MODE_TOOLS = frozenset({"search_tools", "get_tool_schema", "call_tool"}) + + +def _payload(result: typing.Any) -> typing.Any: + assert result.content, f"empty tool result: {result}" + return json.loads(result.content[0].text) + + +def _proxy_session(proxy_server_url: str, **extra_headers: str): + return streamablehttp_client( + url=f"{proxy_server_url}/mcp/proxy", + headers={"Authorization": PROXY_AUTHORIZATION_HEADER, **extra_headers}, + ) + + +class TestProxyMcpSchemaDiscoveryMode: + """Drive /mcp/proxy over the real streamable-HTTP transport with the MCP SDK client: + the fixed three-tool surface, opaque-id discovery, schema-validated execution against + two upstreams that expose the same tool name, and the operations the surface refuses.""" + + @pytest.mark.asyncio + async def test_initialize_and_list_expose_only_discovery_tools(self, proxy_server_url: str) -> None: + async with asyncio.timeout(20): + async with _proxy_session(proxy_server_url) as (read, write, _sid): + async with ClientSession(read, write) as session: + init = await session.initialize() + assert init.capabilities.tools is not None + assert init.capabilities.prompts is None + assert init.capabilities.resources is None + + listed = await session.list_tools() + assert {tool.name for tool in listed.tools} == PROXY_MODE_TOOLS + + @pytest.mark.asyncio + async def test_search_schema_and_call_round_trip_keeps_server_identity(self, proxy_server_url: str) -> None: + async with asyncio.timeout(30): + async with _proxy_session(proxy_server_url) as (read, write, _sid): + async with ClientSession(read, write) as session: + await session.initialize() + + hits = _payload(await session.call_tool("search_tools", arguments={"query": "add"})) + by_name = {hit["name"]: hit for hit in hits} + assert {"math_stdio-add", "math_streamable_http-add"} <= set(by_name) + assert all("inputSchema" not in hit for hit in hits) + assert by_name["math_stdio-add"]["tool_id"] != by_name["math_streamable_http-add"]["tool_id"] + + schema = _payload( + await session.call_tool( + "get_tool_schema", arguments={"tool_id": by_name["math_stdio-add"]["tool_id"]} + ) + ) + assert schema["name"] == "math_stdio-add" + assert set(schema["inputSchema"]["required"]) == {"a", "b"} + assert schema["outputSchema"]["properties"]["result"]["type"] == "integer" + + stdio = await session.call_tool( + "call_tool", + arguments={"tool_id": by_name["math_stdio-add"]["tool_id"], "arguments": {"a": 3, "b": 4}}, + ) + http = await session.call_tool( + "call_tool", + arguments={ + "tool_id": by_name["math_streamable_http-add"]["tool_id"], + "arguments": {"a": 5, "b": 6}, + }, + ) + assert stdio.isError is False and stdio.content[0].text == "7" + assert http.isError is False and http.content[0].text == "11" + + @pytest.mark.asyncio + async def test_server_scope_header_narrows_discovery(self, proxy_server_url: str) -> None: + async with asyncio.timeout(20): + async with _proxy_session(proxy_server_url, **{"x-mcp-servers": "math_streamable_http"}) as ( + read, + write, + _sid, + ): + async with ClientSession(read, write) as session: + await session.initialize() + hits = _payload(await session.call_tool("search_tools", arguments={"query": "add"})) + assert {hit["name"] for hit in hits} == {"math_streamable_http-add"} + + @pytest.mark.asyncio + async def test_rejections_never_reach_upstream(self, proxy_server_url: str) -> None: + from mcp.shared.exceptions import McpError + from mcp.types import METHOD_NOT_FOUND + + async with asyncio.timeout(30): + async with _proxy_session(proxy_server_url) as (read, write, _sid): + async with ClientSession(read, write) as session: + await session.initialize() + hits = _payload(await session.call_tool("search_tools", arguments={"query": "add"})) + tool_id = next(hit["tool_id"] for hit in hits if hit["name"] == "math_stdio-add") + + bad_args = await session.call_tool( + "call_tool", arguments={"tool_id": tool_id, "arguments": {"a": "three", "b": 4}} + ) + assert bad_args.isError is True and "Invalid arguments" in bad_args.content[0].text + + stale = await session.call_tool("get_tool_schema", arguments={"tool_id": "0" * 32}) + assert stale.isError is True and "unauthorized tool_id" in stale.content[0].text + + direct = await session.call_tool("math_stdio-add", arguments={"a": 1, "b": 2}) + assert direct.isError is True and "unavailable on /mcp/proxy" in direct.content[0].text + + for operation in (session.list_prompts, session.list_resources): + with pytest.raises(McpError) as refused: + await operation() + assert refused.value.error.code == METHOD_NOT_FOUND 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 new file mode 100644 index 00000000000..e86fc27fed7 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py @@ -0,0 +1,354 @@ +import json +from collections.abc import Iterator +from unittest.mock import AsyncMock, patch + +import pytest +from mcp.types import CallToolResult, TextContent, Tool + +from litellm.proxy._experimental.mcp_server import server +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing +from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager +from litellm.proxy._experimental.mcp_server.tool_search import ( + MCP_PROXY_CALL_TOOL_NAME, + MCP_PROXY_SCHEMA_TOOL_NAME, + MCP_PROXY_SEARCH_TOOL_NAME, + handle_mcp_proxy_tool, + mcp_proxy_tool_id, + with_mcp_proxy_identity, +) +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth +from litellm.types.mcp import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +TOOL = Tool.model_validate( + { + "name": "math_stdio-add", + "description": "Add two numbers", + "inputSchema": { + "type": "object", + "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}}, + "required": ["a", "b"], + }, + "outputSchema": {"type": "object"}, + "_meta": {"litellm.ai/proxy_tool_identity": {"server_id": "server-1", "tool_name": "math_stdio-add"}}, + } +) +AUTH = UserAPIKeyAuth(api_key="key") + + +def _text(result: CallToolResult) -> object: + return json.loads(result.content[0].text) + + +@pytest.mark.asyncio +async def test_proxy_search_returns_opaque_id_and_schema() -> None: + with ( + patch( # test-quality-ok: isolate authorized catalog owner + "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + new_callable=AsyncMock, + return_value=AggregateToolListing(tools=[TOOL], outcomes={}), + ), + ): + result = await handle_mcp_proxy_tool(MCP_PROXY_SEARCH_TOOL_NAME, {"query": "add"}, AUTH) + + item = _text(result)[0] + assert item["tool_id"] == mcp_proxy_tool_id(TOOL) + assert item["name"] == TOOL.name + assert "inputSchema" not in item + assert "outputSchema" not in item + assert len(item["tool_id"]) == 32 + + +@pytest.mark.asyncio +async def test_proxy_schema_and_call_resolve_current_authorized_catalog() -> None: + executed = CallToolResult(content=[TextContent(type="text", text="3")], isError=False) + with ( + patch( # test-quality-ok: isolate authorized catalog owner + "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + new_callable=AsyncMock, + return_value=AggregateToolListing(tools=[TOOL], outcomes={}), + ), + patch( # test-quality-ok: isolate execution delegate seam + "litellm.proxy._experimental.mcp_server.tool_search.handle_mcp_tool_call", + new_callable=AsyncMock, + return_value=executed, + ) as call, + ): + schema = await handle_mcp_proxy_tool( + MCP_PROXY_SCHEMA_TOOL_NAME, + {"tool_id": mcp_proxy_tool_id(TOOL)}, + AUTH, + ) + result = await handle_mcp_proxy_tool( + MCP_PROXY_CALL_TOOL_NAME, + {"tool_id": mcp_proxy_tool_id(TOOL), "arguments": {"a": 1, "b": 2}}, + AUTH, + ) + + assert _text(schema)["inputSchema"] == TOOL.inputSchema + assert result is executed + assert call.await_args.kwargs["tool_name"] == TOOL.name + assert call.await_args.kwargs["arguments"] == {"a": 1, "b": 2} + assert call.await_args.kwargs["requested_server_id"] == "server-1" + + +@pytest.mark.asyncio +async def test_proxy_rejects_stale_id_and_invalid_arguments_before_dispatch() -> None: + with ( + patch( # test-quality-ok: isolate authorized catalog owner + "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + new_callable=AsyncMock, + return_value=AggregateToolListing(tools=[TOOL], outcomes={}), + ), + patch( # test-quality-ok: isolate execution delegate seam + "litellm.proxy._experimental.mcp_server.tool_search.handle_mcp_tool_call", new_callable=AsyncMock + ) as call, + ): + stale = await handle_mcp_proxy_tool(MCP_PROXY_SCHEMA_TOOL_NAME, {"tool_id": "stale"}, AUTH) + invalid = await handle_mcp_proxy_tool( + MCP_PROXY_CALL_TOOL_NAME, + {"tool_id": mcp_proxy_tool_id(TOOL), "arguments": "wrong"}, + AUTH, + ) + falsy = await handle_mcp_proxy_tool( + MCP_PROXY_CALL_TOOL_NAME, + {"tool_id": mcp_proxy_tool_id(TOOL), "arguments": False}, + AUTH, + ) + invalid_schema = await handle_mcp_proxy_tool( + MCP_PROXY_CALL_TOOL_NAME, + {"tool_id": mcp_proxy_tool_id(TOOL), "arguments": {"a": "wrong"}}, + AUTH, + ) + + assert stale.isError is True + assert invalid.isError is True + assert falsy.isError is True + assert invalid_schema.isError is True + call.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_proxy_call_builds_logging_object() -> None: + from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_proxy_mode + + sentinel = object() + result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + token = _mcp_proxy_mode.set(True) + try: + with ( + patch.object( # test-quality-ok: isolate logging pipeline seam + server, "_build_virtual_call_logging_obj", new_callable=AsyncMock, return_value=sentinel + ) as build, + patch( # test-quality-ok: isolate proxy dispatch seam + "litellm.proxy._experimental.mcp_server.tool_search.handle_mcp_proxy_tool", + new_callable=AsyncMock, + return_value=result, + ) as handle, + ): + actual = await server._dispatch_virtual_mcp_tool( + name=MCP_PROXY_CALL_TOOL_NAME, + arguments={"tool_id": "id", "arguments": {}}, + user_api_key_auth=AUTH, + client_ip=None, + ) + finally: + _mcp_proxy_mode.reset(token) + + assert actual is result + build.assert_awaited_once() + assert handle.await_args.kwargs["litellm_logging_obj"] is sentinel + + +@pytest.mark.asyncio +async def test_proxy_call_rejects_non_proxy_tool_names() -> None: + from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_proxy_mode + + token = _mcp_proxy_mode.set(True) + try: + result = await server._dispatch_virtual_mcp_tool( + name="math_stdio-add", + arguments={"a": 1, "b": 2}, + user_api_key_auth=AUTH, + client_ip=None, + ) + finally: + _mcp_proxy_mode.reset(token) + + assert result is not None + assert result.isError is True + assert "unavailable" in result.content[0].text + + +@pytest.mark.asyncio +async def test_proxy_rejects_non_tool_protocol_operations() -> None: + from mcp.shared.exceptions import McpError + from pydantic import AnyUrl + + from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_proxy_mode + + token = _mcp_proxy_mode.set(True) + try: + with pytest.raises(McpError): + await server.list_prompts() + with pytest.raises(McpError): + await server.get_prompt("prompt", {}) + with pytest.raises(McpError): + await server.list_resources() + with pytest.raises(McpError): + await server.list_resource_templates() + with pytest.raises(McpError): + await server.read_resource(AnyUrl("https://example.com/resource")) + finally: + _mcp_proxy_mode.reset(token) + + +@pytest.mark.asyncio +async def test_proxy_list_mode_has_fixed_definitions_without_search_flag() -> None: + from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_proxy_mode + + token = _mcp_proxy_mode.set(True) + try: + with patch( # test-quality-ok: isolate authenticated MCP context seam + "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", + new_callable=AsyncMock, + return_value=(AUTH, None, None, None, None, None, None), + ): + tools = await server.handle_list_tools() + options = server.server.create_initialization_options() + finally: + _mcp_proxy_mode.reset(token) + + assert {tool.name for tool in tools} == { + MCP_PROXY_SEARCH_TOOL_NAME, + MCP_PROXY_SCHEMA_TOOL_NAME, + MCP_PROXY_CALL_TOOL_NAME, + } + assert options.capabilities.prompts is None + assert options.capabilities.resources is None + assert options.capabilities.tools is not None + + +def _server(server_id: str, name: str, **overrides: object) -> MCPServer: + return MCPServer( + server_id=server_id, + name=name, + server_name=name, + url=f"http://{name}.test", + transport=MCPTransport.http, + **overrides, + ) + + +def _upstream_tool(prefix: str, name: str) -> Tool: + return Tool( + name=f"{prefix}-{name}", + description=f"{name} numbers", + inputSchema={"type": "object", "properties": {"a": {"type": "integer"}}, "required": ["a"]}, + ) + + +def _auth(**object_permission: object) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-scope", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="scope", **object_permission), + ) + + +def _ids(result: CallToolResult) -> dict[str, str]: + return {item["name"]: item["tool_id"] for item in _text(result)} + + +class TestMcpProxyAuthorizationScope: + """The real catalog resolver runs (server grants, tool grants, scope header, sentinel); only the + upstream tools/list fetch and the final upstream dispatch are faked.""" + + ALPHA = _server("srv-alpha", "alpha", tool_name_to_display_name={"add": "Add Numbers"}) + BETA = _server("srv-beta", "beta") + ALPHA_ADD = with_mcp_proxy_identity(_upstream_tool("alpha", "add"), "srv-alpha") + ALPHA_MULTIPLY = with_mcp_proxy_identity(_upstream_tool("alpha", "multiply"), "srv-alpha") + BETA_ADD = with_mcp_proxy_identity(_upstream_tool("beta", "add"), "srv-beta") + + @pytest.fixture + def rig(self) -> Iterator[AsyncMock]: + upstream = { + "srv-alpha": [_upstream_tool("alpha", "add"), _upstream_tool("alpha", "multiply")], + "srv-beta": [_upstream_tool("beta", "add")], + } + + async def fetch(server: MCPServer, **_: object) -> list[Tool]: + return list(upstream[server.server_id]) + + dispatched = AsyncMock( + return_value=CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + ) + global_mcp_server_manager.registry.update({"srv-alpha": self.ALPHA, "srv-beta": self.BETA}) + with ( + patch.object( # test-quality-ok: the upstream MCP server is the only faked collaborator + global_mcp_server_manager, "_get_tools_from_server", new=AsyncMock(side_effect=fetch) + ), + patch.object(global_mcp_server_manager, "call_tool", new=dispatched), # test-quality-ok: dispatch seam + ): + yield dispatched + + async def _proxy( + self, name: str, arguments: dict[str, object], auth: UserAPIKeyAuth, **kwargs: object + ) -> CallToolResult: + return await handle_mcp_proxy_tool(name, arguments, auth, **kwargs) + + @pytest.mark.asyncio + async def test_search_and_schema_are_bounded_by_the_key_server_grant(self, rig: AsyncMock) -> None: + granted = _auth(mcp_servers=["srv-alpha"]) + + assert _ids(await self._proxy(MCP_PROXY_SEARCH_TOOL_NAME, {"query": "numbers"}, granted)) == { + "alpha-add": mcp_proxy_tool_id(self.ALPHA_ADD), + "alpha-multiply": mcp_proxy_tool_id(self.ALPHA_MULTIPLY), + } + denied_schema = await self._proxy( + MCP_PROXY_SCHEMA_TOOL_NAME, {"tool_id": mcp_proxy_tool_id(self.BETA_ADD)}, granted + ) + denied_call = await self._proxy( + MCP_PROXY_CALL_TOOL_NAME, {"tool_id": mcp_proxy_tool_id(self.BETA_ADD), "arguments": {"a": 1}}, granted + ) + assert denied_schema.isError is True and denied_call.isError is True + rig.assert_not_awaited() + + @pytest.mark.asyncio + async def test_no_mcp_servers_sentinel_hides_every_tool(self, rig: AsyncMock) -> None: + result = await self._proxy( + MCP_PROXY_SEARCH_TOOL_NAME, {"query": "numbers"}, _auth(mcp_servers=["no-mcp-servers"]) + ) + assert _text(result) == [] + rig.assert_not_awaited() + + @pytest.mark.asyncio + async def test_tool_grant_hides_ungranted_tools_and_blocks_their_ids(self, rig: AsyncMock) -> None: + scoped = _auth(mcp_servers=["srv-alpha"], mcp_tool_permissions={"srv-alpha": ["add"]}) + + assert set(_ids(await self._proxy(MCP_PROXY_SEARCH_TOOL_NAME, {"query": "numbers"}, scoped))) == {"alpha-add"} + blocked = await self._proxy( + MCP_PROXY_CALL_TOOL_NAME, {"tool_id": mcp_proxy_tool_id(self.ALPHA_MULTIPLY), "arguments": {"a": 1}}, scoped + ) + assert blocked.isError is True + rig.assert_not_awaited() + + @pytest.mark.asyncio + async def test_same_named_tools_keep_distinct_ids_and_dispatch_to_their_own_server(self, rig: AsyncMock) -> None: + both = _auth(mcp_servers=["srv-alpha", "srv-beta"]) + + ids = _ids(await self._proxy(MCP_PROXY_SEARCH_TOOL_NAME, {"query": "add"}, both)) + assert set(ids) == {"alpha-add", "beta-add"}, "display-name overrides must not rename proxy identities" + assert ids["alpha-add"] != ids["beta-add"] + + result = await self._proxy(MCP_PROXY_CALL_TOOL_NAME, {"tool_id": ids["beta-add"], "arguments": {"a": 1}}, both) + assert result.isError is False + rig.assert_awaited_once() + assert rig.await_args.kwargs["server_name"] == "beta" + assert rig.await_args.kwargs["name"] == "add" + + @pytest.mark.asyncio + async def test_server_scope_header_narrows_search_within_the_grant(self, rig: AsyncMock) -> None: + both = _auth(mcp_servers=["srv-alpha", "srv-beta"]) + scoped = await self._proxy(MCP_PROXY_SEARCH_TOOL_NAME, {"query": "add"}, both, mcp_servers=["beta"]) + assert set(_ids(scoped)) == {"beta-add"} + rig.assert_not_awaited() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 56a57f4f7ff..5c7accd9120 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8542,6 +8542,50 @@ export interface paths { patch?: never; trace?: never; }; + "/mcp/proxy": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Proxy Mcp Route + * @description Serve the fixed three-tool MCP proxy surface. + */ + get: operations["proxy_mcp_route_mcp_proxy_get"]; + /** + * Proxy Mcp Route + * @description Serve the fixed three-tool MCP proxy surface. + */ + put: operations["proxy_mcp_route_mcp_proxy_put"]; + /** + * Proxy Mcp Route + * @description Serve the fixed three-tool MCP proxy surface. + */ + post: operations["proxy_mcp_route_mcp_proxy_post"]; + /** + * Proxy Mcp Route + * @description Serve the fixed three-tool MCP proxy surface. + */ + delete: operations["proxy_mcp_route_mcp_proxy_delete"]; + /** + * Proxy Mcp Route + * @description Serve the fixed three-tool MCP proxy surface. + */ + options: operations["proxy_mcp_route_mcp_proxy_options"]; + /** + * Proxy Mcp Route + * @description Serve the fixed three-tool MCP proxy surface. + */ + head: operations["proxy_mcp_route_mcp_proxy_head"]; + /** + * Proxy Mcp Route + * @description Serve the fixed three-tool MCP proxy surface. + */ + patch: operations["proxy_mcp_route_mcp_proxy_patch"]; + trace?: never; + }; "/memory-usage-in-mem-cache": { parameters: { query?: never; @@ -50983,6 +51027,146 @@ export interface operations { }; }; }; + proxy_mcp_route_mcp_proxy_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + proxy_mcp_route_mcp_proxy_put: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + proxy_mcp_route_mcp_proxy_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + proxy_mcp_route_mcp_proxy_delete: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + proxy_mcp_route_mcp_proxy_options: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + proxy_mcp_route_mcp_proxy_head: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + proxy_mcp_route_mcp_proxy_patch: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; memory_usage_in_mem_cache_memory_usage_in_mem_cache_get: { parameters: { query?: never;