From 7fe0d9501e3a1a78e4d74825be54589d9921cf68 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 16 Jul 2026 23:19:06 -0700 Subject: [PATCH 1/8] fix(mcp): route tool calls by server_id so unprefixed names cannot pick an arbitrary server tool_name_to_mcp_server_name_mapping mapped a tool name to a server name and was written with an unqualified, global, last-writer-wins key. Two server entries sharing an upstream URL expose the same tool names, so that key always collided and the winner was whichever server registered last, which meant config declaration order decided which upstream credential left the gateway. server_name is not a sound identity: schema.prisma has no unique constraint on server_name or alias, and the create endpoint only checks server_id. The registry is already keyed by server_id and config servers derive a deterministic one, so server_id was always the real identity. The map now holds frozensets of server ids and accumulates owners instead of overwriting them, so a name served by several reachable servers stays visibly ambiguous and is rejected rather than dispatched to an arbitrary server. Ambiguity is judged against the servers the caller can reach, so a session scoped with x-mcp-servers still resolves its unprefixed names. All three writers now go through one registration chokepoint; they previously disagreed on what the value was. Cleanup withdraws only the departing server's id from each row instead of matching by name, which stopped removing one server from deleting a same-named server's routes. execute_mcp_tool now dispatches the server it resolved rather than passing a name that call_tool re-resolved, closing the same drop on the responses API path. --- .../mcp_server/mcp_server_manager.py | 161 ++++++--- .../proxy/_experimental/mcp_server/server.py | 33 +- .../panw_prisma_airs/panw_prisma_airs.py | 25 +- .../mcp/litellm_proxy_mcp_handler.py | 1 + tests/mcp_tests/test_mcp_logging.py | 53 +-- tests/mcp_tests/test_mcp_server.py | 58 ++-- .../mcp_server/test_mcp_server.py | 315 +++++++++++++++++- .../mcp_server/test_mcp_server_manager.py | 144 ++++++-- 8 files changed, 651 insertions(+), 139 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 115ff2e492c4..b8bf986776ef 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -14,6 +14,7 @@ import re import time from contextlib import asynccontextmanager +from dataclasses import dataclass from typing import Any, AsyncIterator, Callable, Literal, Optional, Union, cast from urllib.parse import urlparse @@ -186,6 +187,31 @@ def validate_tool_name(name: str) -> _ToolNameValidationResult: # type: ignore[ ) +@dataclass(frozen=True, slots=True) +class MCPToolRouteResolved: + """Exactly one MCP server serves the tool name.""" + + server: MCPServer + + +@dataclass(frozen=True, slots=True) +class MCPToolRouteNotFound: + """No MCP server serves the tool name.""" + + tool_name: str + + +@dataclass(frozen=True, slots=True) +class MCPToolRouteAmbiguous: + """Several MCP servers serve the tool name, so it cannot address one of them.""" + + tool_name: str + server_ids: frozenset[str] + + +MCPToolRoute = Union[MCPToolRouteResolved, MCPToolRouteNotFound, MCPToolRouteAmbiguous] + + def _blank_to_none(value: str | None) -> str | None: """Collapse an absent, empty, or whitespace-only string to ``None``. @@ -1035,7 +1061,7 @@ def __init__( # semaphore so an edited limit rebuilds it instead of keeping the old cap # until restart. self._server_call_semaphores: dict[str, tuple[int, asyncio.Semaphore]] = {} - self.tool_name_to_mcp_server_name_mapping: dict[str, str] = {} + self.tool_name_to_mcp_server_ids_mapping: dict[str, frozenset[str]] = {} """ { "gmail_send_email": "zapier_mcp_server", @@ -1367,7 +1393,7 @@ async def load_servers_from_config( verbose_logger.debug(f"Loaded MCP Servers: {json.dumps(self.config_mcp_servers, indent=4, default=str)}") - self.initialize_tool_name_to_mcp_server_name_mapping() + self.initialize_tool_name_to_mcp_server_ids_mapping() async def _register_openapi_tools(self, spec_path: str, server: MCPServer, base_url: str): """ @@ -1462,7 +1488,7 @@ async def _register_openapi_tools(self, spec_path: str, server: MCPServer, base_ resolved_operation = resolve_operation_params(operation, path_item, components) # Generate tool name (without prefix initially) - operation_id = operation.get("operationId", f"{method}_{path.replace('/', '_')}") + operation_id = str(operation.get("operationId", f"{method}_{path.replace('/', '_')}")) base_tool_name = operation_id.replace(" ", "_").lower() # Add server prefix to tool name @@ -1490,9 +1516,9 @@ async def _register_openapi_tools(self, spec_path: str, server: MCPServer, base_ handler=tool_func, ) - # Update tool name to server name mapping (for both prefixed and base names) - self.tool_name_to_mcp_server_name_mapping[base_tool_name] = server_prefix - self.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = server_prefix + # Update tool name to server id mapping (for both prefixed and base names) + self._register_tool_route(base_tool_name, server.server_id) + self._register_tool_route(prefixed_tool_name, server.server_id) registered_count += 1 verbose_logger.debug(f"Registered OpenAPI tool: {prefixed_tool_name} for server {server.name}") @@ -1508,9 +1534,12 @@ def _cleanup_server_tool_routing_artifacts(self, server: MCPServer) -> None: When a server leaves ``self.registry`` (eviction, ``remove_server``, etc.), OpenAPI tools remain in ``global_mcp_tool_registry`` and - ``tool_name_to_mcp_server_name_mapping`` unless removed here. Stale - mappings make ``_get_mcp_server_from_tool_name`` resolve to a prefix that + ``tool_name_to_mcp_server_ids_mapping`` unless removed here. Stale + mappings make ``_get_mcp_server_from_tool_name`` resolve to a server that no longer exists in the live registry. + + Only ``server``'s own id is withdrawn from each routing row, so a tool name + shared with another server keeps routing to that other server. """ from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, @@ -1521,24 +1550,11 @@ def _cleanup_server_tool_routing_artifacts(self, server: MCPServer) -> None: openapi_key_prefix = prefix_root + MCP_TOOL_PREFIX_SEPARATOR global_mcp_tool_registry.unregister_tools_with_prefix(openapi_key_prefix) - owned_raw: set[str] = set() - for p in iter_known_server_prefixes(server): - if p: - owned_raw.add(p) - if server.name: - owned_raw.add(server.name) - - owned_normalized = {normalize_server_name(x) for x in owned_raw} - - stale_mapping_keys: list[str] = [] - for tool_name, mapped_server in list(self.tool_name_to_mcp_server_name_mapping.items()): - if mapped_server in owned_raw: - stale_mapping_keys.append(tool_name) - elif normalize_server_name(str(mapped_server)) in owned_normalized: - stale_mapping_keys.append(tool_name) - - for key in stale_mapping_keys: - del self.tool_name_to_mcp_server_name_mapping[key] + self.tool_name_to_mcp_server_ids_mapping = { + tool_name: remaining + for tool_name, owner_ids in self.tool_name_to_mcp_server_ids_mapping.items() + if (remaining := owner_ids - {server.server_id}) + } def remove_server(self, mcp_server: LiteLLM_MCPServerTable): """ @@ -1946,7 +1962,7 @@ async def _maybe_register_openapi_tools(self, server: MCPServer, *, initialize_m base_url=server.url or "", ) if initialize_mapping: - self.initialize_tool_name_to_mcp_server_name_mapping() + self.initialize_tool_name_to_mcp_server_ids_mapping() async def add_server(self, mcp_server: LiteLLM_MCPServerTable): # The runtime registry is the allowlist for tool calls and health @@ -3818,10 +3834,10 @@ def _create_prefixed_tools(self, tools: list[MCPTool], server: MCPServer, add_pr # Register every known prefix form (alias, server_name, server_id, # short ID) so call_tool can resolve regardless of which form a # caller / cached client is using. - self.tool_name_to_mcp_server_name_mapping[original_name] = prefix + self._register_tool_route(original_name, server.server_id) for known_prefix in iter_known_server_prefixes(server): qualified = add_server_prefix_to_name(original_name, known_prefix) - self.tool_name_to_mcp_server_name_mapping[qualified] = prefix + self._register_tool_route(qualified, server.server_id) verbose_logger.info(f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}") return prefixed_tools @@ -4546,8 +4562,8 @@ def _candidate_matches_server_name(candidate: MCPServer) -> bool: if resolved_by_server_name_only: tool_known = ( - name in self.tool_name_to_mcp_server_name_mapping - or prefixed_tool_name in self.tool_name_to_mcp_server_name_mapping + name in self.tool_name_to_mcp_server_ids_mapping + or prefixed_tool_name in self.tool_name_to_mcp_server_ids_mapping ) if not tool_known: raise ValueError(f"Tool {name} not found") @@ -4658,6 +4674,7 @@ async def call_tool( oauth2_headers: Optional[dict[str, str]] = None, raw_headers: Optional[dict[str, str]] = None, host_progress_callback: Optional[Callable] = None, + resolved_server: MCPServer | None = None, ) -> CallToolResult: """ Call a tool with the given name and arguments @@ -4670,13 +4687,21 @@ async def call_tool( mcp_auth_header: MCP auth header (deprecated) mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} proxy_logging_obj: Optional ProxyLogging object for hook integration + resolved_server: Server the caller already resolved. When supplied it is + dispatched to verbatim, so identity never round-trips through the + non-unique ``server_name``. Callers without one fall back to + resolving by name. Returns: CallToolResult from the MCP server """ start_time = datetime.datetime.now() - mcp_server = self._resolve_mcp_server_for_tool_call(server_name, name) + mcp_server = ( + resolved_server + if resolved_server is not None + else self._resolve_mcp_server_for_tool_call(server_name, name) + ) # Resolved before any hook runs so a missing BYOK credential (401) never # leaves during-hook side effects (audit logging, rate-limit bookkeeping) @@ -4776,19 +4801,19 @@ async def _call_openapi_via_handler(): # End of Methods that call the upstream MCP servers ######################################################### - def initialize_tool_name_to_mcp_server_name_mapping(self): + def initialize_tool_name_to_mcp_server_ids_mapping(self): """ On startup, initialize the tool name to MCP server name mapping """ try: if asyncio.get_running_loop(): - asyncio.create_task(self._initialize_tool_name_to_mcp_server_name_mapping()) + asyncio.create_task(self._initialize_tool_name_to_mcp_server_ids_mapping()) except RuntimeError as e: # no running event loop verbose_logger.exception( f"No running event loop - skipping tool name to MCP server name mapping initialization: {str(e)}" ) - async def _initialize_tool_name_to_mcp_server_name_mapping(self): + async def _initialize_tool_name_to_mcp_server_ids_mapping(self): """ Call list_tools for each server and update the tool name to MCP server name mapping Note: This now handles prefixed tool names @@ -4816,8 +4841,47 @@ async def _initialize_tool_name_to_mcp_server_name_mapping(self): # The tool.name here is already prefixed from _get_tools_from_server # Extract original name for mapping original_name, _ = split_server_prefix_from_name(tool.name) - self.tool_name_to_mcp_server_name_mapping[original_name] = server.name - self.tool_name_to_mcp_server_name_mapping[tool.name] = server.name + self._register_tool_route(original_name, server.server_id) + self._register_tool_route(tool.name, server.server_id) + + def _register_tool_route(self, tool_name: str, server_id: str) -> None: + """Record that ``server_id`` serves ``tool_name``, preserving any other owners. + + Routing rows accumulate rather than overwrite so that a tool name served by + more than one server stays resolvable as ambiguous instead of silently + collapsing to whichever server registered last. + """ + owners = self.tool_name_to_mcp_server_ids_mapping.get(tool_name, frozenset()) + self.tool_name_to_mcp_server_ids_mapping[tool_name] = owners | {server_id} + + def resolve_tool_route( + self, + tool_name: str, + allowed_server_ids: frozenset[str] | None = None, + ) -> "MCPToolRoute": + """Resolve ``tool_name`` to the single MCP server that serves it. + + Ambiguity is judged against ``allowed_server_ids`` when given, because an + unprefixed name is only ambiguous relative to the servers the caller can + actually reach; a session scoped to one server (``x-mcp-servers``) is served + unprefixed names precisely because nothing else is in scope for it. + + Returns ``MCPToolRouteAmbiguous`` when several reachable servers expose the + name, so callers fail loudly instead of dispatching to an arbitrary one. + """ + owner_ids = self.tool_name_to_mcp_server_ids_mapping.get(tool_name) + if owner_ids: + in_scope = owner_ids if allowed_server_ids is None else owner_ids & allowed_server_ids + if len(in_scope) > 1: + return MCPToolRouteAmbiguous(tool_name=tool_name, server_ids=in_scope) + if len(in_scope) == 1: + scoped_server = self.get_mcp_server_by_id(next(iter(in_scope))) + if scoped_server is not None: + return MCPToolRouteResolved(server=scoped_server) + server = self._get_mcp_server_from_tool_name(tool_name) + if server is None: + return MCPToolRouteNotFound(tool_name=tool_name) + return MCPToolRouteResolved(server=server) def _get_mcp_server_from_tool_name(self, tool_name: str) -> Optional[MCPServer]: """ @@ -4827,7 +4891,9 @@ def _get_mcp_server_from_tool_name(self, tool_name: str) -> Optional[MCPServer]: tool_name: Tool name (can be prefixed or non-prefixed) Returns: - MCPServer if found, None otherwise + MCPServer if found, None otherwise. None is also returned when several + servers expose ``tool_name``; use ``resolve_tool_route`` to tell the two + cases apart. """ registry_servers = list(self.get_registry().values()) @@ -4841,14 +4907,11 @@ def _get_mcp_server_from_tool_name(self, tool_name: str) -> Optional[MCPServer]: prefix_to_server.setdefault(normalised, server) # First try with the original tool name - if tool_name in self.tool_name_to_mcp_server_name_mapping: - server_name = self.tool_name_to_mcp_server_name_mapping[tool_name] - normalised_lookup = normalize_server_name(server_name) - if normalised_lookup in prefix_to_server: - return prefix_to_server[normalised_lookup] - for server in registry_servers: - if normalize_server_name(server.name) == normalised_lookup: - return server + owner_ids = self.tool_name_to_mcp_server_ids_mapping.get(tool_name) + if owner_ids is not None and len(owner_ids) == 1: + owned = self.get_mcp_server_by_id(next(iter(owner_ids))) + if owned is not None: + return owned # If not found and tool name is prefixed, extract the prefix and # match against any known form. @@ -4860,8 +4923,8 @@ def _get_mcp_server_from_tool_name(self, tool_name: str) -> Optional[MCPServer]: normalised_prefix = normalize_server_name(server_name_from_prefix) matched_server = prefix_to_server.get(normalised_prefix) if matched_server is not None and ( - original_tool_name in self.tool_name_to_mcp_server_name_mapping - or tool_name in self.tool_name_to_mcp_server_name_mapping + original_tool_name in self.tool_name_to_mcp_server_ids_mapping + or tool_name in self.tool_name_to_mcp_server_ids_mapping ): return matched_server @@ -4964,7 +5027,7 @@ async def reload_servers_from_database(self): self.registry = registered_registry if registered_openapi_tools: - self.initialize_tool_name_to_mcp_server_name_mapping() + self.initialize_tool_name_to_mcp_server_ids_mapping() verbose_logger.debug("MCP registry refreshed (%s servers in registry)", len(registered_registry)) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 68a61b851758..0705c392f7fb 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -358,6 +358,8 @@ def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, + MCPToolRouteAmbiguous, + MCPToolRouteResolved, _caller_authorization_fans_out, _client_forwarded_authorization_headers, _should_strip_caller_authorization, @@ -2622,7 +2624,33 @@ async def execute_mcp_tool( 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 requested_server is None: + route = global_mcp_server_manager.resolve_tool_route( + name, + allowed_server_ids=frozenset(server.server_id for server in allowed_mcp_servers), + ) + if isinstance(route, MCPToolRouteAmbiguous): + candidates = sorted( + server.name + for server in ( + global_mcp_server_manager.get_mcp_server_by_id(server_id) for server_id in route.server_ids + ) + if server is not None + ) + raise HTTPException( + status_code=409, + detail={ + "error": "ambiguous_tool_name", + "message": ( + f"Tool '{name}' is served by more than one MCP server " + f"({', '.join(candidates)}). Call it by its server-prefixed name to " + f"select one." + ), + }, + ) + mcp_server = route.server if isinstance(route, MCPToolRouteResolved) else None + else: + 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( @@ -2820,6 +2848,7 @@ async def execute_mcp_tool( raw_headers=raw_headers, litellm_logging_obj=litellm_logging_obj, host_progress_callback=host_progress_callback, + resolved_server=mcp_server, ) # Fall back to local tool registry with original name (legacy support) @@ -3132,6 +3161,7 @@ async def _handle_managed_mcp_tool( raw_headers: Optional[Dict[str, str]] = None, litellm_logging_obj: Optional[Any] = None, host_progress_callback: Optional[Callable] = None, + resolved_server: MCPServer | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" # Import here to avoid circular import @@ -3148,6 +3178,7 @@ async def _handle_managed_mcp_tool( raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, host_progress_callback=host_progress_callback, + resolved_server=resolved_server, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) return call_tool_result diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 192c8c9bc771..e3317f23e1e9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -467,17 +467,20 @@ def _get_mcp_server_name(request_data: dict, mcp_tool_name: str) -> str: ) server_id = request_data.get("server_id") - if server_id: - server = global_mcp_server_manager.get_mcp_server_by_id(server_id) - if server: - return ( - getattr(server, "alias", None) - or getattr(server, "server_name", None) - or getattr(server, "name", None) - or getattr(server, "server_id", None) - or "unknown" - ) - return global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.get(mcp_tool_name, "unknown") + server = ( + global_mcp_server_manager.get_mcp_server_by_id(server_id) + if server_id + else global_mcp_server_manager._get_mcp_server_from_tool_name(mcp_tool_name) + ) + if server: + return ( + getattr(server, "alias", None) + or getattr(server, "server_name", None) + or getattr(server, "name", None) + or getattr(server, "server_id", None) + or "unknown" + ) + return "unknown" except ImportError: return "unknown" except Exception: diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index e03f0296109f..460be2b04274 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -779,6 +779,7 @@ async def _execute_tool_calls( oauth2_headers=oauth2_headers, raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, + resolved_server=mcp_server, ) if litellm_logging_obj: diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index 55b49aa0d295..d8193c3074f6 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -104,7 +104,7 @@ def mock_client_constructor(*args, **kwargs): litellm.callbacks = [test_logger] # Initialize the tool mapping - await local_mcp_server_manager._initialize_tool_name_to_mcp_server_name_mapping() + await local_mcp_server_manager._initialize_tool_name_to_mcp_server_ids_mapping() # Patch the global manager in both modules where it's used with ( @@ -121,17 +121,18 @@ def mock_client_constructor(*args, **kwargs): _set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids()) print( - "tool_name_to_mcp_server_name_mapping", - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping, + "tool_name_to_mcp_server_ids_mapping", + local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping, ) # Manually add the tool mapping to ensure it's available (since mocking might not capture it properly) - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ + zapier_server_ids = frozenset(local_mcp_server_manager.get_all_mcp_server_ids()) + local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[ "add_tools" - ] = "zapier_gmail_server" - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ + ] = zapier_server_ids + local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[ "zapier_gmail_server-add_tools" - ] = "zapier_gmail_server" + ] = zapier_server_ids # Call mcp tool response = await mcp_server_tool_call( @@ -237,21 +238,22 @@ def mock_client_constructor(*args, **kwargs): litellm.callbacks = [test_logger] # Initialize the tool mapping - await local_mcp_server_manager._initialize_tool_name_to_mcp_server_name_mapping() + await local_mcp_server_manager._initialize_tool_name_to_mcp_server_ids_mapping() # Manually add the tool mapping to ensure it's available (since mocking might not capture it properly) - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ + test_server_ids = frozenset(local_mcp_server_manager.get_all_mcp_server_ids()) + local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[ "expensive_tool" - ] = "test_server" - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ + ] = test_server_ids + local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[ "test_server-expensive_tool" - ] = "test_server" - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["cheap_tool"] = ( - "test_server" + ] = test_server_ids + local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping["cheap_tool"] = ( + test_server_ids ) - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ + local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[ "test_server-cheap_tool" - ] = "test_server" + ] = test_server_ids # Patch the global manager in both modules where it's used with ( @@ -268,8 +270,8 @@ def mock_client_constructor(*args, **kwargs): _set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids()) print( - "tool_name_to_mcp_server_name_mapping", - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping, + "tool_name_to_mcp_server_ids_mapping", + local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping, ) # Test 1: Call expensive_tool - should cost 5.0 @@ -401,15 +403,16 @@ def mock_client_constructor(*args, **kwargs): litellm.callbacks = [test_logger] # Initialize the tool mapping - await local_mcp_server_manager._initialize_tool_name_to_mcp_server_name_mapping() + await local_mcp_server_manager._initialize_tool_name_to_mcp_server_ids_mapping() # Manually add the tool mapping to ensure it's available (since mocking might not capture it properly) - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["add_tools"] = ( - "zapier_gmail_server" + zapier_server_ids = frozenset(local_mcp_server_manager.get_all_mcp_server_ids()) + local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping["add_tools"] = ( + zapier_server_ids ) - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[ + local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[ "zapier_gmail_server-add_tools" - ] = "zapier_gmail_server" + ] = zapier_server_ids # Patch the global manager in both modules where it's used with ( @@ -426,8 +429,8 @@ def mock_client_constructor(*args, **kwargs): _set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids()) print( - "tool_name_to_mcp_server_name_mapping", - local_mcp_server_manager.tool_name_to_mcp_server_name_mapping, + "tool_name_to_mcp_server_ids_mapping", + local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping, ) # Call mcp tool using the correct separator format (- not /) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 515bf1233aa9..e8307f311814 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -105,12 +105,18 @@ def mock_client_constructor(*args, **kwargs): assert tools[0].name == f"{expected_prefix}-gmail_send_email" # Manually set up the tool mapping for the call_tool test - mcp_server_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = ( - expected_prefix + zapier_server_ids = frozenset( + server_id + for server_id, server in mcp_server_manager.get_registry().items() + if server.server_name == "zapier_mcp_server" ) - mcp_server_manager.tool_name_to_mcp_server_name_mapping[ + assert len(zapier_server_ids) == 1, "Expected exactly one configured zapier server" + mcp_server_manager.tool_name_to_mcp_server_ids_mapping["gmail_send_email"] = ( + zapier_server_ids + ) + mcp_server_manager.tool_name_to_mcp_server_ids_mapping[ f"{expected_prefix}-gmail_send_email" - ] = expected_prefix + ] = zapier_server_ids result = await mcp_server_manager.call_tool( server_name="zapier_mcp_server", @@ -221,16 +227,16 @@ def mock_client_constructor(*args, **kwargs): # Verify tool mapping was updated expected_prefix = "test_http_server" assert ( - test_manager.tool_name_to_mcp_server_name_mapping[ + test_manager.tool_name_to_mcp_server_ids_mapping[ f"{expected_prefix}-gmail_send_email" ] - == expected_prefix + == frozenset(allowed_server_ids) ) assert ( - test_manager.tool_name_to_mcp_server_name_mapping[ + test_manager.tool_name_to_mcp_server_ids_mapping[ f"{expected_prefix}-calendar_create_event" ] - == expected_prefix + == frozenset(allowed_server_ids) ) @@ -275,8 +281,8 @@ def mock_client_constructor(*args, **kwargs): ) # Manually set up tool mapping (normally done by list_tools) - test_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = ( - "test_http_server" + test_manager.tool_name_to_mcp_server_ids_mapping["gmail_send_email"] = frozenset( + test_manager.get_registry().keys() ) # Call the tool @@ -341,8 +347,8 @@ def mock_client_constructor(*args, **kwargs): ) # Manually set up tool mapping - test_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = ( - "test_http_server" + test_manager.tool_name_to_mcp_server_ids_mapping["gmail_send_email"] = frozenset( + test_manager.get_registry().keys() ) # Call the tool with invalid data @@ -383,8 +389,8 @@ async def test_mcp_http_transport_tool_not_found(): ) # Mapping populated for this server but not for the requested tool - test_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = ( - "test_http_server" + test_manager.tool_name_to_mcp_server_ids_mapping["gmail_send_email"] = frozenset( + test_manager.get_registry().keys() ) # Try to call a tool that doesn't exist in mapping @@ -1395,12 +1401,12 @@ def mock_client_constructor(*args, **kwargs): # Verify mapping is updated correctly assert ( - test_manager.tool_name_to_mcp_server_name_mapping["send_email"] - == "my_alias" + test_manager.tool_name_to_mcp_server_ids_mapping["send_email"] + == frozenset({"test-server-123"}) ) assert ( - test_manager.tool_name_to_mcp_server_name_mapping["my_alias-send_email"] - == "my_alias" + test_manager.tool_name_to_mcp_server_ids_mapping["my_alias-send_email"] + == frozenset({"test-server-123"}) ) @@ -1455,12 +1461,12 @@ def mock_client_constructor(*args, **kwargs): # Verify mapping is updated correctly assert ( - test_manager.tool_name_to_mcp_server_name_mapping["send_email"] - == "Test Server" + test_manager.tool_name_to_mcp_server_ids_mapping["send_email"] + == frozenset({"test-server-123"}) ) assert ( - test_manager.tool_name_to_mcp_server_name_mapping["Test_Server-send_email"] - == "Test Server" + test_manager.tool_name_to_mcp_server_ids_mapping["Test_Server-send_email"] + == frozenset({"test-server-123"}) ) @@ -1515,14 +1521,14 @@ def mock_client_constructor(*args, **kwargs): # Verify mapping is updated correctly assert ( - test_manager.tool_name_to_mcp_server_name_mapping["send_email"] - == "test-server-123" + test_manager.tool_name_to_mcp_server_ids_mapping["send_email"] + == frozenset({"test-server-123"}) ) assert ( - test_manager.tool_name_to_mcp_server_name_mapping[ + test_manager.tool_name_to_mcp_server_ids_mapping[ "test-server-123-send_email" ] - == "test-server-123" + == frozenset({"test-server-123"}) ) 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 a983ac3ff480..6b878b11a749 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 @@ -39,11 +39,11 @@ def cleanup_mcp_global_state(): # Clear before test global_mcp_server_manager.registry.clear() - global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.clear() + global_mcp_server_manager.tool_name_to_mcp_server_ids_mapping.clear() yield # Clear after test global_mcp_server_manager.registry.clear() - global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.clear() + global_mcp_server_manager.tool_name_to_mcp_server_ids_mapping.clear() except ImportError: # MCP not available, skip cleanup yield @@ -5957,8 +5957,8 @@ async def fake_handle_managed_mcp_tool(**kwargs): with ( patch.dict( - mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_name_mapping, - {"echo": oauth_server.name}, + mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_ids_mapping, + {"echo": frozenset({oauth_server.server_id})}, ), patch.object( mcp_module.global_mcp_server_manager, @@ -6044,8 +6044,8 @@ async def fake_create_mcp_client(server, **kwargs): with ( patch.dict( - mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_name_mapping, - {"echo": collision_server.name}, + mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_ids_mapping, + {"echo": frozenset({collision_server.server_id})}, ), patch.object( mcp_module.global_mcp_server_manager, @@ -6087,6 +6087,309 @@ async def fake_create_mcp_client(server, **kwargs): assert routed.authentication_token != collision_server.authentication_token +@pytest.mark.asyncio +async def test_execute_mcp_tool_jsonrpc_unprefixed_ambiguous_tool_is_rejected(): + """MCP JSON-RPC carries no server_id, so an unprefixed name owned by two servers must 409. + + The REST path disambiguates with requested_server_id, but the MCP protocol has no such + field. Without this the call silently dispatched to whichever server registered the tool + name last, sending that server's upstream credentials. + """ + from litellm.proxy._experimental.mcp_server import server as mcp_module + + alpha = MCPServer( + server_id="alpha-server-id", + name="echo_alpha", + server_name="echo_alpha", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + authentication_token="alpha-secret", + ) + zulu = MCPServer( + server_id="zulu-server-id", + name="echo_zulu", + server_name="echo_zulu", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + authentication_token="zulu-secret", + ) + + injected: dict = {} + + async def fake_create_mcp_client(server, **kwargs): + injected["server"] = server + raise AssertionError("upstream must not be called for an ambiguous tool name") + + with ( + patch.dict( + mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_ids_mapping, + {"echo": frozenset({alpha.server_id, zulu.server_id})}, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "get_registry", + return_value={alpha.server_id: alpha, zulu.server_id: zulu}, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "_create_mcp_client", + new=fake_create_mcp_client, + ), + patch.object(mcp_module.MCPRequestHandler, "is_tool_allowed", return_value=True), + patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=None), + patch("litellm.proxy.proxy_server.proxy_logging_obj", None), + pytest.raises(HTTPException) as exc_info, + ): + await mcp_module.execute_mcp_tool( + name="echo", + arguments={"message": "hello"}, + allowed_mcp_servers=[alpha, zulu], + start_time=datetime.now(), + ) + + assert exc_info.value.status_code == 409 + assert exc_info.value.detail["error"] == "ambiguous_tool_name" + assert "echo_alpha" in exc_info.value.detail["message"] + assert "echo_zulu" in exc_info.value.detail["message"] + assert "server" not in injected + + +@pytest.mark.asyncio +async def test_execute_mcp_tool_jsonrpc_unprefixed_resolves_when_caller_reaches_one_server(): + """A caller scoped to one server is served unprefixed names and must still route. + + Ambiguity is relative to what the caller can reach. Judging it against the whole + registry would reject an x-mcp-servers scoped session, whose unprefixed tool names + only ever had one candidate. + """ + from mcp.types import TextContent + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + alpha = MCPServer( + server_id="alpha-server-id", + name="echo_alpha", + server_name="echo_alpha", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + authentication_token="alpha-secret", + ) + zulu = MCPServer( + server_id="zulu-server-id", + name="echo_zulu", + server_name="echo_zulu", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + authentication_token="zulu-secret", + ) + + fake_client = MagicMock() + fake_client._last_initialize_instructions = None + fake_client.call_tool = AsyncMock( + return_value=mcp_module.CallToolResult( + content=[TextContent(type="text", text="ok")], + isError=False, + ) + ) + + injected: dict = {} + + async def fake_create_mcp_client(server, **kwargs): + injected["server"] = server + return fake_client + + with ( + patch.dict( + mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_ids_mapping, + {"echo": frozenset({alpha.server_id, zulu.server_id})}, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "get_registry", + return_value={alpha.server_id: alpha, zulu.server_id: zulu}, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "_create_mcp_client", + new=fake_create_mcp_client, + ), + patch.object(mcp_module.MCPRequestHandler, "is_tool_allowed", return_value=True), + patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=None), + patch("litellm.proxy.proxy_server.proxy_logging_obj", None), + ): + await mcp_module.execute_mcp_tool( + name="echo", + arguments={"message": "hello"}, + allowed_mcp_servers=[zulu], + start_time=datetime.now(), + ) + + routed = injected["server"] + assert routed.server_id == zulu.server_id + assert routed.authentication_token == "zulu-secret" + + +@pytest.mark.asyncio +async def test_execute_mcp_tool_dispatches_resolved_server_not_a_name_lookup(): + """The server execute_mcp_tool resolved must be the server dispatched to. + + server_name is not unique, so re-deriving the server from it inside call_tool picks the + first registry entry with that name, which defeats server_id resolution when two servers + share a name. + """ + from mcp.types import TextContent + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + first_named_shared = MCPServer( + server_id="first-id", + name="shared", + server_name="shared", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + authentication_token="first-secret", + ) + second_named_shared = MCPServer( + server_id="second-id", + name="shared", + server_name="shared", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + authentication_token="second-secret", + ) + + fake_client = MagicMock() + fake_client._last_initialize_instructions = None + fake_client.call_tool = AsyncMock( + return_value=mcp_module.CallToolResult( + content=[TextContent(type="text", text="ok")], + isError=False, + ) + ) + + injected: dict = {} + + async def fake_create_mcp_client(server, **kwargs): + injected["server"] = server + return fake_client + + with ( + patch.dict( + mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_ids_mapping, + {"echo": frozenset({second_named_shared.server_id})}, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "get_registry", + return_value={ + first_named_shared.server_id: first_named_shared, + second_named_shared.server_id: second_named_shared, + }, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "_create_mcp_client", + new=fake_create_mcp_client, + ), + patch.object(mcp_module.MCPRequestHandler, "is_tool_allowed", return_value=True), + patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=None), + patch("litellm.proxy.proxy_server.proxy_logging_obj", None), + ): + await mcp_module.execute_mcp_tool( + name="echo", + arguments={"message": "hello"}, + allowed_mcp_servers=[first_named_shared, second_named_shared], + start_time=datetime.now(), + requested_server_id=second_named_shared.server_id, + ) + + routed = injected["server"] + assert routed.server_id == second_named_shared.server_id + assert routed.authentication_token == "second-secret" + + +@pytest.mark.asyncio +async def test_execute_mcp_tool_jsonrpc_prefixed_tool_routes_to_its_own_server(): + """The prefixed name stays unambiguous and must still reach exactly its own server.""" + from mcp.types import TextContent + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + alpha = MCPServer( + server_id="alpha-server-id", + name="echo_alpha", + server_name="echo_alpha", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + authentication_token="alpha-secret", + ) + zulu = MCPServer( + server_id="zulu-server-id", + name="echo_zulu", + server_name="echo_zulu", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + authentication_token="zulu-secret", + ) + + fake_client = MagicMock() + fake_client._last_initialize_instructions = None + fake_client.call_tool = AsyncMock( + return_value=mcp_module.CallToolResult( + content=[TextContent(type="text", text="ok")], + isError=False, + ) + ) + + injected: dict = {} + + async def fake_create_mcp_client(server, **kwargs): + injected["server"] = server + return fake_client + + with ( + patch.dict( + mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_ids_mapping, + { + "echo": frozenset({alpha.server_id, zulu.server_id}), + "echo_alpha-echo": frozenset({alpha.server_id}), + "echo_zulu-echo": frozenset({zulu.server_id}), + }, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "get_registry", + return_value={alpha.server_id: alpha, zulu.server_id: zulu}, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "_create_mcp_client", + new=fake_create_mcp_client, + ), + patch.object(mcp_module.MCPRequestHandler, "is_tool_allowed", return_value=True), + patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=None), + patch("litellm.proxy.proxy_server.proxy_logging_obj", None), + ): + await mcp_module.execute_mcp_tool( + name="echo_zulu-echo", + arguments={"message": "hello"}, + allowed_mcp_servers=[alpha, zulu], + start_time=datetime.now(), + ) + + routed = injected["server"] + assert routed.server_id == zulu.server_id + assert routed.authentication_token == "zulu-secret" + + @pytest.mark.asyncio async def test_execute_mcp_tool_rest_prefixed_tool_still_validates_server_id(): """Prefixed REST tool names must still match the requested server_id.""" 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 80bf08a5ebae..646c7a835daf 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 @@ -3005,7 +3005,7 @@ async def test_load_servers_from_config_reads_all_token_exchange_fields(self): assert minimal_spec.config.profile == "rfc8693" @pytest.mark.asyncio - async def test_config_oauth_initialize_tool_name_to_mcp_server_name_mapping(self): + async def test_config_oauth_initialize_tool_name_to_mcp_server_ids_mapping(self): manager = MCPServerManager() config = { @@ -3022,8 +3022,8 @@ async def test_config_oauth_initialize_tool_name_to_mcp_server_name_mapping(self await manager.load_servers_from_config(config) # Initialize the tool mapping - await manager._initialize_tool_name_to_mcp_server_name_mapping() - assert manager.tool_name_to_mcp_server_name_mapping == {} + await manager._initialize_tool_name_to_mcp_server_ids_mapping() + assert manager.tool_name_to_mcp_server_ids_mapping == {} @pytest.mark.asyncio async def test_list_tools_handles_missing_server_alias(self): @@ -3869,8 +3869,8 @@ def test_resolve_mcp_server_for_tool_call_via_prefixed_name(self): transport=MCPTransport.http, ) manager.registry = {"jira": server} - manager.tool_name_to_mcp_server_name_mapping["jira-search_issues"] = "jira" - manager.tool_name_to_mcp_server_name_mapping["search_issues"] = "jira" + manager.tool_name_to_mcp_server_ids_mapping["jira-search_issues"] = frozenset({"jira"}) + manager.tool_name_to_mcp_server_ids_mapping["search_issues"] = frozenset({"jira"}) resolved = manager._resolve_mcp_server_for_tool_call("jira", "search_issues") assert resolved is server @@ -3885,11 +3885,113 @@ def test_resolve_mcp_server_for_tool_call_via_alias(self): transport=MCPTransport.http, ) manager.registry = {"srv-uuid-123": server} - manager.tool_name_to_mcp_server_name_mapping["create_zap"] = "zapier" + manager.tool_name_to_mcp_server_ids_mapping["create_zap"] = frozenset({"srv-uuid-123"}) resolved = manager._resolve_mcp_server_for_tool_call("zapier-alias", "create_zap") assert resolved is server + def test_register_tool_route_accumulates_owners_across_servers(self): + """Two servers exposing one tool name are both recorded, not last-writer-wins.""" + manager = MCPServerManager() + manager._register_tool_route("echo", "id-alpha") + manager._register_tool_route("echo", "id-zulu") + + assert manager.tool_name_to_mcp_server_ids_mapping["echo"] == frozenset({"id-alpha", "id-zulu"}) + + def test_register_tool_route_is_idempotent_for_one_server(self): + """Re-listing the same server must not make its own tool look ambiguous.""" + manager = MCPServerManager() + manager._register_tool_route("echo", "id-alpha") + manager._register_tool_route("echo", "id-alpha") + + assert manager.tool_name_to_mcp_server_ids_mapping["echo"] == frozenset({"id-alpha"}) + + def test_get_mcp_server_from_tool_name_refuses_ambiguous_unprefixed_name(self): + """An unprefixed name owned by several servers resolves to nothing, never to one of them.""" + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) + zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha, "id-zulu": zulu} + manager._register_tool_route("echo", "id-alpha") + manager._register_tool_route("echo", "id-zulu") + + assert manager._get_mcp_server_from_tool_name("echo") is None + + def test_get_mcp_server_from_tool_name_resolves_sole_owner_by_id(self): + """A tool name owned by exactly one server still resolves, addressed by server_id.""" + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha} + manager._register_tool_route("echo", "id-alpha") + + assert manager._get_mcp_server_from_tool_name("echo") is alpha + + def test_get_mcp_server_from_tool_name_refuses_prefixed_name_of_duplicate_named_servers(self): + """server_name is not unique, so a prefix shared by two servers must not silently pick one.""" + manager = MCPServerManager() + first = MCPServer(server_id="id-first", name="shared", transport=MCPTransport.http) + second = MCPServer(server_id="id-second", name="shared", transport=MCPTransport.http) + manager.registry = {"id-first": first, "id-second": second} + manager._register_tool_route("shared-echo", "id-first") + manager._register_tool_route("shared-echo", "id-second") + + assert manager._get_mcp_server_from_tool_name("shared-echo") is None + + def test_resolve_tool_route_names_every_ambiguous_owner(self): + """The ambiguous route carries all owners so callers can report the real candidates.""" + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) + zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha, "id-zulu": zulu} + manager._register_tool_route("echo", "id-alpha") + manager._register_tool_route("echo", "id-zulu") + + route = manager.resolve_tool_route("echo") + + assert type(route).__name__ == "MCPToolRouteAmbiguous" + assert route.server_ids == frozenset({"id-alpha", "id-zulu"}) + + def test_resolve_tool_route_resolves_sole_owner(self): + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha} + manager._register_tool_route("echo", "id-alpha") + + route = manager.resolve_tool_route("echo") + + assert type(route).__name__ == "MCPToolRouteResolved" + assert route.server is alpha + + def test_resolve_tool_route_reports_unknown_tool(self): + manager = MCPServerManager() + + assert type(manager.resolve_tool_route("nothing_serves_this")).__name__ == "MCPToolRouteNotFound" + + def test_cleanup_withdraws_only_departing_server_from_shared_tool_name(self): + """Removing one server must leave a co-owned tool name routable to the survivor.""" + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) + zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha, "id-zulu": zulu} + manager._register_tool_route("echo", "id-alpha") + manager._register_tool_route("echo", "id-zulu") + + manager._cleanup_server_tool_routing_artifacts(zulu) + del manager.registry["id-zulu"] + + assert manager.tool_name_to_mcp_server_ids_mapping["echo"] == frozenset({"id-alpha"}) + assert manager._get_mcp_server_from_tool_name("echo") is alpha + + def test_cleanup_drops_the_route_when_its_last_owner_leaves(self): + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha} + manager._register_tool_route("echo", "id-alpha") + + manager._cleanup_server_tool_routing_artifacts(alpha) + + assert "echo" not in manager.tool_name_to_mcp_server_ids_mapping + def test_resolve_mcp_server_for_tool_call_unknown_tool_with_empty_mapping(self): """Server-name match alone must not let unknown tools through when the mapping has no entries for that server (e.g. listing has not completed @@ -3916,7 +4018,7 @@ def test_resolve_mcp_server_for_tool_call_fallback_to_unprefixed_lookup(self): transport=MCPTransport.http, ) manager.registry = {"linear": server} - manager.tool_name_to_mcp_server_name_mapping["create_issue"] = "linear" + manager.tool_name_to_mcp_server_ids_mapping["create_issue"] = frozenset({"linear"}) # server_name is empty so the fallback unprefixed lookup runs and matches. resolved = manager._resolve_mcp_server_for_tool_call("", "create_issue") @@ -3943,8 +4045,8 @@ def test_resolve_mcp_server_for_tool_call_unknown_tool_with_known_server(self): ) manager.registry = {"github": server} # Mapping has *some* tools for github but not "missing_tool". - manager.tool_name_to_mcp_server_name_mapping["github-list_repos"] = "github" - manager.tool_name_to_mcp_server_name_mapping["list_repos"] = "github" + manager.tool_name_to_mcp_server_ids_mapping["github-list_repos"] = frozenset({"github"}) + manager.tool_name_to_mcp_server_ids_mapping["list_repos"] = frozenset({"github"}) with pytest.raises(ValueError, match="Tool missing_tool not found"): manager._resolve_mcp_server_for_tool_call("github", "missing_tool") @@ -4254,10 +4356,10 @@ def test_create_prefixed_tools_updates_mapping_for_both_forms(self): assert names == ["close_issue", "create_issue"] # Mapping should include both original and prefixed names -> resolves calls either way - assert manager.tool_name_to_mcp_server_name_mapping["create_issue"] == "jira" - assert manager.tool_name_to_mcp_server_name_mapping["jira-create_issue"] == "jira" - assert manager.tool_name_to_mcp_server_name_mapping["close_issue"] == "jira" - assert manager.tool_name_to_mcp_server_name_mapping["jira-close_issue"] == "jira" + assert manager.tool_name_to_mcp_server_ids_mapping["create_issue"] == frozenset({"jira"}) + assert manager.tool_name_to_mcp_server_ids_mapping["jira-create_issue"] == frozenset({"jira"}) + assert manager.tool_name_to_mcp_server_ids_mapping["close_issue"] == frozenset({"jira"}) + assert manager.tool_name_to_mcp_server_ids_mapping["jira-close_issue"] == frozenset({"jira"}) def test_get_mcp_server_from_tool_name_with_prefixed_and_unprefixed(self): """After mapping is populated, manager resolves both prefixed and unprefixed tool names to the same server.""" @@ -4746,8 +4848,8 @@ async def test_call_tool_without_broken_pipe_error(self): # Register the server and map a tool to it manager.registry = {"test-server": server} - manager.tool_name_to_mcp_server_name_mapping["test_tool"] = "test-server" - manager.tool_name_to_mcp_server_name_mapping["test-server-test_tool"] = "test-server" + manager.tool_name_to_mcp_server_ids_mapping["test_tool"] = frozenset({"test-server"}) + manager.tool_name_to_mcp_server_ids_mapping["test-server-test_tool"] = frozenset({"test-server"}) # Create mock client that tracks call_tool usage mock_client = AsyncMock() @@ -5198,8 +5300,8 @@ def test_get_mcp_server_from_tool_name_uses_server_name_not_name(self): prefixed_tool_name = add_server_prefix_to_name(tool_name, "test_server") # Populate the mapping with the original tool name - manager.tool_name_to_mcp_server_name_mapping[tool_name] = "test_server" - manager.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = "test_server" + manager.tool_name_to_mcp_server_ids_mapping[tool_name] = frozenset({server.server_id}) + manager.tool_name_to_mcp_server_ids_mapping[prefixed_tool_name] = frozenset({server.server_id}) # Test: _get_mcp_server_from_tool_name should find the server using server.server_name # even when server.name is different @@ -7250,15 +7352,15 @@ async def _noop_handler(**kwargs): input_schema={"type": "object"}, handler=_noop_handler, ) - manager.tool_name_to_mcp_server_name_mapping["demo_tool"] = prefix - manager.tool_name_to_mcp_server_name_mapping[prefixed] = prefix + manager.tool_name_to_mcp_server_ids_mapping["demo_tool"] = frozenset({server.server_id}) + manager.tool_name_to_mcp_server_ids_mapping[prefixed] = frozenset({server.server_id}) await manager.update_server(self._make_server("evict-openapi", MCPApprovalStatus.rejected)) assert "evict-openapi" not in manager.registry assert prefixed not in global_mcp_tool_registry.tools - assert "demo_tool" not in manager.tool_name_to_mcp_server_name_mapping - assert prefixed not in manager.tool_name_to_mcp_server_name_mapping + assert "demo_tool" not in manager.tool_name_to_mcp_server_ids_mapping + assert prefixed not in manager.tool_name_to_mcp_server_ids_mapping async def test_update_server_noop_for_unregistered_pending(self): # update_server called with a pending row that was never registered From 85eaa03679570a034ccc87a26872eae40c594256 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 17 Jul 2026 03:52:59 -0700 Subject: [PATCH 2/8] fix(mcp): discriminate tool routes on a tag so reloads cannot skip the ambiguity check server.py selected the route variant with isinstance. A module reload rebinds the manager and its classes, so resolve_tool_route returned a reloaded variant while server.py still held the pre-reload class; isinstance was False, the ambiguity branch never ran, and the call fell through to the local registry as a 404 instead of the intended 409. Each variant now carries a Literal kind tag and server.py dispatches on it, which is what makes this a tagged union rather than one that leans on class identity. Withdrawing a server's routes also rebuilt the mapping into a fresh dict, orphaning any holder of the previous one. The initialize task is fired without being awaited and can be holding it, so the rebuild now updates the mapping in place. --- .../proxy/_experimental/mcp_server/mcp_server_manager.py | 7 ++++++- litellm/proxy/_experimental/mcp_server/server.py | 6 ++---- .../_experimental/mcp_server/test_mcp_server_manager.py | 6 +++--- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b8bf986776ef..f107b829a90b 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -192,6 +192,7 @@ class MCPToolRouteResolved: """Exactly one MCP server serves the tool name.""" server: MCPServer + kind: Literal["resolved"] = "resolved" @dataclass(frozen=True, slots=True) @@ -199,6 +200,7 @@ class MCPToolRouteNotFound: """No MCP server serves the tool name.""" tool_name: str + kind: Literal["not_found"] = "not_found" @dataclass(frozen=True, slots=True) @@ -207,6 +209,7 @@ class MCPToolRouteAmbiguous: tool_name: str server_ids: frozenset[str] + kind: Literal["ambiguous"] = "ambiguous" MCPToolRoute = Union[MCPToolRouteResolved, MCPToolRouteNotFound, MCPToolRouteAmbiguous] @@ -1550,11 +1553,13 @@ def _cleanup_server_tool_routing_artifacts(self, server: MCPServer) -> None: openapi_key_prefix = prefix_root + MCP_TOOL_PREFIX_SEPARATOR global_mcp_tool_registry.unregister_tools_with_prefix(openapi_key_prefix) - self.tool_name_to_mcp_server_ids_mapping = { + surviving = { tool_name: remaining for tool_name, owner_ids in self.tool_name_to_mcp_server_ids_mapping.items() if (remaining := owner_ids - {server.server_id}) } + self.tool_name_to_mcp_server_ids_mapping.clear() + self.tool_name_to_mcp_server_ids_mapping.update(surviving) def remove_server(self, mcp_server: LiteLLM_MCPServerTable): """ diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 0705c392f7fb..e18beadd3bfc 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -358,8 +358,6 @@ def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, - MCPToolRouteAmbiguous, - MCPToolRouteResolved, _caller_authorization_fans_out, _client_forwarded_authorization_headers, _should_strip_caller_authorization, @@ -2629,7 +2627,7 @@ async def execute_mcp_tool( name, allowed_server_ids=frozenset(server.server_id for server in allowed_mcp_servers), ) - if isinstance(route, MCPToolRouteAmbiguous): + if route.kind == "ambiguous": candidates = sorted( server.name for server in ( @@ -2648,7 +2646,7 @@ async def execute_mcp_tool( ), }, ) - mcp_server = route.server if isinstance(route, MCPToolRouteResolved) else None + mcp_server = route.server if route.kind == "resolved" else None else: mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) if mcp_server is None and requested_server 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 646c7a835daf..1116e3efaf1a 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 @@ -3948,7 +3948,7 @@ def test_resolve_tool_route_names_every_ambiguous_owner(self): route = manager.resolve_tool_route("echo") - assert type(route).__name__ == "MCPToolRouteAmbiguous" + assert route.kind == "ambiguous" assert route.server_ids == frozenset({"id-alpha", "id-zulu"}) def test_resolve_tool_route_resolves_sole_owner(self): @@ -3959,13 +3959,13 @@ def test_resolve_tool_route_resolves_sole_owner(self): route = manager.resolve_tool_route("echo") - assert type(route).__name__ == "MCPToolRouteResolved" + assert route.kind == "resolved" assert route.server is alpha def test_resolve_tool_route_reports_unknown_tool(self): manager = MCPServerManager() - assert type(manager.resolve_tool_route("nothing_serves_this")).__name__ == "MCPToolRouteNotFound" + assert manager.resolve_tool_route("nothing_serves_this").kind == "not_found" def test_cleanup_withdraws_only_departing_server_from_shared_tool_name(self): """Removing one server must leave a co-owned tool name routable to the survivor.""" From 2f7adb7b1154020ecda6016abd1bd42f7d0d13c9 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 17 Jul 2026 10:18:06 -0700 Subject: [PATCH 3/8] fix(mcp): fail closed when a tool's only owners are outside the caller's scope resolve_tool_route fell through to a scope-blind lookup when the tool had owners but none intersected the caller's allowed set, so a caller scoped to one server could resolve a tool served only by another. The downstream name-based permission check would then misroute the call to a reachable server sharing that name. The scoped branch is now authoritative: known owners with none in scope returns not_found rather than resolving a server the caller cannot reach. Also seed the logging tests from the specific server under test instead of every manager server id, so adding a second fixture server cannot make a tool look ambiguous. --- .../mcp_server/mcp_server_manager.py | 4 ++++ tests/mcp_tests/test_mcp_logging.py | 12 +++++++++--- .../mcp_server/test_mcp_server_manager.py | 17 +++++++++++++++++ 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index f107b829a90b..784b2385d196 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -4883,6 +4883,10 @@ def resolve_tool_route( scoped_server = self.get_mcp_server_by_id(next(iter(in_scope))) if scoped_server is not None: return MCPToolRouteResolved(server=scoped_server) + # The tool exists but no owner is reachable by this caller. Fail closed + # rather than fall through to a scope-blind lookup that could hand back a + # server outside the caller's scope. + return MCPToolRouteNotFound(tool_name=tool_name) server = self._get_mcp_server_from_tool_name(tool_name) if server is None: return MCPToolRouteNotFound(tool_name=tool_name) diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index d8193c3074f6..ba53b0777b12 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -126,7 +126,9 @@ def mock_client_constructor(*args, **kwargs): ) # Manually add the tool mapping to ensure it's available (since mocking might not capture it properly) - zapier_server_ids = frozenset(local_mcp_server_manager.get_all_mcp_server_ids()) + zapier_server = local_mcp_server_manager.get_mcp_server_by_name("zapier_gmail_server") + assert zapier_server is not None + zapier_server_ids = frozenset({zapier_server.server_id}) local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[ "add_tools" ] = zapier_server_ids @@ -241,7 +243,9 @@ def mock_client_constructor(*args, **kwargs): await local_mcp_server_manager._initialize_tool_name_to_mcp_server_ids_mapping() # Manually add the tool mapping to ensure it's available (since mocking might not capture it properly) - test_server_ids = frozenset(local_mcp_server_manager.get_all_mcp_server_ids()) + test_server = local_mcp_server_manager.get_mcp_server_by_name("test_server") + assert test_server is not None + test_server_ids = frozenset({test_server.server_id}) local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping[ "expensive_tool" ] = test_server_ids @@ -406,7 +410,9 @@ def mock_client_constructor(*args, **kwargs): await local_mcp_server_manager._initialize_tool_name_to_mcp_server_ids_mapping() # Manually add the tool mapping to ensure it's available (since mocking might not capture it properly) - zapier_server_ids = frozenset(local_mcp_server_manager.get_all_mcp_server_ids()) + zapier_server = local_mcp_server_manager.get_mcp_server_by_name("zapier_gmail_server") + assert zapier_server is not None + zapier_server_ids = frozenset({zapier_server.server_id}) local_mcp_server_manager.tool_name_to_mcp_server_ids_mapping["add_tools"] = ( zapier_server_ids ) 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 1116e3efaf1a..cfb360f30ac4 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 @@ -3937,6 +3937,23 @@ def test_get_mcp_server_from_tool_name_refuses_prefixed_name_of_duplicate_named_ assert manager._get_mcp_server_from_tool_name("shared-echo") is None + def test_resolve_tool_route_is_not_found_when_no_owner_is_in_scope(self): + """A tool whose only owner is outside the caller's scope must fail closed. + + Falling through to the scope-blind lookup would hand back a server the caller + cannot reach, and the downstream name-based permission check would misroute it + to a reachable same-named server. + """ + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) + zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha, "id-zulu": zulu} + manager._register_tool_route("secret_tool", "id-zulu") + + route = manager.resolve_tool_route("secret_tool", allowed_server_ids=frozenset({"id-alpha"})) + + assert route.kind == "not_found" + def test_resolve_tool_route_names_every_ambiguous_owner(self): """The ambiguous route carries all owners so callers can report the real candidates.""" manager = MCPServerManager() From cd3a92273294a7ff39156d223999ee8445e3e86e Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 18 Jul 2026 19:21:59 -0700 Subject: [PATCH 4/8] fix(mcp): keep out-of-scope servers unreachable through the credential-resolution fallback resolve_tool_route fails closed when a tool's only owners are outside the caller's scope, but execute_mcp_tool then re-resolved a still-unset server with a scope-blind lookup so BYOK and credential injection could run on every path. On the JSON-RPC and tool-search paths that fallback undid the scope decision: because no server_name was set for an out-of-scope tool, the permission gate was skipped, and the fallback dispatched to the out-of-scope server with its credentials. The fallback now only accepts a server the caller is already allowed to reach, matching how the requested_server lookups are guarded by the server-id mismatch check. Cleanup of a departed server's routes withdrew its id per row but did so by clearing the whole mapping and repopulating it, which momentarily emptied the shared dict the un-awaited initialize task writes into. It now withdraws the id key by key, so the dict is never empty and never rebound. --- .../mcp_server/mcp_server_manager.py | 13 ++-- .../proxy/_experimental/mcp_server/server.py | 7 +- .../mcp_server/test_mcp_server.py | 65 +++++++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 22 +++++++ 4 files changed, 99 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index ae790575015e..bfa786f4bce9 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1558,13 +1558,12 @@ def _cleanup_server_tool_routing_artifacts(self, server: MCPServer) -> None: openapi_key_prefix = prefix_root + MCP_TOOL_PREFIX_SEPARATOR global_mcp_tool_registry.unregister_tools_with_prefix(openapi_key_prefix) - surviving = { - tool_name: remaining - for tool_name, owner_ids in self.tool_name_to_mcp_server_ids_mapping.items() - if (remaining := owner_ids - {server.server_id}) - } - self.tool_name_to_mcp_server_ids_mapping.clear() - self.tool_name_to_mcp_server_ids_mapping.update(surviving) + for tool_name in list(self.tool_name_to_mcp_server_ids_mapping): + remaining = self.tool_name_to_mcp_server_ids_mapping[tool_name] - {server.server_id} + if remaining: + self.tool_name_to_mcp_server_ids_mapping[tool_name] = remaining + else: + del self.tool_name_to_mcp_server_ids_mapping[tool_name] def remove_server(self, mcp_server: LiteLLM_MCPServerTable): """ diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 99b1e6b42079..23e853a32d3d 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2690,8 +2690,13 @@ async def execute_mcp_tool( 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). + # Stay within the caller's scope: a scope-blind lookup here would undo the + # scope decision already made above and dispatch to an out-of-scope server. if mcp_server is None: - mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + allowed_server_ids = {server.server_id for server in allowed_mcp_servers} + fallback_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + if fallback_server is not None and fallback_server.server_id in allowed_server_ids: + mcp_server = fallback_server if mcp_server: standard_logging_mcp_tool_call["mcp_server_cost_info"] = (mcp_server.mcp_info or {}).get( 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 1576a4f4270e..2d0bba5efc5a 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 @@ -6161,6 +6161,71 @@ async def fake_create_mcp_client(server, **kwargs): assert "server" not in injected +@pytest.mark.asyncio +async def test_execute_mcp_tool_jsonrpc_never_dispatches_to_an_out_of_scope_server(): + """End to end: a tool whose only owner is outside the caller's scope must not be dispatched. + + resolve_tool_route already fails closed, but a later scope-blind fallback that only runs + when no server resolved would re-resolve the out-of-scope owner and dispatch to it, + bypassing the permission gate because no server_name was set for it to check. + """ + from litellm.proxy._experimental.mcp_server import server as mcp_module + + in_scope = MCPServer( + server_id="in-scope-id", + name="echo_alpha", + server_name="echo_alpha", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + authentication_token="in-scope-secret", + ) + out_of_scope = MCPServer( + server_id="out-of-scope-id", + name="echo_zulu", + server_name="echo_zulu", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + authentication_token="out-of-scope-secret", + ) + + dispatched: dict = {} + + async def fake_create_mcp_client(server, **kwargs): + dispatched["server"] = server + raise AssertionError("must not dispatch to an out-of-scope server") + + with ( + patch.dict( + mcp_module.global_mcp_server_manager.tool_name_to_mcp_server_ids_mapping, + {"secret_tool": frozenset({out_of_scope.server_id})}, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "get_registry", + return_value={in_scope.server_id: in_scope, out_of_scope.server_id: out_of_scope}, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "_create_mcp_client", + new=fake_create_mcp_client, + ), + patch.object(mcp_module.MCPRequestHandler, "is_tool_allowed", return_value=True), + patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=None), + patch("litellm.proxy.proxy_server.proxy_logging_obj", None), + pytest.raises(HTTPException), + ): + await mcp_module.execute_mcp_tool( + name="secret_tool", + arguments={"message": "hello"}, + allowed_mcp_servers=[in_scope], + start_time=datetime.now(), + ) + + assert "server" not in dispatched + + @pytest.mark.asyncio async def test_execute_mcp_tool_jsonrpc_unprefixed_resolves_when_caller_reaches_one_server(): """A caller scoped to one server is served unprefixed names and must still route. 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 fd2e5046c134..cc6cf64b023b 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 @@ -4006,6 +4006,28 @@ def test_cleanup_withdraws_only_departing_server_from_shared_tool_name(self): assert manager.tool_name_to_mcp_server_ids_mapping["echo"] == frozenset({"id-alpha"}) assert manager._get_mcp_server_from_tool_name("echo") is alpha + def test_cleanup_mutates_the_mapping_in_place_without_emptying_it(self): + """Cleanup must keep the same dict object and never clear it wholesale. + + The initialize task is dispatched without being awaited and holds a reference to + this dict, so rebinding it would drop that task's writes; clearing it would let a + concurrent reader briefly see an empty map and mis-route. + """ + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) + zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha, "id-zulu": zulu} + manager._register_tool_route("shared", "id-alpha") + manager._register_tool_route("shared", "id-zulu") + manager._register_tool_route("alpha_only", "id-alpha") + original_map = manager.tool_name_to_mcp_server_ids_mapping + + manager._cleanup_server_tool_routing_artifacts(zulu) + + assert manager.tool_name_to_mcp_server_ids_mapping is original_map + assert manager.tool_name_to_mcp_server_ids_mapping["shared"] == frozenset({"id-alpha"}) + assert manager.tool_name_to_mcp_server_ids_mapping["alpha_only"] == frozenset({"id-alpha"}) + def test_cleanup_drops_the_route_when_its_last_owner_leaves(self): manager = MCPServerManager() alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) From b0908d5345fb3bb75212b689cf0142cad8674f5a Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 20 Jul 2026 18:57:08 -0700 Subject: [PATCH 5/8] test(mcp): seed tool ownership from the specific server under test The call_tool and list_tools mock tests seeded a tool's owning-server set from every id in the registry. That holds only while the fixture has exactly one server; adding a second would make the tool look multi-owned and resolve as ambiguous, failing the test for an unrelated reason. Each site now derives the id from the server it actually loaded, matching the pattern already used in the alias-prefixing tests and the logging tests. --- tests/mcp_tests/test_mcp_server.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 935bfa482e77..7c172ee188a3 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -226,17 +226,20 @@ def mock_client_constructor(*args, **kwargs): # Verify tool mapping was updated expected_prefix = "test_http_server" + test_http_server = test_manager.get_mcp_server_by_name("test_http_server") + assert test_http_server is not None + expected_owner = frozenset({test_http_server.server_id}) assert ( test_manager.tool_name_to_mcp_server_ids_mapping[ f"{expected_prefix}-gmail_send_email" ] - == frozenset(allowed_server_ids) + == expected_owner ) assert ( test_manager.tool_name_to_mcp_server_ids_mapping[ f"{expected_prefix}-calendar_create_event" ] - == frozenset(allowed_server_ids) + == expected_owner ) @@ -281,8 +284,10 @@ def mock_client_constructor(*args, **kwargs): ) # Manually set up tool mapping (normally done by list_tools) + test_http_server = test_manager.get_mcp_server_by_name("test_http_server") + assert test_http_server is not None test_manager.tool_name_to_mcp_server_ids_mapping["gmail_send_email"] = frozenset( - test_manager.get_registry().keys() + {test_http_server.server_id} ) # Call the tool @@ -347,8 +352,10 @@ def mock_client_constructor(*args, **kwargs): ) # Manually set up tool mapping + test_http_server = test_manager.get_mcp_server_by_name("test_http_server") + assert test_http_server is not None test_manager.tool_name_to_mcp_server_ids_mapping["gmail_send_email"] = frozenset( - test_manager.get_registry().keys() + {test_http_server.server_id} ) # Call the tool with invalid data @@ -389,8 +396,10 @@ async def test_mcp_http_transport_tool_not_found(): ) # Mapping populated for this server but not for the requested tool + test_http_server = test_manager.get_mcp_server_by_name("test_http_server") + assert test_http_server is not None test_manager.tool_name_to_mcp_server_ids_mapping["gmail_send_email"] = frozenset( - test_manager.get_registry().keys() + {test_http_server.server_id} ) # Try to call a tool that doesn't exist in mapping From ad888fcf5f8330f9d8385b0e518afc9ba8c1f25a Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 21 Jul 2026 00:11:35 -0700 Subject: [PATCH 6/8] fix(mcp): make the ownership map the single authority when resolving a tool name _get_mcp_server_from_tool_name had two competing sources of truth: the ownership map, and a first-wins prefix-to-server scan of the registry. The prefix path ignored the map's ambiguity, so a prefixed name owned by two servers sharing a prefix resolved to whichever was scanned first instead of None, breaking the function's own contract. It also never checked that the prefix-named server actually owned the tool. Resolution now goes through the ownership map first: a registered tool name with one owner resolves by id and with several owners returns None. Only a name that is not itself registered falls to prefix extraction, and that path resolves to the one server that both matches the prefix and owns the underlying tool, so a shared prefix or a prefix naming a non-owner both stay unresolved. This removes the class of arbitrary-owner resolution that callers were each guarding against individually. The duplicate-prefix regression test set only name, not server_name, so the prefix was never recognized and the prefix path was never exercised; it now sets server_name and a sibling test covers a prefix that names a server which does not own the tool. --- .../mcp_server/mcp_server_manager.py | 63 +++++++++---------- .../mcp_server/test_mcp_server_manager.py | 21 ++++++- 2 files changed, 49 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index bfa786f4bce9..44336cc4ca41 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -4902,40 +4902,39 @@ def _get_mcp_server_from_tool_name(self, tool_name: str) -> Optional[MCPServer]: servers expose ``tool_name``; use ``resolve_tool_route`` to tell the two cases apart. """ - registry_servers = list(self.get_registry().values()) - - # Build prefix → server lookup covering every known form a tool name - # may take (alias / server_name / server_id / short ID). This is what - # makes the short-prefix mode work without breaking historical names. - prefix_to_server: dict[str, MCPServer] = {} - for server in registry_servers: - for known_prefix in iter_known_server_prefixes(server): - normalised = normalize_server_name(known_prefix) - prefix_to_server.setdefault(normalised, server) - - # First try with the original tool name + # The ownership map is the single source of truth. When ``tool_name`` is a + # registered tool (in any of its prefixed or unprefixed forms), a unique + # owner resolves by id and several owners are ambiguous, so we never collapse + # to an arbitrary one. owner_ids = self.tool_name_to_mcp_server_ids_mapping.get(tool_name) - if owner_ids is not None and len(owner_ids) == 1: - owned = self.get_mcp_server_by_id(next(iter(owner_ids))) - if owned is not None: - return owned - - # If not found and tool name is prefixed, extract the prefix and - # match against any known form. - if is_tool_name_prefixed(tool_name, known_server_prefixes=set(prefix_to_server.keys())): - ( - original_tool_name, - server_name_from_prefix, - ) = split_server_prefix_from_name(tool_name) - normalised_prefix = normalize_server_name(server_name_from_prefix) - matched_server = prefix_to_server.get(normalised_prefix) - if matched_server is not None and ( - original_tool_name in self.tool_name_to_mcp_server_ids_mapping - or tool_name in self.tool_name_to_mcp_server_ids_mapping - ): - return matched_server + if owner_ids is not None: + return self.get_mcp_server_by_id(next(iter(owner_ids))) if len(owner_ids) == 1 else None + + # ``tool_name`` is not itself registered. Interpret it as + # ``-`` and resolve to the one server that both matches + # the prefix and owns the underlying tool. A prefix shared by several servers + # (server names are not unique) stays ambiguous and resolves to None. + known_prefixes = { + normalize_server_name(known_prefix) + for server in self.get_registry().values() + for known_prefix in iter_known_server_prefixes(server) + } + if not is_tool_name_prefixed(tool_name, known_server_prefixes=known_prefixes): + return None - return None + original_tool_name, server_name_from_prefix = split_server_prefix_from_name(tool_name) + tool_owner_ids = self.tool_name_to_mcp_server_ids_mapping.get(original_tool_name) + if not tool_owner_ids: + return None + + normalised_prefix = normalize_server_name(server_name_from_prefix) + prefix_owners = [ + server + for server in self.get_registry().values() + if server.server_id in tool_owner_ids + and normalised_prefix in {normalize_server_name(p) for p in iter_known_server_prefixes(server)} + ] + return prefix_owners[0] if len(prefix_owners) == 1 else None async def reload_servers_from_database(self): """Re-synchronize the in-memory MCP server registry with the database.""" 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 cc6cf64b023b..0fdb6be4b44f 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 @@ -3934,16 +3934,31 @@ def test_get_mcp_server_from_tool_name_resolves_sole_owner_by_id(self): assert manager._get_mcp_server_from_tool_name("echo") is alpha def test_get_mcp_server_from_tool_name_refuses_prefixed_name_of_duplicate_named_servers(self): - """server_name is not unique, so a prefix shared by two servers must not silently pick one.""" + """server_name is not unique, so a prefix shared by two servers must not silently pick one. + + server_name is set so ``shared`` is a recognized prefix and the prefix-extraction path + is actually exercised; without it the name would not be seen as prefixed at all. + """ manager = MCPServerManager() - first = MCPServer(server_id="id-first", name="shared", transport=MCPTransport.http) - second = MCPServer(server_id="id-second", name="shared", transport=MCPTransport.http) + first = MCPServer(server_id="id-first", name="shared", server_name="shared", transport=MCPTransport.http) + second = MCPServer(server_id="id-second", name="shared", server_name="shared", transport=MCPTransport.http) manager.registry = {"id-first": first, "id-second": second} manager._register_tool_route("shared-echo", "id-first") manager._register_tool_route("shared-echo", "id-second") assert manager._get_mcp_server_from_tool_name("shared-echo") is None + def test_get_mcp_server_from_tool_name_prefix_must_own_the_tool(self): + """A prefix that names a server which does not own the tool must not resolve to it.""" + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="alpha", server_name="alpha", transport=MCPTransport.http) + beta = MCPServer(server_id="id-beta", name="beta", server_name="beta", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha, "id-beta": beta} + manager._register_tool_route("echo", "id-alpha") + + assert manager._get_mcp_server_from_tool_name("beta-echo") is None + assert manager._get_mcp_server_from_tool_name("alpha-echo") is alpha + def test_resolve_tool_route_is_not_found_when_no_owner_is_in_scope(self): """A tool whose only owner is outside the caller's scope must fail closed. From 0353b2450948cfaa3b1f815b483d504afe20d5bd Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 25 Jul 2026 14:51:54 -0700 Subject: [PATCH 7/8] fix(mcp): make tool-route registration authoritative per server Registering a server's tool routes was union-only, so a routing row could only ever gain owners. Nothing withdrew a tool while its server stayed in the registry: `_cleanup_server_tool_routing_artifacts` withdraws a server's id from every row, but it only runs when the server leaves the registry. So an upstream that stopped exposing a tool left its owner pinned. A name then served by exactly one reachable server kept resolving as ambiguous and kept returning the 409, with no way back short of restarting the proxy. A server's `tools/list` result is its complete listing, so it is the truth rather than an increment. `_replace_server_tool_routes` replaces one server's rows instead of accumulating into them, withdrawing its id from a name it no longer serves and dropping the row once no owner is left. Rows still accumulate across servers, so a genuinely shared name stays ambiguous. Treating a listing as the truth is safe because `_fetch_tools_with_timeout` raises on every failure instead of returning an empty list, and caller-scoped narrowing (`check_allowed_or_banned_tools`, semantic filtering) runs downstream, so neither a failed listing nor one caller's filtered view can evict routes another caller needs. Eviction is the same operation with an empty set, so it delegates rather than keeping its own copy of the withdrawal arithmetic. An earlier cut of this carried a per-server reverse index to avoid scanning the map, which bought under a millisecond on a path that has just made a network round trip, in exchange for a second source of truth that can disagree with the first. It already had: pointing eviction at the index broke `test_update_server_eviction_clears_openapi_routing_artifacts`, which seeds the mapping directly and asserts eviction clears rows however they were written. The scan is exhaustive by construction, so that class is gone. The OpenAPI path replaces once the whole spec has parsed, so a mid-loop failure leaves the previous routes intact rather than committing a partial set that would withdraw operations still served. The startup warm-up no longer re-registers what listing already recorded, which would have re-added names outside the replace and reintroduced rows it cannot withdraw. The OpenAPI registration call site turned out to have no test coverage at all; deleting it left the suite green. It now has three, including the mid-parse failure case. --- .../mcp_server/mcp_server_manager.py | 106 ++++--- .../mcp_server/test_mcp_server_manager.py | 278 ++++++++++++++++-- 2 files changed, 329 insertions(+), 55 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 44336cc4ca41..7bcccd1841da 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -15,7 +15,7 @@ import time from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Any, AsyncIterator, Callable, Literal, Optional, Union, cast +from typing import Any, AsyncIterator, Callable, Collection, Literal, Optional, Union, cast from urllib.parse import urlparse import anyio @@ -1482,6 +1482,7 @@ async def _register_openapi_tools(self, spec_path: str, server: MCPServer, base_ paths = spec.get("paths", {}) components = spec.get("components", {}) registered_count = 0 + route_names: set[str] = set() verbose_logger.debug(f"Processing {len(paths)} paths from OpenAPI spec") @@ -1524,13 +1525,18 @@ async def _register_openapi_tools(self, spec_path: str, server: MCPServer, base_ handler=tool_func, ) - # Update tool name to server id mapping (for both prefixed and base names) - self._register_tool_route(base_tool_name, server.server_id) - self._register_tool_route(prefixed_tool_name, server.server_id) + # Collect tool name to server id routes (both prefixed and base names) + route_names.add(base_tool_name) + route_names.add(prefixed_tool_name) registered_count += 1 verbose_logger.debug(f"Registered OpenAPI tool: {prefixed_tool_name} for server {server.name}") + # Only once the whole spec parsed: a mid-loop raise must leave the + # previous routes intact rather than commit a partial set, which would + # withdraw routes for operations later in the spec that are still served. + self._replace_server_tool_routes(server.server_id, route_names) + verbose_logger.info(f"Successfully registered {registered_count} OpenAPI tools for server {server.name}") except Exception as e: @@ -1558,12 +1564,7 @@ def _cleanup_server_tool_routing_artifacts(self, server: MCPServer) -> None: openapi_key_prefix = prefix_root + MCP_TOOL_PREFIX_SEPARATOR global_mcp_tool_registry.unregister_tools_with_prefix(openapi_key_prefix) - for tool_name in list(self.tool_name_to_mcp_server_ids_mapping): - remaining = self.tool_name_to_mcp_server_ids_mapping[tool_name] - {server.server_id} - if remaining: - self.tool_name_to_mcp_server_ids_mapping[tool_name] = remaining - else: - del self.tool_name_to_mcp_server_ids_mapping[tool_name] + self._replace_server_tool_routes(server.server_id, frozenset()) def remove_server(self, mcp_server: LiteLLM_MCPServerTable): """ @@ -3821,6 +3822,9 @@ def _create_prefixed_tools(self, tools: list[MCPTool], server: MCPServer, add_pr """ prefixed_tools = [] prefix = get_server_prefix(server) + # Server-level, so resolve once instead of per tool. + known_prefixes = tuple(iter_known_server_prefixes(server)) + route_names: set[str] = set() for tool in tools: tool_copy = tool.model_copy(deep=True) @@ -3834,13 +3838,17 @@ def _create_prefixed_tools(self, tools: list[MCPTool], server: MCPServer, add_pr tool_copy.name = name_to_use prefixed_tools.append(tool_copy) - # Register every known prefix form (alias, server_name, server_id, + # Collect every known prefix form (alias, server_name, server_id, # short ID) so call_tool can resolve regardless of which form a # caller / cached client is using. - self._register_tool_route(original_name, server.server_id) - for known_prefix in iter_known_server_prefixes(server): - qualified = add_server_prefix_to_name(original_name, known_prefix) - self._register_tool_route(qualified, server.server_id) + route_names.add(original_name) + for known_prefix in known_prefixes: + route_names.add(add_server_prefix_to_name(original_name, known_prefix)) + + # `tools` is this server's complete upstream listing, so it is authoritative: + # replace rather than accumulate, or a tool the upstream dropped keeps its + # owner and makes the name look ambiguous until restart. + self._replace_server_tool_routes(server.server_id, route_names) verbose_logger.info(f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}") return prefixed_tools @@ -4818,15 +4826,21 @@ def initialize_tool_name_to_mcp_server_ids_mapping(self): async def _initialize_tool_name_to_mcp_server_ids_mapping(self): """ - Call list_tools for each server and update the tool name to MCP server name mapping - Note: This now handles prefixed tool names + Warm the tool-name routing map by listing every server once at startup. + + Listing is what registers the routes: ``_get_tools_from_server`` hands the + server's complete tool list to ``_create_prefixed_tools``, which replaces + that server's routes via ``_replace_server_tool_routes``. OpenAPI servers + registered theirs from the spec in ``_register_openapi_tools``. So nothing is + registered here — doing so would re-add names outside the per-server replace + and reintroduce rows it cannot withdraw. """ for server in self.get_registry().values(): if server.needs_user_oauth_token: # Skip OAuth2 servers that rely on user-provided tokens continue try: - tools = await self._get_tools_from_server(server) + await self._get_tools_from_server(server) except MCPUpstreamAuthError as e: # Pass-through servers expect a user-supplied bearer token; # at startup we have none, so an upstream 401 is normal. @@ -4840,22 +4854,48 @@ async def _initialize_tool_name_to_mcp_server_ids_mapping(self): f"Failed to get tools from server {server.name} during tool name mapping initialization: {str(e)}" ) continue - for tool in tools: - # The tool.name here is already prefixed from _get_tools_from_server - # Extract original name for mapping - original_name, _ = split_server_prefix_from_name(tool.name) - self._register_tool_route(original_name, server.server_id) - self._register_tool_route(tool.name, server.server_id) - - def _register_tool_route(self, tool_name: str, server_id: str) -> None: - """Record that ``server_id`` serves ``tool_name``, preserving any other owners. - - Routing rows accumulate rather than overwrite so that a tool name served by - more than one server stays resolvable as ambiguous instead of silently - collapsing to whichever server registered last. + + def _replace_server_tool_routes(self, server_id: str, route_names: Collection[str]) -> None: + """Make ``route_names`` the complete set of routing rows owned by ``server_id``. + + Rows accumulate *across* servers, so a name several serve stays ambiguous. But + within one server this is authoritative rather than additive: a name it no + longer serves has ``server_id`` withdrawn and the row goes once unowned. + Additive registration pinned owners forever, so a dropped upstream tool kept + its name looking ambiguous, a false 409, until restart. Empty ``route_names`` + means owning nothing, which is how eviction withdraws a departing server. + + Callers must pass a *complete* listing, and both do. Treating one as the truth + is safe because ``_fetch_tools_with_timeout`` raises on every failure instead + of returning an empty list, and caller-scoped narrowing + (``check_allowed_or_banned_tools``, semantic filtering) runs downstream, so + neither a failed listing nor one caller's filtered view can evict another's + routes. Scanning every row beats a per-server index here: it runs once per + server per listing, right after that server's network round trip, so well + under a millisecond amortizes into tens of milliseconds of I/O, and a second + source of truth could disagree with this one. + + Pre-existing and unchanged: ``ClientSession.list_tools()`` is called without a + cursor, so a paginating upstream only yields page one. Every path populating + this map shares that listing, so replacing cannot drop a legitimate row. """ - owners = self.tool_name_to_mcp_server_ids_mapping.get(tool_name, frozenset()) - self.tool_name_to_mcp_server_ids_mapping[tool_name] = owners | {server_id} + desired = frozenset(route_names) + + # Mutated in place: the startup listing task is dispatched without being + # awaited and holds a reference to this dict, so rebinding drops its writes. + for tool_name in list(self.tool_name_to_mcp_server_ids_mapping): + owners = self.tool_name_to_mcp_server_ids_mapping[tool_name] + if server_id not in owners or tool_name in desired: + continue + remaining = owners - {server_id} + if remaining: + self.tool_name_to_mcp_server_ids_mapping[tool_name] = remaining + else: + del self.tool_name_to_mcp_server_ids_mapping[tool_name] + + for tool_name in desired: + owners = self.tool_name_to_mcp_server_ids_mapping.get(tool_name, frozenset()) + self.tool_name_to_mcp_server_ids_mapping[tool_name] = owners | {server_id} def resolve_tool_route( self, 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 0fdb6be4b44f..f9de88610d7d 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 @@ -3562,6 +3562,159 @@ async def tool_func(**kwargs): assert captured["headers"] is not None assert captured["headers"]["Authorization"] == "STATIC token" + @staticmethod + def _openapi_spec_with(operation_ids: list[str]) -> str: + return json.dumps( + { + "openapi": "3.0.0", + "info": {"title": "Demo", "version": "1.0.0"}, + "paths": {f"/{op}": {"get": {"operationId": op, "summary": op}} for op in operation_ids}, + } + ) + + @staticmethod + def _openapi_patches(): + async def tool_func(**kwargs): + return "ok" + + return ( + patch( + "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.create_tool_function", + return_value=tool_func, + ), + patch( + "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.build_input_schema", + return_value={"type": "object", "properties": {}, "required": []}, + ), + patch( + "litellm.proxy._experimental.mcp_server.tool_registry.global_mcp_tool_registry.register_tool", + return_value=None, + ), + ) + + @pytest.mark.asyncio + async def test_register_openapi_tools_records_base_and_prefixed_routes(self, tmp_path): + """OpenAPI operations must be routable by both their base and prefixed names. + + Nothing else covered this call site: deleting the registration entirely left + the whole mcp_server suite green. + """ + manager = MCPServerManager() + spec_path = tmp_path / "openapi.json" + spec_path.write_text(self._openapi_spec_with(["health_check"])) + server = MCPServer( + server_id="openapi-server", + name="openapi-server", + server_name="openapi-server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + create_fn, build_schema, register = self._openapi_patches() + with create_fn, build_schema, register: + await manager._register_openapi_tools( + spec_path=str(spec_path), server=server, base_url="https://example.com" + ) + + assert manager.tool_name_to_mcp_server_ids_mapping["health_check"] == frozenset({"openapi-server"}) + assert manager.tool_name_to_mcp_server_ids_mapping["openapi-server-health_check"] == frozenset( + {"openapi-server"} + ) + + @pytest.mark.asyncio + async def test_reregistering_openapi_tools_withdraws_a_dropped_operation(self, tmp_path): + """Re-parsing a spec that lost an operation must withdraw that route. + + The spec is the server's complete listing, so registration replaces rather + than accumulates here for the same reason it does on the tools/list path. + """ + manager = MCPServerManager() + spec_path = tmp_path / "openapi.json" + server = MCPServer( + server_id="openapi-server", + name="openapi-server", + server_name="openapi-server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + spec_path.write_text(self._openapi_spec_with(["health_check", "legacy_probe"])) + create_fn, build_schema, register = self._openapi_patches() + with create_fn, build_schema, register: + await manager._register_openapi_tools( + spec_path=str(spec_path), server=server, base_url="https://example.com" + ) + assert "legacy_probe" in manager.tool_name_to_mcp_server_ids_mapping + + spec_path.write_text(self._openapi_spec_with(["health_check"])) + create_fn, build_schema, register = self._openapi_patches() + with create_fn, build_schema, register: + await manager._register_openapi_tools( + spec_path=str(spec_path), server=server, base_url="https://example.com" + ) + + assert "legacy_probe" not in manager.tool_name_to_mcp_server_ids_mapping + assert "openapi-server-legacy_probe" not in manager.tool_name_to_mcp_server_ids_mapping + assert manager.tool_name_to_mcp_server_ids_mapping["health_check"] == frozenset({"openapi-server"}) + + @pytest.mark.asyncio + async def test_failed_openapi_reregistration_keeps_the_previous_routes(self, tmp_path): + """A mid-parse failure must not commit a partial set of routes. + + Replacing per operation would withdraw the operations after the failure point + even though the server still serves them. + """ + manager = MCPServerManager() + spec_path = tmp_path / "openapi.json" + server = MCPServer( + server_id="openapi-server", + name="openapi-server", + server_name="openapi-server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + spec_path.write_text(self._openapi_spec_with(["health_check", "legacy_probe"])) + create_fn, build_schema, register = self._openapi_patches() + with create_fn, build_schema, register: + await manager._register_openapi_tools( + spec_path=str(spec_path), server=server, base_url="https://example.com" + ) + + async def tool_func(**kwargs): + return "ok" + + # Fail on the *second* operation, so the first has already been collected. + # Replacing per operation instead of after the loop would commit that partial + # set here and withdraw legacy_probe, which the server still serves. + calls = iter([tool_func, RuntimeError("spec blew up")]) + + def _second_call_explodes(*args, **kwargs): + outcome = next(calls) + if isinstance(outcome, Exception): + raise outcome + return outcome + + _, build_schema, register = self._openapi_patches() + with ( + patch( + "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.create_tool_function", + side_effect=_second_call_explodes, + ), + build_schema, + register, + pytest.raises(RuntimeError), + ): + await manager._register_openapi_tools( + spec_path=str(spec_path), server=server, base_url="https://example.com" + ) + + assert manager.tool_name_to_mcp_server_ids_mapping["health_check"] == frozenset({"openapi-server"}) + assert manager.tool_name_to_mcp_server_ids_mapping["legacy_probe"] == frozenset({"openapi-server"}) + @pytest.mark.asyncio async def test_pre_call_tool_check_allowed_tools_list_allows_tool(self): """Test pre_call_tool_check allows tool when it's in allowed_tools list""" @@ -3897,30 +4050,65 @@ def test_resolve_mcp_server_for_tool_call_via_alias(self): resolved = manager._resolve_mcp_server_for_tool_call("zapier-alias", "create_zap") assert resolved is server - def test_register_tool_route_accumulates_owners_across_servers(self): + def test_replace_server_tool_routes_accumulates_owners_across_servers(self): """Two servers exposing one tool name are both recorded, not last-writer-wins.""" manager = MCPServerManager() - manager._register_tool_route("echo", "id-alpha") - manager._register_tool_route("echo", "id-zulu") + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-zulu", {"echo"}) assert manager.tool_name_to_mcp_server_ids_mapping["echo"] == frozenset({"id-alpha", "id-zulu"}) - def test_register_tool_route_is_idempotent_for_one_server(self): + def test_replace_server_tool_routes_is_idempotent_for_one_server(self): """Re-listing the same server must not make its own tool look ambiguous.""" manager = MCPServerManager() - manager._register_tool_route("echo", "id-alpha") - manager._register_tool_route("echo", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-alpha", {"echo"}) assert manager.tool_name_to_mcp_server_ids_mapping["echo"] == frozenset({"id-alpha"}) + def test_replace_server_tool_routes_withdraws_a_tool_the_upstream_dropped(self): + """A re-listing without a previously served tool must withdraw only that owner. + + Union-only registration would keep id-alpha pinned to "echo" forever, so a + name now served by id-zulu alone would keep resolving as ambiguous (409) + until the proxy restarted. + """ + manager = MCPServerManager() + manager._replace_server_tool_routes("id-alpha", {"echo", "still_there"}) + manager._replace_server_tool_routes("id-zulu", {"echo"}) + + # id-alpha's upstream stops exposing "echo"; "still_there" is unaffected. + manager._replace_server_tool_routes("id-alpha", {"still_there"}) + + assert manager.tool_name_to_mcp_server_ids_mapping["echo"] == frozenset({"id-zulu"}) + assert manager.tool_name_to_mcp_server_ids_mapping["still_there"] == frozenset({"id-alpha"}) + + def test_replace_server_tool_routes_drops_a_row_whose_last_owner_dropped_it(self): + """The row disappears, rather than lingering as an empty owner set.""" + manager = MCPServerManager() + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-alpha", set()) + + assert "echo" not in manager.tool_name_to_mcp_server_ids_mapping + + def test_replace_server_tool_routes_leaves_other_servers_rows_untouched(self): + """Replacing one server's routes must not disturb a row it never owned.""" + manager = MCPServerManager() + manager._replace_server_tool_routes("id-zulu", {"zulu_only"}) + manager._replace_server_tool_routes("id-alpha", {"alpha_only"}) + + manager._replace_server_tool_routes("id-alpha", set()) + + assert manager.tool_name_to_mcp_server_ids_mapping["zulu_only"] == frozenset({"id-zulu"}) + def test_get_mcp_server_from_tool_name_refuses_ambiguous_unprefixed_name(self): """An unprefixed name owned by several servers resolves to nothing, never to one of them.""" manager = MCPServerManager() alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-zulu": zulu} - manager._register_tool_route("echo", "id-alpha") - manager._register_tool_route("echo", "id-zulu") + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-zulu", {"echo"}) assert manager._get_mcp_server_from_tool_name("echo") is None @@ -3929,7 +4117,7 @@ def test_get_mcp_server_from_tool_name_resolves_sole_owner_by_id(self): manager = MCPServerManager() alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha} - manager._register_tool_route("echo", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"echo"}) assert manager._get_mcp_server_from_tool_name("echo") is alpha @@ -3943,8 +4131,8 @@ def test_get_mcp_server_from_tool_name_refuses_prefixed_name_of_duplicate_named_ first = MCPServer(server_id="id-first", name="shared", server_name="shared", transport=MCPTransport.http) second = MCPServer(server_id="id-second", name="shared", server_name="shared", transport=MCPTransport.http) manager.registry = {"id-first": first, "id-second": second} - manager._register_tool_route("shared-echo", "id-first") - manager._register_tool_route("shared-echo", "id-second") + manager._replace_server_tool_routes("id-first", {"shared-echo"}) + manager._replace_server_tool_routes("id-second", {"shared-echo"}) assert manager._get_mcp_server_from_tool_name("shared-echo") is None @@ -3954,7 +4142,7 @@ def test_get_mcp_server_from_tool_name_prefix_must_own_the_tool(self): alpha = MCPServer(server_id="id-alpha", name="alpha", server_name="alpha", transport=MCPTransport.http) beta = MCPServer(server_id="id-beta", name="beta", server_name="beta", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-beta": beta} - manager._register_tool_route("echo", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"echo"}) assert manager._get_mcp_server_from_tool_name("beta-echo") is None assert manager._get_mcp_server_from_tool_name("alpha-echo") is alpha @@ -3970,7 +4158,7 @@ def test_resolve_tool_route_is_not_found_when_no_owner_is_in_scope(self): alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-zulu": zulu} - manager._register_tool_route("secret_tool", "id-zulu") + manager._replace_server_tool_routes("id-zulu", {"secret_tool"}) route = manager.resolve_tool_route("secret_tool", allowed_server_ids=frozenset({"id-alpha"})) @@ -3982,8 +4170,8 @@ def test_resolve_tool_route_names_every_ambiguous_owner(self): alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-zulu": zulu} - manager._register_tool_route("echo", "id-alpha") - manager._register_tool_route("echo", "id-zulu") + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-zulu", {"echo"}) route = manager.resolve_tool_route("echo") @@ -3994,7 +4182,7 @@ def test_resolve_tool_route_resolves_sole_owner(self): manager = MCPServerManager() alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha} - manager._register_tool_route("echo", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"echo"}) route = manager.resolve_tool_route("echo") @@ -4012,8 +4200,8 @@ def test_cleanup_withdraws_only_departing_server_from_shared_tool_name(self): alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-zulu": zulu} - manager._register_tool_route("echo", "id-alpha") - manager._register_tool_route("echo", "id-zulu") + manager._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-zulu", {"echo"}) manager._cleanup_server_tool_routing_artifacts(zulu) del manager.registry["id-zulu"] @@ -4032,9 +4220,8 @@ def test_cleanup_mutates_the_mapping_in_place_without_emptying_it(self): alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha, "id-zulu": zulu} - manager._register_tool_route("shared", "id-alpha") - manager._register_tool_route("shared", "id-zulu") - manager._register_tool_route("alpha_only", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"shared", "alpha_only"}) + manager._replace_server_tool_routes("id-zulu", {"shared"}) original_map = manager.tool_name_to_mcp_server_ids_mapping manager._cleanup_server_tool_routing_artifacts(zulu) @@ -4047,7 +4234,7 @@ def test_cleanup_drops_the_route_when_its_last_owner_leaves(self): manager = MCPServerManager() alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) manager.registry = {"id-alpha": alpha} - manager._register_tool_route("echo", "id-alpha") + manager._replace_server_tool_routes("id-alpha", {"echo"}) manager._cleanup_server_tool_routing_artifacts(alpha) @@ -4422,6 +4609,53 @@ def test_create_prefixed_tools_updates_mapping_for_both_forms(self): assert manager.tool_name_to_mcp_server_ids_mapping["close_issue"] == frozenset({"jira"}) assert manager.tool_name_to_mcp_server_ids_mapping["jira-close_issue"] == frozenset({"jira"}) + def test_relisting_withdraws_routes_for_a_tool_the_upstream_removed(self): + """A tools/list that no longer reports a tool must withdraw that route. + + Registration is authoritative per server, so the second listing is the truth. + With union-only registration "close_issue" would keep resolving to jira after + the upstream dropped it, and a name a second server still served would keep + looking ambiguous (409) until the proxy restarted. + """ + manager = MCPServerManager() + server = MCPServer(server_id="jira", name="jira", transport=MCPTransport.http) + + def _tool(name: str) -> MCPTool: + return MCPTool(name=name, description="", inputSchema={}) + + manager._create_prefixed_tools([_tool("create_issue"), _tool("close_issue")], server) + assert manager.tool_name_to_mcp_server_ids_mapping["close_issue"] == frozenset({"jira"}) + + # Upstream drops close_issue. + manager._create_prefixed_tools([_tool("create_issue")], server) + + assert "close_issue" not in manager.tool_name_to_mcp_server_ids_mapping + assert "jira-close_issue" not in manager.tool_name_to_mcp_server_ids_mapping + assert manager.tool_name_to_mcp_server_ids_mapping["create_issue"] == frozenset({"jira"}) + assert manager.tool_name_to_mcp_server_ids_mapping["jira-create_issue"] == frozenset({"jira"}) + + def test_relisting_clears_a_false_ambiguity_once_one_owner_drops_the_tool(self): + """The reviewed edge case end to end: the 409 must clear without a restart.""" + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) + zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha, "id-zulu": zulu} + scope = frozenset({"id-alpha", "id-zulu"}) + + def _tool(name: str) -> MCPTool: + return MCPTool(name=name, description="", inputSchema={}) + + manager._create_prefixed_tools([_tool("echo")], alpha) + manager._create_prefixed_tools([_tool("echo")], zulu) + assert manager.resolve_tool_route("echo", allowed_server_ids=scope).kind == "ambiguous" + + # alpha's upstream stops exposing echo, so zulu is now the sole owner. + manager._create_prefixed_tools([], alpha) + + route = manager.resolve_tool_route("echo", allowed_server_ids=scope) + assert route.kind == "resolved" + assert route.server is zulu + def test_get_mcp_server_from_tool_name_with_prefixed_and_unprefixed(self): """After mapping is populated, manager resolves both prefixed and unprefixed tool names to the same server.""" manager = MCPServerManager() From e9fd8fccdf2fb8ae064ef6b2b28aa44bc216c2fa Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 25 Jul 2026 17:17:59 -0700 Subject: [PATCH 8/8] fix(mcp): resolve responses-API tool dispatch by server_id and gate it on scope The responses API surface still routed tool calls by server name. It built the caller's reachable set as MCPServer objects, narrowed by both the key's grants and the requested server filter, then discarded that identity by flattening to display names. `tool_server_map` carried a name, and dispatch re-resolved it with `get_mcp_server_by_name`, which walks the whole registry and returns the first match. Server names are not unique, so a tool listed from a reachable server could dispatch to a same-named server the caller cannot reach, sending that server's upstream credential. That is the bug this PR already fixed for MCP JSON-RPC, on the one surface that had not been converted. The fix is the same: stop discarding identity. `tool_server_map` now carries the server_id resolved within the caller's reachable set through `resolve_tool_route`, the same scoped resolver the JSON-RPC path uses, so a name two reachable servers share stays ambiguous here too rather than silently picking one. Dispatch looks the server up by id. A tool with no reachable owner fails closed and reports a result for its tool call, matching how every other failure in that loop is surfaced, rather than being dropped. `resolved_server` is a parameter this PR introduced, and it let a caller hand `call_tool` any server at all. That is the same class of defect one layer down, so the check belongs at the chokepoint rather than at each caller: `call_tool` now takes the reachable set the server was resolved against and refuses to dispatch outside it, covering the caller's server and one it resolves by name itself, which also walks the whole registry. Supplying `resolved_server` without that set is rejected, so caller-supplied identity always arrives with its provenance. Both callers already computed the set, so nothing recomputes it. --- .../mcp_server/mcp_server_manager.py | 22 +++ .../proxy/_experimental/mcp_server/server.py | 3 + .../mcp/litellm_proxy_mcp_handler.py | 67 ++++++--- .../mcp_server/test_mcp_server_manager.py | 57 ++++++++ .../mcp/test_chat_completions_handler.py | 8 + .../mcp/test_litellm_proxy_mcp_handler.py | 138 ++++++++++++++---- .../mcp/test_mcp_streaming_iterator.py | 9 ++ 7 files changed, 252 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7bcccd1841da..91a8c0d473be 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -4686,6 +4686,7 @@ async def call_tool( raw_headers: Optional[dict[str, str]] = None, host_progress_callback: Optional[Callable] = None, resolved_server: MCPServer | None = None, + allowed_server_ids: frozenset[str] | None = None, ) -> CallToolResult: """ Call a tool with the given name and arguments @@ -4702,17 +4703,38 @@ async def call_tool( dispatched to verbatim, so identity never round-trips through the non-unique ``server_name``. Callers without one fall back to resolving by name. + allowed_server_ids: The reachable set ``resolved_server`` was resolved + against. Every dispatch is checked against it, both the caller's + server and one this resolves by name, because ``server_name`` is not + unique and resolving it walks the whole registry. Without it a caller + that resolved scope-blind, or a name collision between a reachable + server and an unreachable one, dispatches the wrong upstream + credential. Required whenever ``resolved_server`` is given, so + caller-supplied identity always arrives with its provenance. Returns: CallToolResult from the MCP server """ start_time = datetime.datetime.now() + if resolved_server is not None and allowed_server_ids is None: + raise ValueError("call_tool requires allowed_server_ids alongside resolved_server") mcp_server = ( resolved_server if resolved_server is not None else self._resolve_mcp_server_for_tool_call(server_name, name) ) + if allowed_server_ids is not None and mcp_server.server_id not in allowed_server_ids: + raise HTTPException( + status_code=403, + detail={ + "error": "server_out_of_scope", + "message": ( + f"Tool '{name}' resolved to MCP server '{mcp_server.name}', which is not among " + f"the servers this caller may reach." + ), + }, + ) # Resolved before any hook runs so a missing BYOK credential (401) never # leaves during-hook side effects (audit logging, rate-limit bookkeeping) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 23e853a32d3d..22f8369895a1 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2841,6 +2841,7 @@ async def execute_mcp_tool( litellm_logging_obj=litellm_logging_obj, host_progress_callback=host_progress_callback, resolved_server=mcp_server, + allowed_server_ids=frozenset(server.server_id for server in allowed_mcp_servers), ) # Fall back to local tool registry with original name (legacy support) @@ -3154,6 +3155,7 @@ async def _handle_managed_mcp_tool( litellm_logging_obj: Optional[Any] = None, host_progress_callback: Optional[Callable] = None, resolved_server: MCPServer | None = None, + allowed_server_ids: frozenset[str] | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" # Import here to avoid circular import @@ -3171,6 +3173,7 @@ async def _handle_managed_mcp_tool( proxy_logging_obj=proxy_logging_obj, host_progress_callback=host_progress_callback, resolved_server=resolved_server, + allowed_server_ids=allowed_server_ids, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) return call_tool_result diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index cae0a38c6e89..2e3e63bc451d 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -192,7 +192,7 @@ async def _get_mcp_tools_from_manager( Returns: List of MCP tools - List names of allowed MCP servers + List of server_ids for the MCP servers this caller may reach """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -282,21 +282,18 @@ async def _get_mcp_tools_from_manager( allowed_mcp_servers=allowed_mcp_servers, ) - server_names: List[str] = [] - for server in allowed_mcp_servers: - if server is None: - continue - server_name = ( - getattr(server, "server_name", None) or getattr(server, "alias", None) or getattr(server, "name", None) - ) - if isinstance(server_name, str): - server_names.append(server_name) + # Carry server_id, not a display name. Names are not unique, so a name handed + # to the dispatch site gets re-resolved against the whole registry and can + # land on a server this caller cannot reach. + server_ids: List[str] = [ + server.server_id for server in allowed_mcp_servers if server is not None and server.server_id + ] - return tools, server_names + return tools, server_ids @staticmethod def _deduplicate_mcp_tools( - mcp_tools: List[MCPTool], allowed_mcp_servers: List[str] + mcp_tools: List[MCPTool], allowed_server_ids_list: List[str] ) -> tuple[List[MCPTool], dict[str, str]]: """ Deduplicate MCP tools by name, keeping the first occurrence of each tool. @@ -306,11 +303,16 @@ def _deduplicate_mcp_tools( Returns: List of deduplicated MCP tools - The returned dictionary maps each tool_name to the server_name + The returned dictionary maps each tool_name to the owning server_id """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + seen_names = set() deduplicated_tools = [] tool_server_map: dict[str, str] = {} + allowed_server_ids = frozenset(allowed_server_ids_list) for tool in mcp_tools: if isinstance(tool, dict): @@ -321,10 +323,16 @@ def _deduplicate_mcp_tools( if tool_name and tool_name not in seen_names: seen_names.add(tool_name) deduplicated_tools.append(tool) - if len(allowed_mcp_servers) == 1: - tool_server_map[tool_name] = allowed_mcp_servers[0] - else: - _, tool_server_map[tool_name] = split_server_prefix_from_name(tool_name) + if len(allowed_server_ids_list) == 1: + tool_server_map[tool_name] = allowed_server_ids_list[0] + continue + # Same scoped resolver the MCP JSON-RPC surface uses, so a name two + # reachable servers share stays ambiguous here too instead of + # silently picking one. Unresolvable names are left out of the map + # and fail closed at dispatch. + route = global_mcp_server_manager.resolve_tool_route(tool_name, allowed_server_ids=allowed_server_ids) + if route.kind == "resolved": + tool_server_map[tool_name] = route.server.server_id return deduplicated_tools, tool_server_map @@ -658,11 +666,25 @@ async def _execute_tool_calls( # Import here to avoid circular import from litellm.proxy.proxy_server import proxy_logging_obj - server_name = tool_server_map[tool_name] - - mcp_server = global_mcp_server_manager.get_mcp_server_by_name( - server_name - ) or global_mcp_server_manager._get_mcp_server_from_tool_name(tool_name) + # tool_server_map carries the server_id resolved within this caller's + # reachable set at listing time, so look it up by identity. A name + # lookup here would walk the whole registry and could land on a + # same-named server the caller cannot reach. + server_id = tool_server_map.get(tool_name) + mcp_server = global_mcp_server_manager.get_mcp_server_by_id(server_id) if server_id else None + if mcp_server is None: + # Fail closed, and report it the way every other failure in this + # loop does so the model still sees a result for its tool call. + verbose_logger.warning(f"No reachable MCP server owns tool {tool_name}") + tool_results.append( + { + "tool_call_id": tool_call_id, + "result": f"Tool call failed: no MCP server reachable by this caller serves '{tool_name}'.", + "name": tool_name, + } + ) + continue + server_name = mcp_server.name resolved_tool_name = ( _resolve_display_name_to_original(tool_name, [mcp_server]) if mcp_server else tool_name ) @@ -781,6 +803,7 @@ async def _execute_tool_calls( raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, resolved_server=mcp_server, + allowed_server_ids=frozenset(tool_server_map.values()), ) if litellm_logging_obj: 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 f9de88610d7d..ba16b70fff05 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 @@ -5121,6 +5121,63 @@ async def test_allowed_tools_with_mixed_prefixed_and_unprefixed_names(self): assert "Tool deletepet is not allowed for server my_api_mcp" in exc_info.value.detail["error"] assert "Contact proxy admin to allow this tool" in exc_info.value.detail["error"] + @pytest.mark.asyncio + async def test_call_tool_refuses_a_resolved_server_outside_the_callers_scope(self): + """Dispatch is gated at the chokepoint, so a caller cannot hand over any server. + + ``resolved_server`` exists so identity does not round-trip through the + non-unique ``server_name``, but a caller that resolved it scope-blind would + otherwise send an unreachable server's upstream credential. + """ + manager = MCPServerManager() + zulu = MCPServer(server_id="id-zulu", name="echo_zulu", transport=MCPTransport.http) + manager.registry = {"id-zulu": zulu} + + with pytest.raises(HTTPException) as excinfo: + await manager.call_tool( + server_name="echo_zulu", + name="echo", + arguments={}, + resolved_server=zulu, + allowed_server_ids=frozenset({"id-alpha"}), + ) + + assert excinfo.value.status_code == 403 + assert excinfo.value.detail["error"] == "server_out_of_scope" + + @pytest.mark.asyncio + async def test_call_tool_refuses_a_name_resolved_server_outside_the_callers_scope(self): + """The gate covers name resolution too, which walks the whole registry.""" + manager = MCPServerManager() + zulu = MCPServer(server_id="id-zulu", name="echo_zulu", server_name="echo_zulu", transport=MCPTransport.http) + manager.registry = {"id-zulu": zulu} + manager._replace_server_tool_routes("id-zulu", {"echo", "echo_zulu-echo"}) + + with pytest.raises(HTTPException) as excinfo: + await manager.call_tool( + server_name="echo_zulu", + name="echo", + arguments={}, + allowed_server_ids=frozenset({"id-alpha"}), + ) + + assert excinfo.value.status_code == 403 + + @pytest.mark.asyncio + async def test_call_tool_requires_the_scope_a_resolved_server_was_resolved_against(self): + """Caller-supplied identity must arrive with its provenance, or it is unverifiable.""" + manager = MCPServerManager() + alpha = MCPServer(server_id="id-alpha", name="echo_alpha", transport=MCPTransport.http) + manager.registry = {"id-alpha": alpha} + + with pytest.raises(ValueError, match="allowed_server_ids"): + await manager.call_tool( + server_name="echo_alpha", + name="echo", + arguments={}, + resolved_server=alpha, + ) + @pytest.mark.asyncio async def test_call_tool_without_broken_pipe_error(self): """ diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py index ab4c5185057e..3cef036214fd 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -1069,6 +1069,14 @@ async def mock_call_tool(**kwargs): "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager.call_tool", mock_call_tool, ) + # tool_server_map carries server_ids, so dispatch resolves the server by identity. + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServer + from litellm.types.mcp import MCPTransport + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager.get_mcp_server_by_id", + lambda sid: MCPServer(server_id=sid, name=sid, server_name=sid, transport=MCPTransport.http), + ) # Create test data tool_calls = [ diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index a1347aa111c0..44c1883343de 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -27,11 +27,28 @@ def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: proxy_module = types.SimpleNamespace(proxy_logging_obj=object()) monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) + # tool_server_map carries server_ids, so dispatch resolves by identity. A real + # MCPServer is used rather than a stub so the prefix-stripping and display-name + # paths downstream see the attributes they actually read. + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServer + from litellm.types.mcp import MCPTransport + + def _server_by_id(server_id): + if server_id is None: + return None + return MCPServer( + server_id=server_id, + name=server_id, + server_name=server_id, + transport=MCPTransport.http, + ) + fake_manager = types.SimpleNamespace( call_tool=AsyncMock(return_value=_DummyMCPResult()), # Newer logging path calls this to enrich spend logs metadata _get_mcp_server_from_tool_name=MagicMock(return_value=None), get_mcp_server_by_name=MagicMock(return_value=None), + get_mcp_server_by_id=MagicMock(side_effect=_server_by_id), ) monkeypatch.setattr( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", @@ -40,6 +57,21 @@ def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: return fake_manager.call_tool +def _mcp_server(server_id: str, server_name: str, alias=None, tool_name_to_display_name=None): + """A real MCPServer, so tests exercise the attributes production actually reads.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServer + from litellm.types.mcp import MCPTransport + + return MCPServer( + server_id=server_id, + name=server_name, + server_name=server_name, + alias=alias, + transport=MCPTransport.http, + tool_name_to_display_name=tool_name_to_display_name, + ) + + def _setup_proxy_logging(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: """Patch proxy_logging_obj so failure hook can be asserted.""" proxy_logging_obj = MagicMock() @@ -61,22 +93,75 @@ def test_deduplicate_mcp_tools_single_allowed_server(): assert server_map == {"search": "everything"} -@pytest.mark.parametrize( - "tool_name,expected_server", - [ - ("alpha-tool", "alpha"), - ("beta-another_tool", "beta"), - ], -) -def test_deduplicate_mcp_tools_prefixed_names(tool_name, expected_server): - tools = [{"name": tool_name}] +def _manager_with(monkeypatch, servers, routes): + """A real manager holding `servers`, with `routes` as {server_id: {tool names}}.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + manager = MCPServerManager() + manager.registry = {server.server_id: server for server in servers} + for server_id, names in routes.items(): + manager._replace_server_tool_routes(server_id, names) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + manager, + ) + return manager + + +def test_deduplicate_mcp_tools_records_the_owning_server_id(monkeypatch): + """A multi-server listing records server_ids, not names, since names are not unique.""" + _manager_with( + monkeypatch, + [_mcp_server(server_id="id-alpha", server_name="alpha"), _mcp_server(server_id="id-beta", server_name="beta")], + {"id-alpha": {"alpha-tool"}, "id-beta": {"beta-another_tool"}}, + ) _, server_map = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools( - tools, - ["alpha", "beta"], + [{"name": "alpha-tool"}, {"name": "beta-another_tool"}], + ["id-alpha", "id-beta"], ) - assert server_map[tool_name] == expected_server + assert server_map == {"alpha-tool": "id-alpha", "beta-another_tool": "id-beta"} + + +def test_deduplicate_mcp_tools_ignores_a_same_named_server_out_of_scope(monkeypatch): + """The reviewed case: a reachable server sharing its name with an unreachable one. + + Recording the name would let dispatch re-resolve it against the whole registry and + land on the unreachable server, sending that server's upstream credential. + """ + reachable = _mcp_server(server_id="id-reachable", server_name="shared") + unreachable = _mcp_server(server_id="id-unreachable", server_name="shared") + _manager_with( + monkeypatch, + [reachable, unreachable], + {"id-reachable": {"echo"}, "id-unreachable": {"echo"}}, + ) + + _, server_map = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools( + [{"name": "echo"}], + ["id-reachable"], + ) + + assert server_map == {"echo": "id-reachable"} + + +def test_deduplicate_mcp_tools_omits_a_name_two_reachable_servers_share(monkeypatch): + """Ambiguity within the caller's own scope must not silently pick one server.""" + _manager_with( + monkeypatch, + [_mcp_server(server_id="id-alpha", server_name="alpha"), _mcp_server(server_id="id-beta", server_name="beta")], + {"id-alpha": {"echo"}, "id-beta": {"echo"}}, + ) + + _, server_map = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools( + [{"name": "echo"}], + ["id-alpha", "id-beta"], + ) + + assert "echo" not in server_map def test_extract_tool_calls_from_chat_response_handles_tool_calls(): @@ -288,19 +373,14 @@ async def test_execute_tool_calls_strips_prefix_when_alias_differs_from_server_n monkeypatch, ): call_tool_mock = _setup_mcp_call_environment(monkeypatch) - fake_server = types.SimpleNamespace( + fake_server = _mcp_server( alias="my_deepwiki", server_name="deepwiki_test", server_id="test-server-id", - short_prefix=None, - mcp_info=None, - tool_name_to_display_name=None, ) from litellm.proxy._experimental.mcp_server import mcp_server_manager as _msm - _msm.global_mcp_server_manager._get_mcp_server_from_tool_name = MagicMock( - return_value=fake_server - ) + _msm.global_mcp_server_manager.get_mcp_server_by_id = MagicMock(return_value=fake_server) tool_name = "my_deepwiki-read_wiki_structure" tool_calls = [ @@ -311,7 +391,7 @@ async def test_execute_tool_calls_strips_prefix_when_alias_differs_from_server_n ] await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( - tool_server_map={tool_name: "deepwiki_test"}, + tool_server_map={tool_name: "test-server-id"}, tool_calls=tool_calls, user_api_key_auth=None, ) @@ -324,26 +404,24 @@ async def test_execute_tool_calls_strips_prefix_when_alias_differs_from_server_n @pytest.mark.asyncio async def test_execute_tool_calls_reverse_maps_display_name(monkeypatch): call_tool_mock = _setup_mcp_call_environment(monkeypatch) - colliding_server = types.SimpleNamespace( - alias=None, + colliding_server = _mcp_server( server_name="other_mcp", server_id="other-server-id", - short_prefix=None, - mcp_info=None, tool_name_to_display_name={"search": "search_docs"}, ) - fake_server = types.SimpleNamespace( - alias=None, + fake_server = _mcp_server( server_name="deepwiki_mcp", server_id="test-server-id", - short_prefix=None, - mcp_info=None, tool_name_to_display_name={"read_wiki_structure": "browse_repo_docs"}, ) from litellm.proxy._experimental.mcp_server import mcp_server_manager as _msm + # Dispatch resolves the server_id recorded at listing time. A name lookup would + # be free to return the collider instead, which is the bug this guards. + by_id = {"test-server-id": fake_server, "other-server-id": colliding_server} + _msm.global_mcp_server_manager.get_mcp_server_by_id = MagicMock(side_effect=by_id.get) _msm.global_mcp_server_manager._get_mcp_server_from_tool_name = MagicMock(return_value=colliding_server) - _msm.global_mcp_server_manager.get_mcp_server_by_name = MagicMock(return_value=fake_server) + _msm.global_mcp_server_manager.get_mcp_server_by_name = MagicMock(return_value=colliding_server) tool_name = "browse_repo_docs" tool_calls = [ @@ -354,7 +432,7 @@ async def test_execute_tool_calls_reverse_maps_display_name(monkeypatch): ] await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( - tool_server_map={tool_name: "deepwiki_mcp"}, + tool_server_map={tool_name: "test-server-id"}, tool_calls=tool_calls, user_api_key_auth=None, ) diff --git a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py index 24edf12fffe9..38e58dc14405 100644 --- a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py @@ -68,10 +68,19 @@ def _text_only_stream(text: str, response_id: str = "resp-1") -> _FakeAsyncStrea def _mock_mcp_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: """Patch the MCP tool-call plumbing so _execute_tool_calls can run in tests.""" call_tool = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)) + # tool_server_map carries server_ids now, so dispatch resolves by identity. + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServer + from litellm.types.mcp import MCPTransport + fake_manager = types.SimpleNamespace( call_tool=call_tool, _get_mcp_server_from_tool_name=MagicMock(return_value=None), get_mcp_server_by_name=MagicMock(return_value=None), + get_mcp_server_by_id=MagicMock( + side_effect=lambda sid: ( + MCPServer(server_id=sid, name=sid, server_name=sid, transport=MCPTransport.http) if sid else None + ) + ), ) monkeypatch.setattr( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",