From ec5f945390ff6cec3260565824e814dfdb2b2534 Mon Sep 17 00:00:00 2001 From: Fernando Cucci Date: Thu, 20 Aug 2026 10:00:15 +0200 Subject: [PATCH] fix(mcp): map only OpenAPI servers when a registry reload registers their tools reload_servers_from_database() re-registers OpenAPI tools on every run and then called initialize_tool_name_to_mcp_server_name_mapping() for the whole registry, which calls list_tools upstream on every remote MCP server. The reload is on a timer (proxy_config_reload_interval_seconds, 30s by default), so any deployment with at least one OpenAPI server polls every other MCP server every 30 seconds. Servers that take a per-user credential at call time (forwarded via the x-mcp--authorization header) are called with no credential at all there and answer 401, so the proxy log fills with 'Error listing tools from : 401' twice a minute and the upstream sees a constant stream of unauthenticated requests. OpenAPI tools are served from the local tool registry, so the mapping refresh those registrations need only covers the OpenAPI servers themselves. Restrict the sweep to them; startup still maps the full registry. --- .../mcp_server/mcp_server_manager.py | 31 +++++++++----- .../mcp_server/test_mcp_server_manager.py | 42 +++++++++++++++++++ 2 files changed, 62 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 26a6f8d1251..153fdbb30e4 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2731,7 +2731,10 @@ class MCPServerManager: base_url=server.url or "", ) if initialize_mapping: - self.initialize_tool_name_to_mcp_server_name_mapping() + # OpenAPI tools are served from the local tool registry, so only this + # server needs re-mapping. Sweeping the whole registry would call + # list_tools upstream on every other (remote) MCP server. + self.initialize_tool_name_to_mcp_server_name_mapping(servers=(server,)) async def add_server(self, mcp_server: LiteLLM_MCPServerTable): # The runtime registry is the allowlist for tool calls and health @@ -5828,24 +5831,28 @@ class MCPServerManager: # End of Methods that call the upstream MCP servers ######################################################### - def initialize_tool_name_to_mcp_server_name_mapping(self): + def initialize_tool_name_to_mcp_server_name_mapping(self, servers: Sequence[MCPServer] | None = None): """ - On startup, initialize the tool name to MCP server name mapping + On startup, initialize the tool name to MCP server name mapping. + + ``servers`` restricts the sweep to those servers. Callers that only registered + OpenAPI tools pass the affected servers so a re-registration does not re-list + tools upstream from every remote server in the registry. """ try: if asyncio.get_running_loop(): - asyncio.create_task(self._initialize_tool_name_to_mcp_server_name_mapping()) + asyncio.create_task(self._initialize_tool_name_to_mcp_server_name_mapping(servers)) except RuntimeError as e: # no running event loop verbose_logger.exception( "No running event loop - skipping tool name to MCP server name mapping initialization: %s", e ) - async def _initialize_tool_name_to_mcp_server_name_mapping(self): + async def _initialize_tool_name_to_mcp_server_name_mapping(self, servers: Sequence[MCPServer] | None = None): """ Call list_tools for each server and update the tool name to MCP server name mapping Note: This now handles prefixed tool names """ - for server in self.get_registry().values(): + for server in servers if servers is not None else tuple(self.get_registry().values()): if self._oauth_discovery_slot(server.server_id) is not None: continue if server.needs_user_oauth_token: @@ -5995,7 +6002,6 @@ class MCPServerManager: # Assign short prefixes against the full candidate set without # publishing the staged registry to concurrent callers. registered_registry: Final[dict[str, MCPServer]] = {} - registered_openapi_tools = False for server_id, new_server in new_registry.items(): try: self._assign_unique_short_prefix(new_server, registry=new_registry) @@ -6004,8 +6010,6 @@ class MCPServerManager: # prefix that lookups will use. await self._maybe_register_openapi_tools(new_server, initialize_mapping=False) registered_registry[server_id] = new_server - if new_server.spec_path: - registered_openapi_tools = True except Exception as e: verbose_logger.exception( "Skipping MCP server %s (%s) during DB reload: %s", @@ -6026,8 +6030,13 @@ class MCPServerManager: registered_servers: Final = tuple(registered_registry.values()) self._reconcile_oauth_discovery_slots_for_servers(registered_servers) self._prime_oauth_metadata_discovery_for_servers(registered_servers) - if registered_openapi_tools: - self.initialize_tool_name_to_mcp_server_name_mapping() + openapi_servers: Final = tuple(server for server in registered_servers if server.spec_path) + if openapi_servers: + # This reload runs on a timer (``proxy_config_reload_interval_seconds``, + # 30s by default). Map only the OpenAPI servers, whose tools come from the + # local registry: a full sweep would list tools upstream from every remote + # MCP server on every reload. + self.initialize_tool_name_to_mcp_server_name_mapping(servers=openapi_servers) verbose_logger.debug("MCP registry refreshed (%s servers in registry)", len(registered_registry)) 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 5ee8143fb8e..5aa057da91c 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 @@ -928,6 +928,48 @@ class TestMCPServerManager: get_tools.assert_not_awaited() + @pytest.mark.asyncio + async def test_database_reload_maps_only_openapi_servers(self): + """The DB reload runs on a timer, so the tool mapping it triggers for freshly + registered OpenAPI servers must not re-list tools upstream from every remote + server in the registry.""" + manager = MCPServerManager() + openapi_row = LiteLLM_MCPServerTable( + server_id="openapi-1", + server_name="openapi_server", + alias="openapi_server", + url="https://openapi.example.com", + transport=MCPTransport.http, + spec_path="https://openapi.example.com/openapi.json", + ) + remote_row = LiteLLM_MCPServerTable( + server_id="remote-1", + server_name="remote_server", + alias="remote_server", + url="https://remote.example.com/mcp", + transport=MCPTransport.http, + ) + repository = MagicMock() + repository.table.find_many = AsyncMock(return_value=[openapi_row, remote_row]) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + return_value=repository, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch.object(manager, "_register_openapi_tools", new=AsyncMock()), + patch.object(manager, "initialize_tool_name_to_mcp_server_name_mapping") as init_mapping, + ): + await manager.reload_servers_from_database() + + init_mapping.assert_called_once() + mapped = init_mapping.call_args.kwargs["servers"] + assert [server.server_id for server in mapped] == ["openapi-1"] + @pytest.mark.asyncio async def test_load_servers_from_config_requires_oauth2_flow(self): """auth_type oauth2 without an explicit oauth2_flow is a config error: the