diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 02799b90566..96cf6ee11aa 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -4531,7 +4531,7 @@ class MCPServerManager: ) else: page: Final = ( - await client.list_tools_page(params) + await self._fetch_tools_page_with_timeout(client, server, params) if params is not None else ListToolsResult(tools=await self._fetch_tools_with_timeout(client, server.name)) ) @@ -5621,6 +5621,19 @@ class MCPServerManager: ) raise_classified_list_failure(e, server_name) + async def _fetch_tools_page_with_timeout( + self, client: MCPClient, server: MCPServer, params: PaginatedRequestParams + ) -> ListToolsResult: + deadline: Final = ( + MCP_TOOL_LISTING_TIMEOUT if server.timeout is None else max(server.timeout, MCP_TOOL_LISTING_TIMEOUT) + ) + try: + with anyio.fail_after(deadline): + return await client.list_tools_page(params) + except TimeoutError as e: + verbose_logger.warning("Timeout while listing tools from %s", server.name) + raise MCPServerListError(ServerListFault(tag="timeout"), server.name) from e + _SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024 def assign_unique_short_prefix( diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 691da35b8c4..d47b2f73a48 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -20056,6 +20056,74 @@ async def test_disabled_stdio_tools_are_empty_initially_and_reject_continuation( assert page.next_cursor is None +@pytest.mark.asyncio +async def test_aggregate_bounds_each_paginated_server_by_its_own_listing_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from mcp.types import JSONRPCRequest, PaginatedRequestParams + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import catalog, operations + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module + from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR, get_server_prefix + + monkeypatch.setenv("LITELLM_SALT_KEY", "paginated-listing-deadline-test") + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(manager_module, "MCP_TOOL_LISTING_TIMEOUT", 0.0) + extended: Final = MCPServer( + server_id="extended", + name="extended", + transport=MCPTransport.http, + url="https://extended.example/mcp", + timeout=60.0, + ) + global_deadline: Final = MCPServer( + server_id="global", name="global", transport=MCPTransport.http, url="https://global.example/mcp" + ) + servers: Final = (extended, global_deadline) + manager: Final = MCPServerManager() + manager.registry = {server.server_id: server for server in servers} + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + context: Final = operations.prepare_context(UserAPIKeyAuth(api_key="owned-caller", user_id="alice")) + + async def respond(request: httpx2.Request) -> httpx2.Response: + if request.method == "DELETE": + return httpx2.Response(200) + payload: Final = _JSONRPC_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + result: Final = { + "initialize": { + "protocolVersion": "2025-03-26", + "serverInfo": {"name": "test", "version": "1"}, + "capabilities": {"tools": {}}, + }, + "server/discover": { + "supportedVersions": ["2026-07-28"], + "capabilities": {"tools": {}}, + "ttlMs": 60000, + "cacheScope": "private", + "resultType": "complete", + }, + "tools/list": { + "tools": [{"name": "echo", "inputSchema": {"type": "object"}}], + "ttlMs": 60000, + "cacheScope": "private", + "resultType": "complete", + }, + }[payload.method] + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + + with _mcp_upstream(respond): + listing: Final = await catalog.aggregate_gateway_tools( + context, PaginatedRequestParams(), servers, {}, record_listing=True, enforce_rate_limits=False + ) + + assert [tool.name for tool in listing.tools] == [f"{get_server_prefix(extended)}{MCP_TOOL_PREFIX_SEPARATOR}echo"] + assert listing.outcomes["extended"].tag == "ok" + assert listing.outcomes["global"].tag == "timeout" + + @pytest.mark.asyncio async def test_failed_aggregate_continuation_preserves_only_delivered_tool_metadata(monkeypatch): from mcp.shared.exceptions import MCPError