Skip to content

Commit 2b78338

Browse files
committed
Harden the drain guard against cycling pagination cursors
Comparing only against the immediately preceding cursor catches a server that echoes the cursor back but not one that alternates between cursors (a, b, a, ...), which would still page forever. Track every cursor seen during the drain and raise on any repeat, in both the Client drains and the ClientSessionGroup aggregation drain.
1 parent 6e250ac commit 2b78338

5 files changed

Lines changed: 80 additions & 33 deletions

File tree

docs/advanced/pagination.md

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -65,10 +65,11 @@ That loop is the same one in every client that pages, so `Client` ships it. The
6565
`ClientSessionGroup` aggregation drains the same way, so a group fronting several servers reports the full collection instead of each server's first page. That aggregator is **[Session groups](../client/session-groups.md)**.
6666

6767
!!! warning
68-
A drain trusts the server to advance the cursor. A server that keeps returning the same
69-
`next_cursor` it was handed would page forever, so the drains stop and raise `RuntimeError`
70-
the moment a cursor fails to move. A page that does not advance is a broken server, and a
71-
loud failure beats a silent hang or a half-read list.
68+
A drain trusts the server to advance the cursor. A server that echoes back the
69+
`next_cursor` it was handed, or cycles through a longer loop of them, would page forever,
70+
so the drains remember every cursor they have seen and raise `RuntimeError` the moment one
71+
repeats. A repeated cursor is a broken server, and a loud failure beats a silent hang or a
72+
half-read list.
7273

7374
## The three rules
7475

src/mcp/client/client.py

Lines changed: 32 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -955,19 +955,20 @@ async def iter_all_tools(self, *, meta: RequestParamsMeta | None = None) -> Asyn
955955
materializing the full list in memory.
956956
957957
Raises:
958-
RuntimeError: The server returned a pagination cursor that did not advance.
958+
RuntimeError: The server returned a pagination cursor it already
959+
returned, which would page forever.
959960
"""
961+
seen_cursors: set[str] = set()
960962
cursor: str | None = None
961963
while True:
962964
result = await self.list_tools(cursor=cursor, meta=meta)
963965
for tool in result.tools:
964966
yield tool
965967
if result.next_cursor is None:
966968
return
967-
if result.next_cursor == cursor:
968-
raise RuntimeError(
969-
"Server returned a pagination cursor that did not advance; refusing to page forever."
970-
)
969+
if result.next_cursor in seen_cursors:
970+
raise RuntimeError("Server returned a pagination cursor it already returned; refusing to page forever.")
971+
seen_cursors.add(result.next_cursor)
971972
cursor = result.next_cursor
972973

973974
async def list_all_tools(self, *, meta: RequestParamsMeta | None = None) -> list[Tool]:
@@ -978,61 +979,66 @@ async def list_all_tools(self, *, meta: RequestParamsMeta | None = None) -> list
978979
list.
979980
980981
Raises:
981-
RuntimeError: The server returned a pagination cursor that did not advance.
982+
RuntimeError: The server returned a pagination cursor it already
983+
returned, which would page forever.
982984
"""
983985
return [tool async for tool in self.iter_all_tools(meta=meta)]
984986

985987
async def iter_all_prompts(self, *, meta: RequestParamsMeta | None = None) -> AsyncIterator[Prompt]:
986988
"""Yield every prompt from the server, paging through `next_cursor`.
987989
988990
Raises:
989-
RuntimeError: The server returned a pagination cursor that did not advance.
991+
RuntimeError: The server returned a pagination cursor it already
992+
returned, which would page forever.
990993
"""
994+
seen_cursors: set[str] = set()
991995
cursor: str | None = None
992996
while True:
993997
result = await self.list_prompts(cursor=cursor, meta=meta)
994998
for prompt in result.prompts:
995999
yield prompt
9961000
if result.next_cursor is None:
9971001
return
998-
if result.next_cursor == cursor:
999-
raise RuntimeError(
1000-
"Server returned a pagination cursor that did not advance; refusing to page forever."
1001-
)
1002+
if result.next_cursor in seen_cursors:
1003+
raise RuntimeError("Server returned a pagination cursor it already returned; refusing to page forever.")
1004+
seen_cursors.add(result.next_cursor)
10021005
cursor = result.next_cursor
10031006

10041007
async def list_all_prompts(self, *, meta: RequestParamsMeta | None = None) -> list[Prompt]:
10051008
"""List every prompt from the server, draining `next_cursor` across pages.
10061009
10071010
Raises:
1008-
RuntimeError: The server returned a pagination cursor that did not advance.
1011+
RuntimeError: The server returned a pagination cursor it already
1012+
returned, which would page forever.
10091013
"""
10101014
return [prompt async for prompt in self.iter_all_prompts(meta=meta)]
10111015

10121016
async def iter_all_resources(self, *, meta: RequestParamsMeta | None = None) -> AsyncIterator[Resource]:
10131017
"""Yield every resource from the server, paging through `next_cursor`.
10141018
10151019
Raises:
1016-
RuntimeError: The server returned a pagination cursor that did not advance.
1020+
RuntimeError: The server returned a pagination cursor it already
1021+
returned, which would page forever.
10171022
"""
1023+
seen_cursors: set[str] = set()
10181024
cursor: str | None = None
10191025
while True:
10201026
result = await self.list_resources(cursor=cursor, meta=meta)
10211027
for resource in result.resources:
10221028
yield resource
10231029
if result.next_cursor is None:
10241030
return
1025-
if result.next_cursor == cursor:
1026-
raise RuntimeError(
1027-
"Server returned a pagination cursor that did not advance; refusing to page forever."
1028-
)
1031+
if result.next_cursor in seen_cursors:
1032+
raise RuntimeError("Server returned a pagination cursor it already returned; refusing to page forever.")
1033+
seen_cursors.add(result.next_cursor)
10291034
cursor = result.next_cursor
10301035

10311036
async def list_all_resources(self, *, meta: RequestParamsMeta | None = None) -> list[Resource]:
10321037
"""List every resource from the server, draining `next_cursor` across pages.
10331038
10341039
Raises:
1035-
RuntimeError: The server returned a pagination cursor that did not advance.
1040+
RuntimeError: The server returned a pagination cursor it already
1041+
returned, which would page forever.
10361042
"""
10371043
return [resource async for resource in self.iter_all_resources(meta=meta)]
10381044

@@ -1042,26 +1048,28 @@ async def iter_all_resource_templates(
10421048
"""Yield every resource template from the server, paging through `next_cursor`.
10431049
10441050
Raises:
1045-
RuntimeError: The server returned a pagination cursor that did not advance.
1051+
RuntimeError: The server returned a pagination cursor it already
1052+
returned, which would page forever.
10461053
"""
1054+
seen_cursors: set[str] = set()
10471055
cursor: str | None = None
10481056
while True:
10491057
result = await self.list_resource_templates(cursor=cursor, meta=meta)
10501058
for template in result.resource_templates:
10511059
yield template
10521060
if result.next_cursor is None:
10531061
return
1054-
if result.next_cursor == cursor:
1055-
raise RuntimeError(
1056-
"Server returned a pagination cursor that did not advance; refusing to page forever."
1057-
)
1062+
if result.next_cursor in seen_cursors:
1063+
raise RuntimeError("Server returned a pagination cursor it already returned; refusing to page forever.")
1064+
seen_cursors.add(result.next_cursor)
10581065
cursor = result.next_cursor
10591066

10601067
async def list_all_resource_templates(self, *, meta: RequestParamsMeta | None = None) -> list[ResourceTemplate]:
10611068
"""List every resource template from the server, draining `next_cursor` across pages.
10621069
10631070
Raises:
1064-
RuntimeError: The server returned a pagination cursor that did not advance.
1071+
RuntimeError: The server returned a pagination cursor it already
1072+
returned, which would page forever.
10651073
"""
10661074
return [template async for template in self.iter_all_resource_templates(meta=meta)]
10671075

src/mcp/client/session_group.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,9 +78,11 @@ async def _drain_paginated(
7878
the list attribute on the result (e.g. `"tools"`, `"prompts"`).
7979
8080
Raises:
81-
RuntimeError: The server returned a pagination cursor that did not advance.
81+
RuntimeError: The server returned a pagination cursor it already
82+
returned, which would page forever.
8283
"""
8384
items: list[Any] = []
85+
seen_cursors: set[str] = set()
8486
cursor: str | None = None
8587
while True:
8688
params = types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None
@@ -89,8 +91,9 @@ async def _drain_paginated(
8991
next_cursor = getattr(result, "next_cursor", None)
9092
if next_cursor is None:
9193
return items
92-
if next_cursor == cursor:
93-
raise RuntimeError("Server returned a pagination cursor that did not advance; refusing to page forever.")
94+
if next_cursor in seen_cursors:
95+
raise RuntimeError("Server returned a pagination cursor it already returned; refusing to page forever.")
96+
seen_cursors.add(next_cursor)
9497
cursor = next_cursor
9598

9699

tests/client/test_list_all_pagination.py

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,25 @@ async def handler(_ctx: ServerRequestContext, _params: types.PaginatedRequestPar
7070
return handler
7171

7272

73+
def _cycling_cursor_handler(
74+
make_item: Callable[[str], ItemT],
75+
result_cls: Callable[..., ResultT],
76+
items_field: str,
77+
) -> Callable[[ServerRequestContext, types.PaginatedRequestParams | None], Awaitable[ResultT]]:
78+
"""Build a malformed handler that alternates between two cursors forever.
79+
80+
The cursor always advances relative to the previous page, so a guard that
81+
only compares against the immediately preceding cursor would page forever.
82+
"""
83+
84+
async def handler(_ctx: ServerRequestContext, params: types.PaginatedRequestParams | None) -> ResultT:
85+
cursor = params.cursor if params else None
86+
next_cursor = "b" if cursor == "a" else "a"
87+
return result_cls(**{items_field: [make_item("x")]}, next_cursor=next_cursor)
88+
89+
return handler
90+
91+
7392
def _make_tool(name: str) -> types.Tool:
7493
return types.Tool(name=name, input_schema={"type": "object"})
7594

@@ -269,5 +288,21 @@ async def test_drain_raises_when_cursor_does_not_advance(
269288
server = build_server()
270289

271290
async with Client(server) as client:
272-
with pytest.raises(RuntimeError, match="did not advance"):
291+
with pytest.raises(RuntimeError, match="already returned"):
273292
await getattr(client, client_method)()
293+
294+
295+
async def test_drain_raises_when_cursors_cycle():
296+
"""A server whose cursors cycle (a, b, a, ...) must fail loudly, not loop forever.
297+
298+
Each cursor differs from the one before it, so this specifically exercises
299+
the seen-set guard rather than the simpler stuck-cursor case above.
300+
"""
301+
server = Server(
302+
"cycling-tools",
303+
on_list_tools=_cycling_cursor_handler(_make_tool, types.ListToolsResult, "tools"),
304+
)
305+
306+
async with Client(server) as client:
307+
with pytest.raises(RuntimeError, match="already returned"):
308+
await client.list_all_tools()

tests/client/test_session_group.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -463,5 +463,5 @@ async def test_client_session_group_rejects_non_advancing_cursor(
463463

464464
group = ClientSessionGroup(exit_stack=mock_exit_stack)
465465
with mock.patch.object(group, "_establish_session", return_value=(mock_server_info, mock_session)):
466-
with pytest.raises(RuntimeError, match="did not advance"):
466+
with pytest.raises(RuntimeError, match="already returned"):
467467
await group.connect_to_server(StdioServerParameters(command="test"))

0 commit comments

Comments
 (0)