fix: fetch MCP REST tools concurrently

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -829,21 +829,20 @@ if MCP_AVAILABLE:
                     else {}
                 )
 
-                # Query all servers the user has access to
-                errors = []
-                for allowed_server_id in allowed_server_ids:
-                    server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id)
-                    if server is None:
-                        continue
-
-                    server_auth_header = _get_server_auth_header(server, mcp_server_auth_headers, mcp_auth_header)
-                    user_oauth_extra_headers = await _get_user_oauth_extra_headers(
-                        server,
-                        user_api_key_dict,
-                        prefetched_creds=prefetched_oauth_creds,
-                    )
+                servers = tuple(
+                    server
+                    for allowed_server_id in allowed_server_ids
+                    if (server := global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id)) is not None
+                )
 
+                async def _fetch_server_tools(server):
                     try:
+                        server_auth_header = _get_server_auth_header(server, mcp_server_auth_headers, mcp_auth_header)
+                        user_oauth_extra_headers = await _get_user_oauth_extra_headers(
+                            server,
+                            user_api_key_dict,
+                            prefetched_creds=prefetched_oauth_creds,
+                        )
                         tools_result = await _get_tools_for_single_server(
                             server,
                             server_auth_header,
@@ -852,15 +851,19 @@ if MCP_AVAILABLE:
                             extra_headers=user_oauth_extra_headers,
                             apply_tool_filters=apply_tool_filters,
                         )
-                        list_tools_result.extend(tools_result)
+                        return tools_result, None
                     except Exception as e:
                         verbose_logger.exception(f"Error getting tools from {server.name}: {e}")
-                        errors.append(
+                        error = (
                             f"{get_server_prefix(server)}: {classify_list_exception(e).tag}"
                             if isinstance(e, (MCPServerListError, MCPUpstreamAuthError))
                             else f"{get_server_prefix(server)}: {str(e)}"
                         )
-                        continue
+                        return [], error
+
+                results = await asyncio.gather(*(_fetch_server_tools(server) for server in servers))
+                list_tools_result = [tool for tools_result, _ in results for tool in tools_result]
+                errors = [error for _, error in results if error is not None]
 
                 if errors and not list_tools_result:
                     error_message = "Failed to get tools from servers: " + "; ".join(errors)

tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py

@@ -1039,6 +1039,66 @@ class TestListToolsRestAPI:
         assert result["tools"] == ["good-tool"]
         assert result["error"] is None
 
+    async def test_aggregate_list_fetches_servers_concurrently(self, monkeypatch):
+        class StubServer:
+            def __init__(self, name):
+                self.alias = name
+                self.server_name = name
+                self.name = name
+                self.allowed_tools = None
+                self.mcp_info = {"server_name": name}
+                self.available_on_public_internet = True
+
+        first = StubServer("first")
+        second = StubServer("second")
+        both_started = asyncio.Event()
+        release = asyncio.Event()
+        started = 0
+
+        async def fake_contexts(user_api_key_auth):
+            return [user_api_key_auth]
+
+        async def fake_get_allowed_mcp_servers(*args, **kwargs):
+            return ["first", "second"]
+
+        async def fake_get_tools(server, *args, **kwargs):
+            nonlocal started
+            started += 1
+            if started == 2:
+                both_started.set()
+            await release.wait()
+            return [f"{server.name}-tool"]
+
+        monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False)
+        monkeypatch.setattr(
+            rest_endpoints.global_mcp_server_manager,
+            "get_allowed_mcp_servers",
+            fake_get_allowed_mcp_servers,
+            raising=False,
+        )
+        monkeypatch.setattr(
+            rest_endpoints.global_mcp_server_manager,
+            "get_mcp_server_by_id",
+            lambda server_id: {"first": first, "second": second}.get(server_id),
+            raising=False,
+        )
+        monkeypatch.setattr(rest_endpoints, "_get_tools_for_single_server", fake_get_tools, raising=False)
+
+        request = _build_request(path="/mcp-rest/tools/list", method="GET")
+        listing_task = asyncio.create_task(
+            rest_endpoints.list_tool_rest_api(
+                request,
+                server_id=None,
+                user_api_key_dict=UserAPIKeyAuth(),
+            )
+        )
+
+        await asyncio.wait_for(both_started.wait(), timeout=0.1)
+        release.set()
+        result = await listing_task
+
+        assert set(result["tools"]) == {"first-tool", "second-tool"}
+
     async def test_name_resolution_finds_server_by_uuid(self, monkeypatch):
         """When server_id is a name string, it should be resolved to its UUID
         and used for the tools lookup when the UUID is in allowed_server_ids."""