Skip to content

Commit e401eae

Browse files
committed
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 <dongjiang1989@126.com>
1 parent 837ef90 commit e401eae

3 files changed

Lines changed: 342 additions & 1 deletion

File tree

src/mcp/client/__init__.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
InMemoryResponseCacheStore,
1111
ResponseCacheStore,
1212
)
13-
from mcp.client.client import Client
13+
from mcp.client.client import Client, CursorCycleError, PaginationExceededError
1414
from mcp.client.context import ClientRequestContext
1515
from mcp.client.extension import (
1616
ClaimContext,
@@ -32,9 +32,11 @@
3232
"ClientExtension",
3333
"ClientRequestContext",
3434
"ClientSession",
35+
"CursorCycleError",
3536
"InMemoryResponseCacheStore",
3637
"InputRequiredRoundsExceededError",
3738
"NotificationBinding",
39+
"PaginationExceededError",
3840
"ResponseCacheStore",
3941
"ResultClaim",
4042
"Transport",

src/mcp/client/client.py

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
ListToolsResult,
3333
LoggingLevel,
3434
PaginatedRequestParams,
35+
PaginatedResult,
3536
PromptReference,
3637
ReadResourceResult,
3738
RequestParamsMeta,
@@ -81,6 +82,37 @@
8182
_T = TypeVar("_T")
8283
_ResultT = TypeVar("_ResultT")
8384
_CacheableT = TypeVar("_CacheableT", bound=CacheableResult)
85+
_PaginatedCacheableT = TypeVar("_PaginatedCacheableT", bound=PaginatedResult)
86+
87+
_DEFAULT_LIST_MAX_PAGES = 64
88+
89+
90+
class PaginationExceededError(RuntimeError):
91+
"""A ``list_all_*`` walk exceeded the configured page limit.
92+
93+
Raised when automatic pagination does not terminate within
94+
``Client.list_max_pages`` pages, indicating a server that keeps
95+
returning ``next_cursor`` without end.
96+
"""
97+
98+
def __init__(self, method: str, max_pages: int) -> None:
99+
super().__init__(f"{method}: exceeded list_max_pages ({max_pages}); server pagination did not terminate")
100+
self.method = method
101+
self.max_pages = max_pages
102+
103+
104+
class CursorCycleError(RuntimeError):
105+
"""A ``list_all_*`` walk detected a repeated pagination cursor.
106+
107+
Raised when the server returns a ``next_cursor`` that was already
108+
seen during the current pagination walk, indicating a cursor cycle.
109+
"""
110+
111+
def __init__(self, method: str, cursor: str) -> None:
112+
super().__init__(f"{method}: pagination detected cursor cycle: {cursor!r}")
113+
self.method = method
114+
self.cursor = cursor
115+
84116

85117
_Connector = Callable[[AsyncExitStack, ConnectMode, bool], Awaitable["Dispatcher[Any]"]]
86118
"""Resolved at ``__post_init__`` from the shape of ``server`` alone: enter whatever resources
@@ -358,6 +390,14 @@ async def main():
358390
`meta` always reach the server. A `CacheConfig` with a custom `store` requires
359391
`target_id` when the server is not a URL (no identity can be derived)."""
360392

393+
list_max_pages: int = _DEFAULT_LIST_MAX_PAGES
394+
"""Maximum number of pages to fetch during ``list_all_*`` auto-pagination.
395+
396+
Defaults to 64 (matching the TypeScript SDK). A value of 0 disables the cap
397+
(unlimited pagination). A negative value also means unlimited. Raises
398+
`PaginationExceededError` when the limit is exceeded, and `CursorCycleError`
399+
when a repeated cursor is detected."""
400+
361401
_entered: bool = field(init=False, default=False)
362402
_session: ClientSession | None = field(init=False, default=None)
363403
_exit_stack: AsyncExitStack | None = field(init=False, default=None)
@@ -933,6 +973,102 @@ async def list_tools(
933973
),
934974
)
935975

976+
async def _drain_all_pages(
977+
self,
978+
method: str,
979+
fetch_page: Callable[[str | None], Awaitable[_PaginatedCacheableT]],
980+
get_items: Callable[[_PaginatedCacheableT], list[Any]],
981+
set_items: Callable[[_PaginatedCacheableT, list[Any]], None],
982+
) -> _PaginatedCacheableT:
983+
"""Drain all pages of a paginated list method into a single aggregated result.
984+
985+
Safety guards:
986+
- Cursor cycle detection via a seen-set; raises `CursorCycleError`.
987+
- Page cap enforced via ``self.list_max_pages``; raises `PaginationExceededError`.
988+
"""
989+
first = await fetch_page(None)
990+
cursor = first.next_cursor
991+
seen: set[str] = set()
992+
pages = 1
993+
max_pages = self.list_max_pages
994+
995+
# Use an accumulator to avoid quadratic list copying
996+
accumulator = get_items(first)
997+
998+
while cursor is not None:
999+
if cursor in seen:
1000+
raise CursorCycleError(method, cursor)
1001+
seen.add(cursor)
1002+
if max_pages > 0 and pages >= max_pages:
1003+
raise PaginationExceededError(method, max_pages)
1004+
page = await fetch_page(cursor)
1005+
accumulator.extend(get_items(page))
1006+
cursor = page.next_cursor
1007+
pages += 1
1008+
1009+
set_items(first, accumulator)
1010+
# Strip the terminal cursor from the aggregated result.
1011+
first.next_cursor = None
1012+
return first
1013+
1014+
async def list_all_tools(self) -> ListToolsResult:
1015+
"""Fetch all tools from the server, draining pagination automatically.
1016+
1017+
Raises:
1018+
PaginationExceededError: The walk exceeded ``list_max_pages``.
1019+
CursorCycleError: The server returned a repeated cursor.
1020+
"""
1021+
result = await self._drain_all_pages(
1022+
"tools/list",
1023+
fetch_page=lambda c: self.list_tools(cursor=c, cache_mode="bypass"),
1024+
get_items=lambda r: list(r.tools),
1025+
set_items=lambda r, items: setattr(r, "tools", items), # noqa: B010
1026+
)
1027+
# Absorb the final aggregated result as complete to prune stale session state
1028+
return self.session._absorb_tool_listing(result, complete=True) # pyright: ignore[reportPrivateUsage]
1029+
1030+
async def list_all_resources(self) -> ListResourcesResult:
1031+
"""Fetch all resources from the server, draining pagination automatically.
1032+
1033+
Raises:
1034+
PaginationExceededError: The walk exceeded ``list_max_pages``.
1035+
CursorCycleError: The server returned a repeated cursor.
1036+
"""
1037+
return await self._drain_all_pages(
1038+
"resources/list",
1039+
fetch_page=lambda c: self.list_resources(cursor=c, cache_mode="bypass"),
1040+
get_items=lambda r: list(r.resources),
1041+
set_items=lambda r, items: setattr(r, "resources", items), # noqa: B010
1042+
)
1043+
1044+
async def list_all_resource_templates(self) -> ListResourceTemplatesResult:
1045+
"""Fetch all resource templates from the server, draining pagination automatically.
1046+
1047+
Raises:
1048+
PaginationExceededError: The walk exceeded ``list_max_pages``.
1049+
CursorCycleError: The server returned a repeated cursor.
1050+
"""
1051+
return await self._drain_all_pages(
1052+
"resources/templates/list",
1053+
fetch_page=lambda c: self.list_resource_templates(cursor=c, cache_mode="bypass"),
1054+
get_items=lambda r: list(r.resource_templates),
1055+
set_items=lambda r, items: setattr(r, "resource_templates", items), # noqa: B010
1056+
)
1057+
1058+
async def list_all_prompts(self) -> ListPromptsResult:
1059+
"""Fetch all prompts from the server, draining pagination automatically.
1060+
1061+
Raises:
1062+
PaginationExceededError: The walk exceeded ``list_max_pages``.
1063+
CursorCycleError: The server returned a repeated cursor.
1064+
"""
1065+
return await self._drain_all_pages(
1066+
"prompts/list",
1067+
fetch_page=lambda c: self.list_prompts(cursor=c, cache_mode="bypass"),
1068+
get_items=lambda r: list(r.prompts),
1069+
set_items=lambda r, items: setattr(r, "prompts", items), # noqa: B010
1070+
)
1071+
9361072
@deprecated("The roots capability is deprecated as of 2026-07-28 (SEP-2577).", category=MCPDeprecationWarning)
9371073
async def send_roots_list_changed(self) -> None:
9381074
"""Send a notification that the roots list has changed."""
Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,203 @@
1+
"""Tests for list_all_* auto-pagination helpers with safety guards."""
2+
3+
from __future__ import annotations
4+
5+
import pytest
6+
from mcp_types import (
7+
ListPromptsResult,
8+
ListResourcesResult,
9+
ListResourceTemplatesResult,
10+
ListToolsResult,
11+
RequestParamsMeta,
12+
)
13+
14+
from mcp.client.caching import CacheMode
15+
from mcp.client.client import Client, CursorCycleError, PaginationExceededError
16+
from mcp.server.mcpserver import MCPServer
17+
18+
pytestmark = pytest.mark.anyio
19+
20+
21+
@pytest.fixture
22+
def multi_item_server() -> MCPServer:
23+
"""Server with several tools, resources, prompts, and templates."""
24+
server = MCPServer("paginated-test")
25+
26+
@server.tool()
27+
def tool_a() -> str: # pragma: no cover
28+
"""Tool A."""
29+
return "a"
30+
31+
@server.tool()
32+
def tool_b() -> str: # pragma: no cover
33+
"""Tool B."""
34+
return "b"
35+
36+
@server.tool()
37+
def tool_c() -> str: # pragma: no cover
38+
"""Tool C."""
39+
return "c"
40+
41+
@server.resource("test://r1")
42+
def res1() -> str: # pragma: no cover
43+
"""Resource 1."""
44+
return "r1"
45+
46+
@server.resource("test://r2")
47+
def res2() -> str: # pragma: no cover
48+
"""Resource 2."""
49+
return "r2"
50+
51+
@server.prompt()
52+
def prompt_a() -> str: # pragma: no cover
53+
"""Prompt A."""
54+
return "pa"
55+
56+
@server.resource("test://tmpl/{id}")
57+
def tmpl(id: str) -> str: # pragma: no cover
58+
"""Template."""
59+
return f"t-{id}"
60+
61+
return server
62+
63+
64+
async def test_list_all_tools_returns_all(multi_item_server: MCPServer) -> None:
65+
"""list_all_tools drains all pages and returns a flat result."""
66+
async with Client(multi_item_server, mode="legacy") as client:
67+
result = await client.list_all_tools()
68+
assert isinstance(result, ListToolsResult)
69+
assert len(result.tools) == 3
70+
assert result.next_cursor is None
71+
72+
73+
async def test_list_all_resources_returns_all(multi_item_server: MCPServer) -> None:
74+
"""list_all_resources drains all pages."""
75+
async with Client(multi_item_server, mode="legacy") as client:
76+
result = await client.list_all_resources()
77+
assert isinstance(result, ListResourcesResult)
78+
assert len(result.resources) == 2
79+
assert result.next_cursor is None
80+
81+
82+
async def test_list_all_prompts_returns_all(multi_item_server: MCPServer) -> None:
83+
"""list_all_prompts drains all pages."""
84+
async with Client(multi_item_server, mode="legacy") as client:
85+
result = await client.list_all_prompts()
86+
assert isinstance(result, ListPromptsResult)
87+
assert len(result.prompts) == 1
88+
assert result.next_cursor is None
89+
90+
91+
async def test_list_all_resource_templates_returns_all(multi_item_server: MCPServer) -> None:
92+
"""list_all_resource_templates drains all pages."""
93+
async with Client(multi_item_server, mode="legacy") as client:
94+
result = await client.list_all_resource_templates()
95+
assert isinstance(result, ListResourceTemplatesResult)
96+
assert len(result.resource_templates) == 1
97+
assert result.next_cursor is None
98+
99+
100+
async def test_list_all_empty_server() -> None:
101+
"""list_all_* on a server with no items returns empty lists."""
102+
server = MCPServer("empty-test")
103+
async with Client(server, mode="legacy") as client:
104+
tools = await client.list_all_tools()
105+
assert tools.tools == []
106+
assert tools.next_cursor is None
107+
108+
resources = await client.list_all_resources()
109+
assert resources.resources == []
110+
111+
prompts = await client.list_all_prompts()
112+
assert prompts.prompts == []
113+
114+
templates = await client.list_all_resource_templates()
115+
assert templates.resource_templates == []
116+
117+
118+
async def test_list_all_max_pages_exceeded() -> None:
119+
"""PaginationExceededError when the server doesn't terminate within max_pages."""
120+
server = MCPServer("infinite-test")
121+
122+
@server.tool()
123+
def dummy() -> str: # pragma: no cover
124+
"""A dummy tool."""
125+
return "x"
126+
127+
async with Client(server, mode="legacy", list_max_pages=2) as client:
128+
original = client.list_tools
129+
call_count = 0
130+
131+
async def infinite_list_tools(
132+
*,
133+
cursor: str | None = None,
134+
meta: RequestParamsMeta | None = None,
135+
cache_mode: CacheMode = "use",
136+
) -> ListToolsResult:
137+
nonlocal call_count
138+
call_count += 1
139+
result = await original(cursor=cursor, meta=meta, cache_mode=cache_mode)
140+
result.next_cursor = f"cursor-{call_count}"
141+
return result
142+
143+
client.list_tools = infinite_list_tools # type: ignore[method-assign]
144+
with pytest.raises(PaginationExceededError, match=r"exceeded list_max_pages \(2\)") as exc_info:
145+
await client.list_all_tools()
146+
assert exc_info.value.method == "tools/list"
147+
assert exc_info.value.max_pages == 2
148+
149+
150+
async def test_list_all_cursor_cycle_detected() -> None:
151+
"""CursorCycleError when the server returns a repeated cursor."""
152+
server = MCPServer("cycle-test")
153+
154+
@server.tool()
155+
def dummy() -> str: # pragma: no cover
156+
"""A dummy tool."""
157+
return "x"
158+
159+
async with Client(server, mode="legacy", list_max_pages=0) as client:
160+
original = client.list_tools
161+
cursors_seq = ["cursorA", "cursorB", "cursorA", "cursorA", "cursorA"]
162+
call_count = 0
163+
164+
async def cycling_list_tools(
165+
*,
166+
cursor: str | None = None,
167+
meta: RequestParamsMeta | None = None,
168+
cache_mode: CacheMode = "use",
169+
) -> ListToolsResult:
170+
nonlocal call_count
171+
call_count += 1
172+
result = await original(cursor=cursor, meta=meta, cache_mode=cache_mode)
173+
result.next_cursor = cursors_seq[call_count - 1] # pragma: no branch
174+
return result
175+
176+
client.list_tools = cycling_list_tools # type: ignore[method-assign]
177+
with pytest.raises(CursorCycleError, match=r"cursor cycle.*cursorA") as exc_info:
178+
await client.list_all_tools()
179+
assert exc_info.value.method == "tools/list"
180+
assert exc_info.value.cursor == "cursorA"
181+
182+
183+
async def test_list_all_unlimited_with_zero_max_pages() -> None:
184+
"""list_max_pages=0 disables the page cap (unlimited)."""
185+
server = MCPServer("unlimited-test")
186+
187+
@server.tool()
188+
def dummy() -> str: # pragma: no cover
189+
"""A dummy tool."""
190+
return "x"
191+
192+
async with Client(server, mode="legacy", list_max_pages=0) as client:
193+
result = await client.list_all_tools()
194+
assert isinstance(result, ListToolsResult)
195+
assert len(result.tools) == 1
196+
197+
198+
async def test_list_all_strips_terminal_cursor() -> None:
199+
"""The aggregated result has next_cursor=None."""
200+
server = MCPServer("strip-test")
201+
async with Client(server, mode="legacy") as client:
202+
result = await client.list_all_tools()
203+
assert result.next_cursor is None

0 commit comments

Comments
 (0)