diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 94271c54f4b..b122f7871a9 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/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) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index d4ba66c4381..ecad9ed75e3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/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."""