From 8b56e51e39b876d13d1112efa4130554ddf5f173 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 21:00:12 -0700 Subject: [PATCH] fix(mcp): resolve call_tool by registry without requiring tool map Multi-worker reloads put MCP servers in the registry from the DB but do not re-run tools/list on every process. Gating call_tool on tool_name_to_mcp_server_name_mapping made cold workers 500 with Tool not found after another worker had already listed the tool. Treat a registry match on server id/name/alias as enough; upstream rejects unknown tools --- .../mcp_server/mcp_server_manager.py | 22 +++++----- .../mcp_server/test_mcp_server_manager.py | 40 +++++++++++++------ 2 files changed, 38 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 24a1909aa04..d86847eeebc 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -4892,14 +4892,23 @@ class MCPServerManager: server_name: str, name: str, ) -> MCPServer: - """Resolve MCP server for call_tool (prefixed name, registry, fallback).""" + """Resolve the MCP server that should handle call_tool. + + Prefer the process-local tool→server map when this worker has listed + tools, then fall back to a registry match on server id / name / alias. + A registry hit is enough: do not require the tool map. That map is only + filled after tools/list on *this* process; multi-worker reloads load the + server row from the DB without re-listing, so gating on the map made + cold workers 500 "Tool X not found" for tools another worker already + served. Unknown tool names are rejected by the upstream MCP server. + """ prefixed_tool_name = add_server_prefix_to_name(name, server_name) mcp_server = self._get_mcp_server_from_tool_name(prefixed_tool_name) - resolved_by_server_name_only = False normalized_server_name = normalize_server_name(server_name) def _candidate_matches_server_name(candidate: MCPServer) -> bool: for identifier in ( + candidate.server_id, candidate.alias, candidate.server_name, candidate.name, @@ -4912,7 +4921,6 @@ class MCPServerManager: for candidate in self.get_registry().values(): if _candidate_matches_server_name(candidate): mcp_server = candidate - resolved_by_server_name_only = True break if mcp_server is None: fallback = self._get_mcp_server_from_tool_name(name) @@ -4921,14 +4929,6 @@ class MCPServerManager: if mcp_server is None: raise ValueError(f"Tool {name} not found") - 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 - ) - if not tool_known: - raise ValueError(f"Tool {name} not found") - return mcp_server async def has_user_oauth_token(self, server: MCPServer, user_api_key_auth: Optional[UserAPIKeyAuth]) -> bool: 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 42b6cbee1c4..377a84f229e 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 @@ -4098,10 +4098,13 @@ class TestMCPServerManager: resolved = manager._resolve_mcp_server_for_tool_call("zapier-alias", "create_zap") assert resolved is server - 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 - or the server is OAuth2 and the user has not yet listed tools). + def test_resolve_mcp_server_for_tool_call_empty_mapping_still_resolves_server(self): + """Registry name match is enough when this process never ran tools/list. + + Multi-worker deployments reload MCP servers from the DB without + re-listing tools, so the in-process tool map is empty while the + server row is present. call_tool must still resolve the server; + upstream rejects tools that do not exist. """ manager = MCPServerManager() server = MCPServer( @@ -4112,8 +4115,21 @@ class TestMCPServerManager: ) manager.registry = {"srv-uuid-123": server} - with pytest.raises(ValueError, match="Tool create_zap not found"): - manager._resolve_mcp_server_for_tool_call("zapier-alias", "create_zap") + resolved = manager._resolve_mcp_server_for_tool_call("zapier-alias", "create_zap") + assert resolved is server + + def test_resolve_mcp_server_for_tool_call_by_server_id_without_tool_map(self): + """REST clients pass server_id; registry id match must not need the tool map.""" + manager = MCPServerManager() + server = MCPServer( + server_id="srv-uuid-abc", + name="datadog", + transport=MCPTransport.http, + ) + manager.registry = {"srv-uuid-abc": server} + + resolved = manager._resolve_mcp_server_for_tool_call("srv-uuid-abc", "search_datadog_logs") + assert resolved is server def test_resolve_mcp_server_for_tool_call_fallback_to_unprefixed_lookup(self): """Fallback to unprefixed _get_mcp_server_from_tool_name when other paths fail.""" @@ -4137,11 +4153,10 @@ class TestMCPServerManager: manager._resolve_mcp_server_for_tool_call("nonexistent", "ghost_tool") def test_resolve_mcp_server_for_tool_call_unknown_tool_with_known_server(self): - """Server-name match alone must not let unknown tools slip through. + """A known server still resolves even when the tool is not in the map. - If the registry has tools for this server but neither the prefixed nor - unprefixed tool name is in the mapping, raise rather than returning the - server (would otherwise allow tool enumeration via name spoofing). + The tool map is a routing cache filled by tools/list, not an allow-list. + Unknown tool names are rejected by the upstream MCP server. """ manager = MCPServerManager() server = MCPServer( @@ -4150,12 +4165,11 @@ class TestMCPServerManager: transport=MCPTransport.http, ) 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" - with pytest.raises(ValueError, match="Tool missing_tool not found"): - manager._resolve_mcp_server_for_tool_call("github", "missing_tool") + resolved = manager._resolve_mcp_server_for_tool_call("github", "missing_tool") + assert resolved is server @pytest.mark.asyncio async def test_resolve_oauth2_headers_skipped_when_not_user_oauth(self):