scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(mcp): resolve call_tool by registry without requiring tool map

Browse source 
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
This commit is contained in:
mubashir1osmani 2026-07-28 21:00:12 -07:00
parent 998a372417
commit 8b56e51e39
2 changed files with 38 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -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:

40
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -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):