diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 4b456710057..49434befd4e 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -7,12 +7,13 @@ import base64 import hashlib import json import os -from collections.abc import Awaitable, Callable, Generator +from collections.abc import Awaitable, Callable, Generator, Sequence from contextlib import AbstractAsyncContextManager from functools import partial from types import MappingProxyType from typing import Any, Final, TypeAlias, TypeVar +import anyio import httpx2 from httpx2._client import UseClientDefault from httpx2._types import AuthTypes @@ -38,6 +39,8 @@ from mcp.types import ( ListPromptsResult, ListResourcesResult, ListResourceTemplatesResult, + PaginatedRequestParams, + PaginatedResult, Prompt, ResourceTemplate, ServerNotification, @@ -49,7 +52,12 @@ from mcp.types import Tool as MCPTool from pydantic import AnyUrl from litellm._logging import verbose_logger -from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT +from litellm.constants import ( + MCP_CLIENT_TIMEOUT, + MCP_NPM_CACHE_DIR, + MCP_TOOL_LISTING_MAX_PAGES, + MCP_TOOL_LISTING_TIMEOUT, +) from litellm.experimental_mcp_client.tools import list_tools_with_pagination from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.proxy._experimental.mcp_server.mcp_debug import capture_upstream_error_response @@ -147,6 +155,8 @@ def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None: TSessionResult = TypeVar("TSessionResult") +_ListPage = TypeVar("_ListPage", bound=PaginatedResult) +_ListItem = TypeVar("_ListItem") class _MCPHTTPClient(httpx2.AsyncClient): @@ -793,6 +803,33 @@ class MCPClient: # Return a default error result instead of raising return self.error_tool_result(e) + async def _list_optional_pages( + self, + fetch_page: Callable[[PaginatedRequestParams | None], Awaitable[_ListPage]], + items_of: Callable[[_ListPage], Sequence[_ListItem]], + ) -> list[_ListItem]: # mutable-ok: existing list discovery API + items: Final[list[_ListItem]] = [] # mutable-ok: bounded iterative page accumulation + cursors: Final[set[str]] = set() # mutable-ok: constant-time detection of cursor cycles + cursor: str | None = None # rebind-ok: iterative traversal avoids recursion at the existing page cap + with anyio.fail_after(max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)): + for page_index in range(MCP_TOOL_LISTING_MAX_PAGES): + try: + page = await fetch_page( # rebind-ok: each SDK page replaces the previous one + None if cursor is None else PaginatedRequestParams(cursor=cursor) + ) + except MCPError as error: + if page_index > 0 and error.error.code == METHOD_NOT_FOUND: + raise RuntimeError("MCP list operation became unavailable during pagination") from error + raise + items.extend(items_of(page)) + if not page.next_cursor: + return items + if page.next_cursor in cursors: + raise RuntimeError("MCP list pagination repeated a cursor") + cursors.add(page.next_cursor) + cursor = page.next_cursor + raise RuntimeError(f"MCP list pagination exceeded {MCP_TOOL_LISTING_MAX_PAGES} pages") + async def list_prompts(self, *, raise_on_error: bool = False) -> list[Prompt]: """List available prompts from the server.""" verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") @@ -802,7 +839,11 @@ class MCPClient: if capabilities is not None and capabilities.prompts is None: return ListPromptsResult(prompts=[]) try: - return await session.list_prompts() + return ListPromptsResult( + prompts=await self._list_optional_pages( + lambda params: session.list_prompts(params=params), lambda page: page.prompts + ) + ) except MCPError as error: if error.error.code != METHOD_NOT_FOUND: raise @@ -892,7 +933,11 @@ class MCPClient: if capabilities is not None and capabilities.resources is None: return ListResourcesResult(resources=[]) try: - return await session.list_resources() + return ListResourcesResult( + resources=await self._list_optional_pages( + lambda params: session.list_resources(params=params), lambda page: page.resources + ) + ) except MCPError as error: if error.error.code != METHOD_NOT_FOUND: raise @@ -941,7 +986,12 @@ class MCPClient: if capabilities is not None and capabilities.resources is None: return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload try: - return await session.list_resource_templates() + return ListResourceTemplatesResult( + resource_templates=await self._list_optional_pages( + lambda params: session.list_resource_templates(params=params), + lambda page: page.resource_templates, + ) + ) except MCPError as error: if error.error.code != METHOD_NOT_FOUND: raise diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 4b698f1258d..6c20ef135ba 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -2036,6 +2036,15 @@ async def test_optional_discovery_preserves_cancellation(method: str) -> None: }, }, ) + if not (payload.params or {}).get("cursor"): + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + return httpx2.Response( + 200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [], "nextCursor": "pending-page"}} + ) ready.set() await pending.wait() return httpx2.Response(202) @@ -2055,6 +2064,255 @@ async def test_optional_discovery_preserves_cancellation(method: str) -> None: await asyncio.wait_for(task, timeout=3) +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +@pytest.mark.parametrize("session_id", (None, "pagination-session")) +@pytest.mark.parametrize("empty_middle", (False, True)) +async def test_optional_discovery_collects_all_pages(method: str, session_id: str | None, empty_middle: bool) -> None: + from mcp.types import Prompt, PromptArgument, Resource, ResourceTemplate + + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + entries: Final = tuple( + { + "prompts/list": Prompt( + name=f"item-{index}", + description="prompt description", + arguments=[PromptArgument(name="query", required=True)], + ), + "resources/list": Resource( + name=f"item-{index}", + uri=f"test://item/{index}", + mime_type="text/plain", + description="resource description", + ), + "resources/templates/list": ResourceTemplate( + name=f"item-{index}", uri_template=f"test://item/{index}/{{query}}", mime_type="text/plain" + ), + }[method] + for index in range(5) + ) + + def respond(request: httpx2.Request) -> httpx2.Response: + if request.method == "GET": + return httpx2.Response(405) + if request.method == "DELETE": + return httpx2.Response(200) + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + if payload.method == "initialize": + return httpx2.Response( + 200, + headers={"mcp-session-id": session_id} if session_id else {}, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": payload.params["protocolVersion"], + "capabilities": {"prompts": {}, "resources": {}}, + "serverInfo": {"name": "paged", "version": "1"}, + }, + }, + ) + assert payload.method == method + assert request.headers.get("mcp-session-id") == session_id + cursor: Final = (payload.params or {}).get("cursor") + assert cursor in (None, "opaque:/second+page", "opaque:/last+page") + page: Final = ( + entries[:3] if cursor is None else (() if empty_middle and cursor == "opaque:/second+page" else entries[3:]) + ) + next_cursor: Final = ( + "opaque:/second+page" + if cursor is None + else "opaque:/last+page" + if empty_middle and cursor == "opaque:/second+page" + else "" + ) + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + field: [item.model_dump(mode="json", by_alias=True) for item in page], + "nextCursor": next_cursor, + }, + }, + ) + + responder: Final = Mock(side_effect=respond) + client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + assert await operation(raise_on_error=True) == list(entries) + requests: Final = tuple( + _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content) + for call in responder.call_args_list + if call.args[0].method == "POST" + ) + assert sum(isinstance(request, JSONRPCRequest) and request.method == "initialize" for request in requests) == 1 + assert tuple( + (request.params or {}).get("cursor") + for request in requests + if isinstance(request, JSONRPCRequest) and request.method == method + ) == ((None, "opaque:/second+page", "opaque:/last+page") if empty_middle else (None, "opaque:/second+page")) + assert sum(call.args[0].method == "DELETE" for call in responder.call_args_list) == (1 if session_id else 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +@pytest.mark.parametrize( + "failure", ("repeat", "cycle", "cap", "method_not_found", "internal_error", "unauthorized", "deadline") +) +@pytest.mark.parametrize("strict", (False, True)) +async def test_optional_discovery_rejects_incomplete_walks( + method: str, failure: str, strict: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_MAX_PAGES", 3 if failure == "cycle" else 2, raising=False) + monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_TIMEOUT", 0.05) + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + entry: Final = { + "prompts/list": {"name": "first"}, + "resources/list": {"name": "first", "uri": "test://first"}, + "resources/templates/list": {"name": "first", "uriTemplate": "test://{name}"}, + }[method] + cancelled: Final = asyncio.Event() + + async def respond(request: httpx2.Request) -> httpx2.Response: + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + if payload.method == "initialize": + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": payload.params["protocolVersion"], + "capabilities": {"prompts": {}, "resources": {}}, + "serverInfo": {"name": "interrupted", "version": "1"}, + }, + }, + ) + assert payload.method == method + cursor: Final = (payload.params or {}).get("cursor") + if cursor is not None: + if failure == "deadline": + try: + await asyncio.Event().wait() + finally: + cancelled.set() + if failure == "unauthorized": + return httpx2.Response(401) + if failure in ("method_not_found", "internal_error"): + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "error": { + "code": -32601 if failure == "method_not_found" else -32603, + "message": "Later page unavailable", + }, + }, + ) + next_cursor: Final = ( + "private-cursor-2" if cursor == "private-cursor-1" and failure != "repeat" else "private-cursor-1" + ) + return httpx2.Response( + 200, json={"jsonrpc": "2.0", "id": payload.id, "result": {field: [entry], "nextCursor": next_cursor}} + ) + + responder: Final = AsyncMock(side_effect=respond) + client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp", timeout=0.2) + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + if strict: + error_type: Final = { + "internal_error": MCPError, + "unauthorized": httpx2.HTTPStatusError, + "deadline": TimeoutError, + }.get(failure, RuntimeError) + with pytest.raises(error_type): + await operation(raise_on_error=True) + else: + assert await operation() == [] + assert len( + tuple( + payload + for call in responder.call_args_list + if isinstance(payload := _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content), JSONRPCRequest) + and payload.method == method + ) + ) == (3 if failure == "cycle" else 2) + assert "private-cursor" not in caplog.text + if failure == "deadline": + assert cancelled.is_set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ("prompts/list", "resources/list", "resources/templates/list")) +async def test_optional_discovery_allows_exhaustion_at_page_cap(method: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(mcp_client_module, "MCP_TOOL_LISTING_MAX_PAGES", 2, raising=False) + field: Final = { + "prompts/list": "prompts", + "resources/list": "resources", + "resources/templates/list": "resourceTemplates", + }[method] + + def respond(request: httpx2.Request) -> httpx2.Response: + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + if payload.method == "initialize": + result: Final = { + "protocolVersion": payload.params["protocolVersion"], + "capabilities": {"prompts": {}, "resources": {}}, + "serverInfo": {"name": "empty-pages", "version": "1"}, + } + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + assert payload.method == method + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": {field: [], "nextCursor": None if (payload.params or {}).get("cursor") else "last-page"}, + }, + ) + + responder: Final = Mock(side_effect=respond) + client: Final = _MockTransportClient(responder, server_url="https://example.com/mcp") + operation: Final = { + "prompts/list": client.list_prompts, + "resources/list": client.list_resources, + "resources/templates/list": client.list_resource_templates, + }[method] + assert await operation(raise_on_error=True) == [] + assert ( + sum( + isinstance(payload := _JSONRPC_MESSAGE_ADAPTER.validate_json(call.args[0].content), JSONRPCRequest) + and payload.method == method + for call in responder.call_args_list + ) + == 2 + ) + def test_client_import_before_proxy_credentials_succeeds_in_fresh_process(): import subprocess 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 9140ac61f1a..7418cf67e5f 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 @@ -13182,6 +13182,8 @@ class _DiscoveryUpstream: await self.release.wait() if self.outcome == "failure": return httpx2.Response(503) + if self.outcome == "paged_failure" and (payload.params or {}).get("cursor"): + return httpx2.Response(503) if self.outcome == "cancelled": raise asyncio.CancelledError() if self.outcome == "rejected": @@ -13196,7 +13198,12 @@ class _DiscoveryUpstream: }, "tools/list": {"tools": []}, }[payload.method] - return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + continuation: Final = ( + {"nextCursor": "last-page"} + if self.outcome in ("paged", "paged_failure") and not (payload.params or {}).get("cursor") + else {} + ) + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {**result, **continuation}}) @property def initializes(self) -> int: @@ -13262,6 +13269,29 @@ async def test_discovery_cache_empty_results_and_failures(kind: str, outcome: st assert upstream.initializes == 3 +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ("prompts", "resources", "templates")) +async def test_discovery_cache_retries_failed_pagination_before_caching_complete_list(kind: str) -> None: + manager: Final = MCPServerManager() + upstream: Final = _DiscoveryUpstream() + upstream.outcome = "paged_failure" + operation: Final = { + "prompts": manager.get_prompts_from_server, + "resources": manager.get_resources_from_server, + "templates": manager.get_resource_templates_from_server, + }[kind] + with _mcp_upstream(upstream.respond): + assert await operation(_discovery_server(), None) == [] + assert upstream.initializes == 1 + upstream.outcome = "paged" + recovered: Final = await operation(_discovery_server(), None) + assert [item.name for item in recovered] == ["discovery-example", "discovery-example"] + assert upstream.initializes == 2 + requests_after_recovery: Final = upstream.requests + assert await operation(_discovery_server(), None) == recovered + assert upstream.requests == requests_after_recovery + + @pytest.mark.asyncio async def test_discovery_cache_isolates_forwarded_credentials_and_shares_static_auth() -> None: import respx