From 85eaa03679570a034ccc87a26872eae40c594256 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 17 Jul 2026 03:52:59 -0700 Subject: [PATCH] fix(mcp): discriminate tool routes on a tag so reloads cannot skip the ambiguity check server.py selected the route variant with isinstance. A module reload rebinds the manager and its classes, so resolve_tool_route returned a reloaded variant while server.py still held the pre-reload class; isinstance was False, the ambiguity branch never ran, and the call fell through to the local registry as a 404 instead of the intended 409. Each variant now carries a Literal kind tag and server.py dispatches on it, which is what makes this a tagged union rather than one that leans on class identity. Withdrawing a server's routes also rebuilt the mapping into a fresh dict, orphaning any holder of the previous one. The initialize task is fired without being awaited and can be holding it, so the rebuild now updates the mapping in place. --- .../proxy/_experimental/mcp_server/mcp_server_manager.py | 7 ++++++- litellm/proxy/_experimental/mcp_server/server.py | 6 ++---- .../_experimental/mcp_server/test_mcp_server_manager.py | 6 +++--- 3 files changed, 11 insertions(+), 8 deletions(-) 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."""