From e401eae543782fa8b8610c07a562e903c98291ec Mon Sep 17 00:00:00 2001 From: dongjiang Date: Fri, 24 Jul 2026 18:55:47 +0800 Subject: [PATCH] mcp/client: add list_all_* helpers with pagination safety guards Add four auto-pagination helpers (list_all_tools, list_all_resources, list_all_resource_templates, list_all_prompts) that drain all pages into a single aggregated result with safety guards: - Cursor cycle detection via a seen-set; raises CursorCycleError - Configurable page cap via list_max_pages (default 64); raises PaginationExceededError - Use accumulator with extend() to avoid quadratic list copying - Absorb final aggregated tools result as complete to prune stale session state (x-mcp-header maps and output schemas) Aligns with TypeScript SDK's _listAllPages helper (DEFAULT_LIST_MAX_PAGES=64). Signed-off-by: dongjiang --- src/mcp/client/__init__.py | 4 +- src/mcp/client/client.py | 136 +++++++++++++++ tests/client/test_list_all_pagination.py | 203 +++++++++++++++++++++++ 3 files changed, 342 insertions(+), 1 deletion(-) create mode 100644 tests/client/test_list_all_pagination.py diff --git a/src/mcp/client/__init__.py b/src/mcp/client/__init__.py index 21581749d0..de6a2f1837 100644 --- a/src/mcp/client/__init__.py +++ b/src/mcp/client/__init__.py @@ -10,7 +10,7 @@ InMemoryResponseCacheStore, ResponseCacheStore, ) -from mcp.client.client import Client +from mcp.client.client import Client, CursorCycleError, PaginationExceededError from mcp.client.context import ClientRequestContext from mcp.client.extension import ( ClaimContext, @@ -32,9 +32,11 @@ "ClientExtension", "ClientRequestContext", "ClientSession", + "CursorCycleError", "InMemoryResponseCacheStore", "InputRequiredRoundsExceededError", "NotificationBinding", + "PaginationExceededError", "ResponseCacheStore", "ResultClaim", "Transport", diff --git a/src/mcp/client/client.py b/src/mcp/client/client.py index e840675011..827e2678b4 100644 --- a/src/mcp/client/client.py +++ b/src/mcp/client/client.py @@ -32,6 +32,7 @@ ListToolsResult, LoggingLevel, PaginatedRequestParams, + PaginatedResult, PromptReference, ReadResourceResult, RequestParamsMeta, @@ -81,6 +82,37 @@ _T = TypeVar("_T") _ResultT = TypeVar("_ResultT") _CacheableT = TypeVar("_CacheableT", bound=CacheableResult) +_PaginatedCacheableT = TypeVar("_PaginatedCacheableT", bound=PaginatedResult) + +_DEFAULT_LIST_MAX_PAGES = 64 + + +class PaginationExceededError(RuntimeError): + """A ``list_all_*`` walk exceeded the configured page limit. + + Raised when automatic pagination does not terminate within + ``Client.list_max_pages`` pages, indicating a server that keeps + returning ``next_cursor`` without end. + """ + + def __init__(self, method: str, max_pages: int) -> None: + super().__init__(f"{method}: exceeded list_max_pages ({max_pages}); server pagination did not terminate") + self.method = method + self.max_pages = max_pages + + +class CursorCycleError(RuntimeError): + """A ``list_all_*`` walk detected a repeated pagination cursor. + + Raised when the server returns a ``next_cursor`` that was already + seen during the current pagination walk, indicating a cursor cycle. + """ + + def __init__(self, method: str, cursor: str) -> None: + super().__init__(f"{method}: pagination detected cursor cycle: {cursor!r}") + self.method = method + self.cursor = cursor + _Connector = Callable[[AsyncExitStack, ConnectMode, bool], Awaitable["Dispatcher[Any]"]] """Resolved at ``__post_init__`` from the shape of ``server`` alone: enter whatever resources @@ -358,6 +390,14 @@ async def main(): `meta` always reach the server. A `CacheConfig` with a custom `store` requires `target_id` when the server is not a URL (no identity can be derived).""" + list_max_pages: int = _DEFAULT_LIST_MAX_PAGES + """Maximum number of pages to fetch during ``list_all_*`` auto-pagination. + + Defaults to 64 (matching the TypeScript SDK). A value of 0 disables the cap + (unlimited pagination). A negative value also means unlimited. Raises + `PaginationExceededError` when the limit is exceeded, and `CursorCycleError` + when a repeated cursor is detected.""" + _entered: bool = field(init=False, default=False) _session: ClientSession | None = field(init=False, default=None) _exit_stack: AsyncExitStack | None = field(init=False, default=None) @@ -933,6 +973,102 @@ async def list_tools( ), ) + async def _drain_all_pages( + self, + method: str, + fetch_page: Callable[[str | None], Awaitable[_PaginatedCacheableT]], + get_items: Callable[[_PaginatedCacheableT], list[Any]], + set_items: Callable[[_PaginatedCacheableT, list[Any]], None], + ) -> _PaginatedCacheableT: + """Drain all pages of a paginated list method into a single aggregated result. + + Safety guards: + - Cursor cycle detection via a seen-set; raises `CursorCycleError`. + - Page cap enforced via ``self.list_max_pages``; raises `PaginationExceededError`. + """ + first = await fetch_page(None) + cursor = first.next_cursor + seen: set[str] = set() + pages = 1 + max_pages = self.list_max_pages + + # Use an accumulator to avoid quadratic list copying + accumulator = get_items(first) + + while cursor is not None: + if cursor in seen: + raise CursorCycleError(method, cursor) + seen.add(cursor) + if max_pages > 0 and pages >= max_pages: + raise PaginationExceededError(method, max_pages) + page = await fetch_page(cursor) + accumulator.extend(get_items(page)) + cursor = page.next_cursor + pages += 1 + + set_items(first, accumulator) + # Strip the terminal cursor from the aggregated result. + first.next_cursor = None + return first + + async def list_all_tools(self) -> ListToolsResult: + """Fetch all tools from the server, draining pagination automatically. + + Raises: + PaginationExceededError: The walk exceeded ``list_max_pages``. + CursorCycleError: The server returned a repeated cursor. + """ + result = await self._drain_all_pages( + "tools/list", + fetch_page=lambda c: self.list_tools(cursor=c, cache_mode="bypass"), + get_items=lambda r: list(r.tools), + set_items=lambda r, items: setattr(r, "tools", items), # noqa: B010 + ) + # Absorb the final aggregated result as complete to prune stale session state + return self.session._absorb_tool_listing(result, complete=True) # pyright: ignore[reportPrivateUsage] + + async def list_all_resources(self) -> ListResourcesResult: + """Fetch all resources from the server, draining pagination automatically. + + Raises: + PaginationExceededError: The walk exceeded ``list_max_pages``. + CursorCycleError: The server returned a repeated cursor. + """ + return await self._drain_all_pages( + "resources/list", + fetch_page=lambda c: self.list_resources(cursor=c, cache_mode="bypass"), + get_items=lambda r: list(r.resources), + set_items=lambda r, items: setattr(r, "resources", items), # noqa: B010 + ) + + async def list_all_resource_templates(self) -> ListResourceTemplatesResult: + """Fetch all resource templates from the server, draining pagination automatically. + + Raises: + PaginationExceededError: The walk exceeded ``list_max_pages``. + CursorCycleError: The server returned a repeated cursor. + """ + return await self._drain_all_pages( + "resources/templates/list", + fetch_page=lambda c: self.list_resource_templates(cursor=c, cache_mode="bypass"), + get_items=lambda r: list(r.resource_templates), + set_items=lambda r, items: setattr(r, "resource_templates", items), # noqa: B010 + ) + + async def list_all_prompts(self) -> ListPromptsResult: + """Fetch all prompts from the server, draining pagination automatically. + + Raises: + PaginationExceededError: The walk exceeded ``list_max_pages``. + CursorCycleError: The server returned a repeated cursor. + """ + return await self._drain_all_pages( + "prompts/list", + fetch_page=lambda c: self.list_prompts(cursor=c, cache_mode="bypass"), + get_items=lambda r: list(r.prompts), + set_items=lambda r, items: setattr(r, "prompts", items), # noqa: B010 + ) + @deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning) async def send_roots_list_changed(self) -> None: """Send a notification that the roots list has changed.""" diff --git a/tests/client/test_list_all_pagination.py b/tests/client/test_list_all_pagination.py new file mode 100644 index 0000000000..61d1e739d6 --- /dev/null +++ b/tests/client/test_list_all_pagination.py @@ -0,0 +1,203 @@ +"""Tests for list_all_* auto-pagination helpers with safety guards.""" + +from __future__ import annotations + +import pytest +from mcp_types import ( + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, + ListToolsResult, + RequestParamsMeta, +) + +from mcp.client.caching import CacheMode +from mcp.client.client import Client, CursorCycleError, PaginationExceededError +from mcp.server.mcpserver import MCPServer + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def multi_item_server() -> MCPServer: + """Server with several tools, resources, prompts, and templates.""" + server = MCPServer("paginated-test") + + @server.tool() + def tool_a() -> str: # pragma: no cover + """Tool A.""" + return "a" + + @server.tool() + def tool_b() -> str: # pragma: no cover + """Tool B.""" + return "b" + + @server.tool() + def tool_c() -> str: # pragma: no cover + """Tool C.""" + return "c" + + @server.resource("test://r1") + def res1() -> str: # pragma: no cover + """Resource 1.""" + return "r1" + + @server.resource("test://r2") + def res2() -> str: # pragma: no cover + """Resource 2.""" + return "r2" + + @server.prompt() + def prompt_a() -> str: # pragma: no cover + """Prompt A.""" + return "pa" + + @server.resource("test://tmpl/{id}") + def tmpl(id: str) -> str: # pragma: no cover + """Template.""" + return f"t-{id}" + + return server + + +async def test_list_all_tools_returns_all(multi_item_server: MCPServer) -> None: + """list_all_tools drains all pages and returns a flat result.""" + async with Client(multi_item_server, mode="legacy") as client: + result = await client.list_all_tools() + assert isinstance(result, ListToolsResult) + assert len(result.tools) == 3 + assert result.next_cursor is None + + +async def test_list_all_resources_returns_all(multi_item_server: MCPServer) -> None: + """list_all_resources drains all pages.""" + async with Client(multi_item_server, mode="legacy") as client: + result = await client.list_all_resources() + assert isinstance(result, ListResourcesResult) + assert len(result.resources) == 2 + assert result.next_cursor is None + + +async def test_list_all_prompts_returns_all(multi_item_server: MCPServer) -> None: + """list_all_prompts drains all pages.""" + async with Client(multi_item_server, mode="legacy") as client: + result = await client.list_all_prompts() + assert isinstance(result, ListPromptsResult) + assert len(result.prompts) == 1 + assert result.next_cursor is None + + +async def test_list_all_resource_templates_returns_all(multi_item_server: MCPServer) -> None: + """list_all_resource_templates drains all pages.""" + async with Client(multi_item_server, mode="legacy") as client: + result = await client.list_all_resource_templates() + assert isinstance(result, ListResourceTemplatesResult) + assert len(result.resource_templates) == 1 + assert result.next_cursor is None + + +async def test_list_all_empty_server() -> None: + """list_all_* on a server with no items returns empty lists.""" + server = MCPServer("empty-test") + async with Client(server, mode="legacy") as client: + tools = await client.list_all_tools() + assert tools.tools == [] + assert tools.next_cursor is None + + resources = await client.list_all_resources() + assert resources.resources == [] + + prompts = await client.list_all_prompts() + assert prompts.prompts == [] + + templates = await client.list_all_resource_templates() + assert templates.resource_templates == [] + + +async def test_list_all_max_pages_exceeded() -> None: + """PaginationExceededError when the server doesn't terminate within max_pages.""" + server = MCPServer("infinite-test") + + @server.tool() + def dummy() -> str: # pragma: no cover + """A dummy tool.""" + return "x" + + async with Client(server, mode="legacy", list_max_pages=2) as client: + original = client.list_tools + call_count = 0 + + async def infinite_list_tools( + *, + cursor: str | None = None, + meta: RequestParamsMeta | None = None, + cache_mode: CacheMode = "use", + ) -> ListToolsResult: + nonlocal call_count + call_count += 1 + result = await original(cursor=cursor, meta=meta, cache_mode=cache_mode) + result.next_cursor = f"cursor-{call_count}" + return result + + client.list_tools = infinite_list_tools # type: ignore[method-assign] + with pytest.raises(PaginationExceededError, match=r"exceeded list_max_pages \(2\)") as exc_info: + await client.list_all_tools() + assert exc_info.value.method == "tools/list" + assert exc_info.value.max_pages == 2 + + +async def test_list_all_cursor_cycle_detected() -> None: + """CursorCycleError when the server returns a repeated cursor.""" + server = MCPServer("cycle-test") + + @server.tool() + def dummy() -> str: # pragma: no cover + """A dummy tool.""" + return "x" + + async with Client(server, mode="legacy", list_max_pages=0) as client: + original = client.list_tools + cursors_seq = ["cursorA", "cursorB", "cursorA", "cursorA", "cursorA"] + call_count = 0 + + async def cycling_list_tools( + *, + cursor: str | None = None, + meta: RequestParamsMeta | None = None, + cache_mode: CacheMode = "use", + ) -> ListToolsResult: + nonlocal call_count + call_count += 1 + result = await original(cursor=cursor, meta=meta, cache_mode=cache_mode) + result.next_cursor = cursors_seq[call_count - 1] # pragma: no branch + return result + + client.list_tools = cycling_list_tools # type: ignore[method-assign] + with pytest.raises(CursorCycleError, match=r"cursor cycle.*cursorA") as exc_info: + await client.list_all_tools() + assert exc_info.value.method == "tools/list" + assert exc_info.value.cursor == "cursorA" + + +async def test_list_all_unlimited_with_zero_max_pages() -> None: + """list_max_pages=0 disables the page cap (unlimited).""" + server = MCPServer("unlimited-test") + + @server.tool() + def dummy() -> str: # pragma: no cover + """A dummy tool.""" + return "x" + + async with Client(server, mode="legacy", list_max_pages=0) as client: + result = await client.list_all_tools() + assert isinstance(result, ListToolsResult) + assert len(result.tools) == 1 + + +async def test_list_all_strips_terminal_cursor() -> None: + """The aggregated result has next_cursor=None.""" + server = MCPServer("strip-test") + async with Client(server, mode="legacy") as client: + result = await client.list_all_tools() + assert result.next_cursor is None