Skip to content

Commit 36a8bdd

Browse files
committed
Extend the non-advancing cursor guard to ClientSessionGroup drains
The pagination guide says every drain fails loudly when a server returns the same cursor it was handed, but the group aggregator's drain loop had no such guard and would page forever. Guard it the same way the Client drains are guarded.
1 parent 36ea094 commit 36a8bdd

2 files changed

Lines changed: 28 additions & 2 deletions

File tree

src/mcp/client/session_group.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,9 @@ async def _drain_paginated(
7676
`fetch_page` is one of the ClientSession `list_*` methods that takes a
7777
`params=PaginatedRequestParams(...)` keyword. `attribute` is the name of
7878
the list attribute on the result (e.g. `"tools"`, `"prompts"`).
79+
80+
Raises:
81+
RuntimeError: The server returned a pagination cursor that did not advance.
7982
"""
8083
items: list[Any] = []
8184
cursor: str | None = None
@@ -86,6 +89,8 @@ async def _drain_paginated(
8689
next_cursor = getattr(result, "next_cursor", None)
8790
if next_cursor is None:
8891
return items
92+
if next_cursor == cursor:
93+
raise RuntimeError("Server returned a pagination cursor that did not advance; refusing to page forever.")
8994
cursor = next_cursor
9095

9196

@@ -435,8 +440,8 @@ async def _aggregate_components(self, server_info: types.Implementation, session
435440
tools_temp[name] = tool
436441
tool_to_session_temp[name] = session
437442
component_names.tools.add(name)
438-
except MCPError as err: # pragma: no cover
439-
logging.warning(f"Could not fetch tools: {err}")
443+
except MCPError as err:
444+
logging.warning(f"Could not fetch tools: {err}") # pragma: no cover
440445

441446
# Clean up exit stack for session if we couldn't retrieve anything
442447
# from the server.

tests/client/test_session_group.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -444,3 +444,24 @@ async def test_client_session_group_aggregates_paginated_tools(
444444
second_call_kwargs = mock_session.list_tools.await_args_list[1].kwargs
445445
assert second_call_kwargs["params"] is not None
446446
assert second_call_kwargs["params"].cursor == "page-2"
447+
448+
449+
@pytest.mark.anyio
450+
async def test_client_session_group_rejects_non_advancing_cursor(
451+
mock_exit_stack: contextlib.AsyncExitStack,
452+
):
453+
"""A server that keeps returning the cursor it was handed must fail loudly, not page forever."""
454+
mock_server_info = mock.Mock(spec=types.Implementation)
455+
mock_server_info.name = "StuckCursorServer"
456+
mock_session = mock.AsyncMock(spec=mcp.ClientSession)
457+
458+
stuck_tool = mock.Mock(spec=types.Tool)
459+
stuck_tool.name = "tool_a"
460+
mock_session.list_tools.return_value = mock.AsyncMock(tools=[stuck_tool], next_cursor="page-2")
461+
mock_session.list_resources.return_value = mock.AsyncMock(resources=[], next_cursor=None)
462+
mock_session.list_prompts.return_value = mock.AsyncMock(prompts=[], next_cursor=None)
463+
464+
group = ClientSessionGroup(exit_stack=mock_exit_stack)
465+
with mock.patch.object(group, "_establish_session", return_value=(mock_server_info, mock_session)):
466+
with pytest.raises(RuntimeError, match="did not advance"):
467+
await group.connect_to_server(StdioServerParameters(command="test"))

0 commit comments

Comments
 (0)