From 80ec93d0f0656c1eef6d3705ebf2dec3bdabe755 Mon Sep 17 00:00:00 2001 From: "Claude (greptile subagent)" Date: Wed, 20 May 2026 20:11:31 +0000 Subject: [PATCH] fix(mcp): always reject unknown tools in server-name fallback Defense-in-depth: _resolve_mcp_server_for_tool_call previously skipped the unknown-tool check whenever the per-server mapping had no entries yet (cold start, OAuth2 lazy listing, or upstream listing failure), allowing arbitrary tool names to reach upstream servers. Tighten the check so the server-name fallback always rejects tool names not present in the mapping. Callers must call list_tools first (standard MCP flow) before tools/call can resolve. Removes the now-unused _mapping_has_tools_for_server helper and adds an explicit empty-mapping rejection test alongside the existing populated-mapping rejection test. Co-authored-by: Sameer Kankute --- .../mcp_server/mcp_server_manager.py | 17 +---------------- .../mcp_server/test_mcp_server_manager.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 4da24eca573..bbf40f6e9ef 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2876,7 +2876,7 @@ class MCPServerManager: 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 and self._mapping_has_tools_for_server(mcp_server): + if not tool_known: raise ValueError(f"Tool {name} not found") return mcp_server @@ -3078,21 +3078,6 @@ class MCPServerManager: 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 _mapping_has_tools_for_server(self, server: MCPServer) -> bool: - """Return True if the tool mapping lists any tool for this server.""" - server_identifiers = { - normalize_server_name(server.name), - } - if server.server_name: - server_identifiers.add(normalize_server_name(server.server_name)) - if server.alias: - server_identifiers.add(normalize_server_name(server.alias)) - - for mapped_server_name in self.tool_name_to_mcp_server_name_mapping.values(): - if normalize_server_name(mapped_server_name) in server_identifiers: - return True - return False - def _get_mcp_server_from_tool_name(self, tool_name: str) -> Optional[MCPServer]: """ Get the MCP Server from the tool name (handles both prefixed and non-prefixed names) 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 eb7878337a1..d7078412a44 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 @@ -1877,12 +1877,30 @@ class TestMCPServerManager: transport=MCPTransport.http, ) manager.registry = {"srv-uuid-123": server} + manager.tool_name_to_mcp_server_name_mapping["create_zap"] = "zapier" 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). + """ + manager = MCPServerManager() + server = MCPServer( + server_id="srv-uuid-123", + name="zapier", + alias="zapier-alias", + transport=MCPTransport.http, + ) + 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") + 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.""" manager = MCPServerManager()