Merge pull request #39189 from BerriAI/litellm_mcp_list_pagination_lit5594

fix(mcp): paginate prompt and resource discovery

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

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

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