diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b8bf986776e..f107b829a90 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -192,6 +192,7 @@ class MCPToolRouteResolved: """Exactly one MCP server serves the tool name.""" server: MCPServer + kind: Literal["resolved"] = "resolved" @dataclass(frozen=True, slots=True) @@ -199,6 +200,7 @@ class MCPToolRouteNotFound: """No MCP server serves the tool name.""" tool_name: str + kind: Literal["not_found"] = "not_found" @dataclass(frozen=True, slots=True) @@ -207,6 +209,7 @@ class MCPToolRouteAmbiguous: tool_name: str server_ids: frozenset[str] + kind: Literal["ambiguous"] = "ambiguous" MCPToolRoute = Union[MCPToolRouteResolved, MCPToolRouteNotFound, MCPToolRouteAmbiguous] @@ -1550,11 +1553,13 @@ class MCPServerManager: openapi_key_prefix = prefix_root + MCP_TOOL_PREFIX_SEPARATOR global_mcp_tool_registry.unregister_tools_with_prefix(openapi_key_prefix) - self.tool_name_to_mcp_server_ids_mapping = { + surviving = { tool_name: remaining for tool_name, owner_ids in self.tool_name_to_mcp_server_ids_mapping.items() if (remaining := owner_ids - {server.server_id}) } + self.tool_name_to_mcp_server_ids_mapping.clear() + self.tool_name_to_mcp_server_ids_mapping.update(surviving) def remove_server(self, mcp_server: LiteLLM_MCPServerTable): """ diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 0705c392f7f..e18beadd3bf 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -358,8 +358,6 @@ if MCP_AVAILABLE: ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, - MCPToolRouteAmbiguous, - MCPToolRouteResolved, _caller_authorization_fans_out, _client_forwarded_authorization_headers, _should_strip_caller_authorization, @@ -2629,7 +2627,7 @@ if MCP_AVAILABLE: name, allowed_server_ids=frozenset(server.server_id for server in allowed_mcp_servers), ) - if isinstance(route, MCPToolRouteAmbiguous): + if route.kind == "ambiguous": candidates = sorted( server.name for server in ( @@ -2648,7 +2646,7 @@ if MCP_AVAILABLE: ), }, ) - mcp_server = route.server if isinstance(route, MCPToolRouteResolved) else None + mcp_server = route.server if route.kind == "resolved" else None else: mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) if mcp_server is None and requested_server is not None: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 646c7a835da..1116e3efaf1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -3948,7 +3948,7 @@ class TestMCPServerManager: route = manager.resolve_tool_route("echo") - assert type(route).__name__ == "MCPToolRouteAmbiguous" + assert route.kind == "ambiguous" assert route.server_ids == frozenset({"id-alpha", "id-zulu"}) def test_resolve_tool_route_resolves_sole_owner(self): @@ -3959,13 +3959,13 @@ class TestMCPServerManager: route = manager.resolve_tool_route("echo") - assert type(route).__name__ == "MCPToolRouteResolved" + assert route.kind == "resolved" assert route.server is alpha def test_resolve_tool_route_reports_unknown_tool(self): manager = MCPServerManager() - assert type(manager.resolve_tool_route("nothing_serves_this")).__name__ == "MCPToolRouteNotFound" + assert manager.resolve_tool_route("nothing_serves_this").kind == "not_found" def test_cleanup_withdraws_only_departing_server_from_shared_tool_name(self): """Removing one server must leave a co-owned tool name routable to the survivor."""