diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index ed6dde23d9e3..91a8c0d473be 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -14,7 +14,8 @@ import re import time from contextlib import asynccontextmanager -from typing import Any, AsyncIterator, Callable, Literal, Optional, Union, cast +from dataclasses import dataclass +from typing import Any, AsyncIterator, Callable, Collection, Literal, Optional, Union, cast from urllib.parse import urlparse import anyio @@ -194,6 +195,34 @@ 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 + kind: Literal["resolved"] = "resolved" + + +@dataclass(frozen=True, slots=True) +class MCPToolRouteNotFound: + """No MCP server serves the tool name.""" + + tool_name: str + kind: Literal["not_found"] = "not_found" + + +@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] + kind: Literal["ambiguous"] = "ambiguous" + + +MCPToolRoute = Union[MCPToolRouteResolved, MCPToolRouteNotFound, MCPToolRouteAmbiguous] + + def _blank_to_none(value: str | None) -> str | None: """Collapse an absent, empty, or whitespace-only string to ``None``. @@ -1008,7 +1037,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", @@ -1350,7 +1379,7 @@ async def load_servers_from_config( await self._hydrate_config_servers_dcr_clients() - self.initialize_tool_name_to_mcp_server_name_mapping() + self.initialize_tool_name_to_mcp_server_ids_mapping() async def _hydrate_config_servers_dcr_clients(self) -> None: """Overlay each config-declared server's persisted DCR client (from the server-scoped @@ -1453,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") @@ -1467,7 +1497,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 @@ -1495,13 +1525,18 @@ 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 + # 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: @@ -1513,9 +1548,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, @@ -1526,24 +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) - 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._replace_server_tool_routes(server.server_id, frozenset()) def remove_server(self, mcp_server: LiteLLM_MCPServerTable): """ @@ -1951,7 +1972,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 @@ -3801,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) @@ -3814,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.tool_name_to_mcp_server_name_mapping[original_name] = prefix - 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 + 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 @@ -4545,8 +4573,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") @@ -4657,6 +4685,8 @@ 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, + allowed_server_ids: frozenset[str] | None = None, ) -> CallToolResult: """ Call a tool with the given name and arguments @@ -4669,13 +4699,42 @@ 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. + 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() - mcp_server = self._resolve_mcp_server_for_tool_call(server_name, name) + 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) @@ -4775,29 +4834,35 @@ 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 + 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. @@ -4811,12 +4876,81 @@ async def _initialize_tool_name_to_mcp_server_name_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.tool_name_to_mcp_server_name_mapping[original_name] = server.name - self.tool_name_to_mcp_server_name_mapping[tool.name] = server.name + + 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. + """ + 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, + 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) + # 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) + return MCPToolRouteResolved(server=server) def _get_mcp_server_from_tool_name(self, tool_name: str) -> Optional[MCPServer]: """ @@ -4826,45 +4960,43 @@ 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()) - - # 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 - 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 - - # 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_name_mapping - or tool_name in self.tool_name_to_mcp_server_name_mapping - ): - return matched_server + # 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: + 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.""" @@ -4963,7 +5095,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 a8ab0937124c..22f8369895a1 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2611,7 +2611,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 route.kind == "ambiguous": + 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 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: for known_prefix in iter_known_server_prefixes(requested_server): candidate = global_mcp_server_manager._get_mcp_server_from_tool_name( @@ -2664,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( @@ -2809,6 +2840,8 @@ 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, + allowed_server_ids=frozenset(server.server_id for server in allowed_mcp_servers), ) # Fall back to local tool registry with original name (legacy support) @@ -3121,6 +3154,8 @@ 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, + allowed_server_ids: frozenset[str] | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" # Import here to avoid circular import @@ -3137,6 +3172,8 @@ 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, + allowed_server_ids=allowed_server_ids, ) 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 392bb7bcab27..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 ) @@ -780,6 +802,8 @@ async def _execute_tool_calls( oauth2_headers=oauth2_headers, 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/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index 55b49aa0d295..ba53b0777b12 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,20 @@ 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 = 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_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 +240,24 @@ 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 = 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" - 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 +274,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 +407,18 @@ 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 = 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 ) - 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 +435,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 f1e6539439ac..7c172ee188a3 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", @@ -220,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_name_mapping[ + test_manager.tool_name_to_mcp_server_ids_mapping[ f"{expected_prefix}-gmail_send_email" ] - == expected_prefix + == expected_owner ) 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 + == expected_owner ) @@ -275,8 +284,10 @@ 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_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_http_server.server_id} ) # Call the tool @@ -341,8 +352,10 @@ 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_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_http_server.server_id} ) # Call the tool with invalid data @@ -383,8 +396,10 @@ 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_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_http_server.server_id} ) # Try to call a tool that doesn't exist in mapping @@ -1395,12 +1410,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 +1470,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 +1530,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 ae4f12fc1e17..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 @@ -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 @@ -5962,8 +5962,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, @@ -6049,8 +6049,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, @@ -6092,6 +6092,374 @@ 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_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. + + 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 55b6bbbdbc29..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 @@ -3012,7 +3012,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 = { @@ -3029,8 +3029,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): @@ -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""" @@ -3876,8 +4029,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 @@ -3892,11 +4045,201 @@ 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_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._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_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._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._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 + + 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._replace_server_tool_routes("id-alpha", {"echo"}) + + 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 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", 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._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 + + 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._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 + + 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._replace_server_tool_routes("id-zulu", {"secret_tool"}) + + 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() + 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._replace_server_tool_routes("id-alpha", {"echo"}) + manager._replace_server_tool_routes("id-zulu", {"echo"}) + + route = manager.resolve_tool_route("echo") + + assert route.kind == "ambiguous" + 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._replace_server_tool_routes("id-alpha", {"echo"}) + + route = manager.resolve_tool_route("echo") + + assert route.kind == "resolved" + assert route.server is alpha + + def test_resolve_tool_route_reports_unknown_tool(self): + manager = MCPServerManager() + + 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.""" + 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._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"] + + 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._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) + + 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) + manager.registry = {"id-alpha": alpha} + manager._replace_server_tool_routes("id-alpha", {"echo"}) + + 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 @@ -3923,7 +4266,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") @@ -3950,8 +4293,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") @@ -4261,10 +4604,57 @@ 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_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.""" @@ -4731,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): """ @@ -4753,8 +5200,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() @@ -5205,8 +5652,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 @@ -7257,15 +7704,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 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",