From 5732e7d350708ef2b8f4c339b41e9f49f9117c34 Mon Sep 17 00:00:00 2001 From: namtran1812 <158846154+namtran1812@users.noreply.github.com> Date: Sat, 31 Oct 2026 13:34:39 -0400 Subject: [PATCH] fix(client): drain paginated session-group catalogs --- docs/client/session-groups.md | 11 ++ src/mcp/client/session_group.py | 55 +++++--- tests/client/test_session_group.py | 201 +++++++++++++++++++++++++++-- 3 files changed, 237 insertions(+), 30 deletions(-) diff --git a/docs/client/session-groups.md b/docs/client/session-groups.md index de43e74890..aa2947f3e5 100644 --- a/docs/client/session-groups.md +++ b/docs/client/session-groups.md @@ -66,6 +66,17 @@ Run it again. `print(sorted(group.tools))` now shows both: If you already hold a connected `ClientSession` (`Client.session` is one), hand it to `await group.connect_with_session(server_info, session)` instead of opening a new transport. It aggregates the same way. The group never closes a session it didn't open. `server_info` names the server for component prefixes; on a 2026-era connection `client.server_info` can be `None` (identity is optional), so pass your own `Implementation(name=..., version=...)` in that case. +## Paginated catalogs + +Connecting a server follows every continuation cursor in its tools, resources, and +prompts listings. Naming hooks and collision checks apply to the complete catalog, +and tools discovered on later pages route to the same owning session. + +A failure on a continuation page, a repeated cursor, or a name collision aborts +registration without adding any of that server's components to the group. Existing +servers remain registered. A first-page listing error still leaves that component +kind unavailable while allowing the other kinds to be discovered. + ## The classic handshake `ClientSessionGroup` is built on `ClientSession`, not on `Client`. Each `connect_to_server` runs the classic `initialize` handshake. It never sends the `server/discover` probe described in **[Protocol versions](../protocol-versions.md)**. Every MCP server understands that handshake, so this costs you compatibility with nothing; it only means a group takes the older, slower path to a server that could do better. diff --git a/src/mcp/client/session_group.py b/src/mcp/client/session_group.py index 13d8928ed4..c4e7423177 100644 --- a/src/mcp/client/session_group.py +++ b/src/mcp/client/session_group.py @@ -8,10 +8,10 @@ import contextlib import logging -from collections.abc import Callable +from collections.abc import AsyncIterator, Awaitable, Callable from dataclasses import dataclass from types import TracebackType -from typing import Any, Literal, Self, TypeAlias, overload +from typing import Any, Literal, Self, TypeAlias, TypeVar, overload import anyio import httpx2 @@ -27,6 +27,31 @@ from mcp.shared.dispatcher import ProgressFnT from mcp.shared.exceptions import MCPError +_ListingResult = TypeVar("_ListingResult", bound=types.PaginatedResult) + + +async def _list_pages(list_page: Callable[..., Awaitable[_ListingResult]], kind: str) -> AsyncIterator[_ListingResult]: + cursor: str | None = None + seen_cursors: set[str] = set() + while True: + try: + if cursor is None: + page = await list_page() + else: + page = await list_page(params=types.PaginatedRequestParams(cursor=cursor)) + except MCPError as err: + if cursor is not None: + raise + logging.warning(f"Could not fetch {kind}: {err}") + return + yield page + cursor = page.next_cursor + if cursor is None: + return + if cursor in seen_cursors: + raise MCPError(code=types.INVALID_PARAMS, message=f"Repeated {kind} pagination cursor") + seen_cursors.add(cursor) + class SseServerParameters(BaseModel): """Parameters for initializing an sse_client.""" @@ -175,7 +200,7 @@ async def __aexit__( @property def sessions(self) -> list[mcp.ClientSession]: """Returns the list of sessions being managed.""" - return list(self._sessions.keys()) # pragma: no cover + return list(self._sessions.keys()) @property def prompts(self) -> dict[str, types.Prompt]: @@ -386,36 +411,24 @@ async def _aggregate_components(self, server_info: types.Implementation, session tools_temp: dict[str, types.Tool] = {} tool_to_session_temp: dict[str, mcp.ClientSession] = {} - # Query the server for its prompts and aggregate to list. - try: - prompts = (await session.list_prompts()).prompts - for prompt in prompts: + async for prompt_page in _list_pages(session.list_prompts, "prompts"): + for prompt in prompt_page.prompts: name = self._component_name(prompt.name, server_info) prompts_temp[name] = prompt component_names.prompts.add(name) - except MCPError as err: # pragma: no cover - logging.warning(f"Could not fetch prompts: {err}") - # Query the server for its resources and aggregate to list. - try: - resources = (await session.list_resources()).resources - for resource in resources: + async for resource_page in _list_pages(session.list_resources, "resources"): + for resource in resource_page.resources: name = self._component_name(resource.name, server_info) resources_temp[name] = resource component_names.resources.add(name) - except MCPError as err: # pragma: no cover - logging.warning(f"Could not fetch resources: {err}") - # Query the server for its tools and aggregate to list. - try: - tools = (await session.list_tools()).tools - for tool in tools: + async for tool_page in _list_pages(session.list_tools, "tools"): + for tool in tool_page.tools: name = self._component_name(tool.name, server_info) tools_temp[name] = tool tool_to_session_temp[name] = session component_names.tools.add(name) - except MCPError as err: # pragma: no cover - logging.warning(f"Could not fetch tools: {err}") # Clean up exit stack for session if we couldn't retrieve anything # from the server. diff --git a/tests/client/test_session_group.py b/tests/client/test_session_group.py index 922e1eebcb..182c1f296c 100644 --- a/tests/client/test_session_group.py +++ b/tests/client/test_session_group.py @@ -1,11 +1,13 @@ import contextlib from unittest import mock +import anyio import httpx2 import mcp_types as types import pytest import mcp +from mcp import Client from mcp.client.session_group import ( ClientSessionGroup, ClientSessionParameters, @@ -13,6 +15,7 @@ StreamableHttpParameters, ) from mcp.client.stdio import StdioServerParameters +from mcp.server import Server, ServerRequestContext from mcp.shared.exceptions import MCPError @@ -116,9 +119,9 @@ async def test_client_session_group_connect_to_server(mock_exit_stack: contextli mock_resource1.name = "resource_b" mock_prompt1 = mock.Mock(spec=types.Prompt) mock_prompt1.name = "prompt_c" - mock_session.list_tools.return_value = mock.AsyncMock(tools=[mock_tool1]) - mock_session.list_resources.return_value = mock.AsyncMock(resources=[mock_resource1]) - mock_session.list_prompts.return_value = mock.AsyncMock(prompts=[mock_prompt1]) + mock_session.list_tools.return_value = mock.AsyncMock(tools=[mock_tool1], next_cursor=None) + mock_session.list_resources.return_value = mock.AsyncMock(resources=[mock_resource1], next_cursor=None) + mock_session.list_prompts.return_value = mock.AsyncMock(prompts=[mock_prompt1], next_cursor=None) # --- Test Execution --- group = ClientSessionGroup(exit_stack=mock_exit_stack) @@ -151,9 +154,9 @@ async def test_client_session_group_connect_to_server_with_name_hook(mock_exit_s mock_session = mock.AsyncMock(spec=mcp.ClientSession) mock_tool = mock.Mock(spec=types.Tool) mock_tool.name = "base_tool" - mock_session.list_tools.return_value = mock.AsyncMock(tools=[mock_tool]) - mock_session.list_resources.return_value = mock.AsyncMock(resources=[]) - mock_session.list_prompts.return_value = mock.AsyncMock(prompts=[]) + mock_session.list_tools.return_value = mock.AsyncMock(tools=[mock_tool], next_cursor=None) + mock_session.list_resources.return_value = mock.AsyncMock(resources=[], next_cursor=None) + mock_session.list_prompts.return_value = mock.AsyncMock(prompts=[], next_cursor=None) # --- Test Setup --- def name_hook(name: str, server_info: types.Implementation) -> str: @@ -262,10 +265,10 @@ async def test_client_session_group_connect_to_server_duplicate_tool_raises_erro # Configure the new session to return a tool with the *same name* duplicate_tool = mock.Mock(spec=types.Tool) duplicate_tool.name = existing_tool_name - mock_session_new.list_tools.return_value = mock.AsyncMock(tools=[duplicate_tool]) + mock_session_new.list_tools.return_value = mock.AsyncMock(tools=[duplicate_tool], next_cursor=None) # Keep other lists empty for simplicity - mock_session_new.list_resources.return_value = mock.AsyncMock(resources=[]) - mock_session_new.list_prompts.return_value = mock.AsyncMock(prompts=[]) + mock_session_new.list_resources.return_value = mock.AsyncMock(resources=[], next_cursor=None) + mock_session_new.list_prompts.return_value = mock.AsyncMock(prompts=[], next_cursor=None) # --- Test Execution and Assertion --- with pytest.raises(MCPError) as excinfo: @@ -405,3 +408,183 @@ async def test_client_session_group_establish_session_parameterized( # 3. Assert returned values assert returned_server_info is mock_initialize_result.server_info assert returned_session is mock_entered_session + + +@pytest.mark.anyio +@pytest.mark.parametrize("prefixed", [False, True]) +async def test_group_discovers_all_pages_routes_later_tools_and_removes_them(prefixed: bool) -> None: + """SDK-defined: aggregation drains each catalog and keeps later-page ownership for routing and removal.""" + cursors: dict[str, list[str | None]] = {kind: [] for kind in ("tools", "resources", "prompts")} + + async def tools(ctx: ServerRequestContext, params: types.PaginatedRequestParams | None) -> types.ListToolsResult: + cursor = params.cursor if params else None + cursors["tools"].append(cursor) + return types.ListToolsResult( + tools=[types.Tool(name="first" if cursor is None else "second", input_schema={"type": "object"})], + next_cursor="tools-2" if cursor is None else None, + ) + + async def resources( + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None + ) -> types.ListResourcesResult: + cursor = params.cursor if params else None + cursors["resources"].append(cursor) + name = "first" if cursor is None else "second" + return types.ListResourcesResult( + resources=[types.Resource(name=name, uri=f"memory://{name}")], + next_cursor="resources-2" if cursor is None else None, + ) + + async def prompts( + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None + ) -> types.ListPromptsResult: + cursor = params.cursor if params else None + cursors["prompts"].append(cursor) + return types.ListPromptsResult( + prompts=[types.Prompt(name="first" if cursor is None else "second")], + next_cursor="prompts-2" if cursor is None else None, + ) + + async def call(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult: + assert params.name == "second" + return types.CallToolResult(content=[types.TextContent(text=params.name)]) + + server = Server( + "catalog", on_list_tools=tools, on_list_resources=resources, on_list_prompts=prompts, on_call_tool=call + ) + + def hook(name: str, info: types.Implementation) -> str: + return f"{info.name}.{name}" + + group = ClientSessionGroup(component_name_hook=hook if prefixed else None) + with anyio.fail_after(5): + async with Client(server) as client: + session = client.session + await group.connect_with_session(types.Implementation(name="catalog", version="1"), session) + expected = {"catalog.first", "catalog.second"} if prefixed else {"first", "second"} + assert set(group.tools) == expected + assert set(group.resources) == expected + assert set(group.prompts) == expected + result = await group.call_tool("catalog.second" if prefixed else "second") + assert result.content == [types.TextContent(text="second")] + await group.disconnect_from_server(session) + assert group.tools == {} + assert group.resources == {} + assert group.prompts == {} + assert group.sessions == [] + assert cursors == {kind: [None, f"{kind}-2"] for kind in cursors} + + +@pytest.mark.anyio +@pytest.mark.parametrize("kind", ["tools", "resources", "prompts"]) +@pytest.mark.parametrize("repeated_cursor", [False, True]) +async def test_failed_later_catalog_page_leaves_the_group_unchanged(kind: str, repeated_cursor: bool) -> None: + """SDK-defined: later-page errors and cursor cycles cannot publish a partial server catalog.""" + calls: dict[str, list[str | None]] = {key: [] for key in ("tools", "resources", "prompts")} + + def cursor_for(key: str, params: types.PaginatedRequestParams | None) -> str | None: + cursor = params.cursor if params else None + calls[key].append(cursor) + if key != kind: + return None + if cursor is not None and not repeated_cursor: + raise MCPError(code=types.INTERNAL_ERROR, message="page failed") + return "repeat" + + async def tools(ctx: ServerRequestContext, params: types.PaginatedRequestParams | None) -> types.ListToolsResult: + return types.ListToolsResult( + tools=[types.Tool(name="candidate", input_schema={"type": "object"})], + next_cursor=cursor_for("tools", params), + ) + + async def resources( + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None + ) -> types.ListResourcesResult: + return types.ListResourcesResult( + resources=[types.Resource(name="candidate", uri="memory://candidate")], + next_cursor=cursor_for("resources", params), + ) + + async def prompts( + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None + ) -> types.ListPromptsResult: + return types.ListPromptsResult( + prompts=[types.Prompt(name="candidate")], next_cursor=cursor_for("prompts", params) + ) + + async def existing_tools( + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None + ) -> types.ListToolsResult: + return types.ListToolsResult(tools=[types.Tool(name="existing", input_schema={"type": "object"})]) + + existing = Server("existing", on_list_tools=existing_tools) + candidate = Server("candidate", on_list_tools=tools, on_list_resources=resources, on_list_prompts=prompts) + group = ClientSessionGroup() + with anyio.fail_after(5): + async with Client(existing) as first, Client(candidate) as second: + await group.connect_with_session(types.Implementation(name="existing", version="1"), first.session) + before = (dict(group.tools), dict(group.resources), dict(group.prompts), list(group.sessions)) + with pytest.raises(MCPError) as error: + await group.connect_with_session(types.Implementation(name="candidate", version="1"), second.session) + assert error.value.code == (types.INVALID_PARAMS if repeated_cursor else types.INTERNAL_ERROR) + assert (group.tools, group.resources, group.prompts, group.sessions) == before + assert calls[kind] == [None, "repeat"] + + +@pytest.mark.anyio +async def test_tool_collision_on_a_later_page_does_not_register_the_candidate_server() -> None: + """SDK-defined: cross-server collisions are checked against the complete paginated catalog.""" + + async def existing_tools( + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None + ) -> types.ListToolsResult: + return types.ListToolsResult(tools=[types.Tool(name="existing", input_schema={"type": "object"})]) + + async def candidate_tools( + ctx: ServerRequestContext, params: types.PaginatedRequestParams | None + ) -> types.ListToolsResult: + later = params is not None and params.cursor is not None + return types.ListToolsResult( + tools=[types.Tool(name="existing" if later else "new", input_schema={"type": "object"})], + next_cursor=None if later else "second", + ) + + group = ClientSessionGroup() + with anyio.fail_after(5): + async with ( + Client(Server("existing", on_list_tools=existing_tools)) as first, + Client(Server("candidate", on_list_tools=candidate_tools)) as second, + ): + await group.connect_with_session(types.Implementation(name="existing", version="1"), first.session) + with pytest.raises(MCPError) as error: + await group.connect_with_session(types.Implementation(name="candidate", version="1"), second.session) + assert error.value.code == types.INVALID_PARAMS + assert set(group.tools) == {"existing"} + assert group.sessions == [first.session] + + +@pytest.mark.anyio +async def test_group_follows_an_empty_cursor_and_empty_pages_until_cursor_is_absent() -> None: + """SDK-defined: opaque cursors, not page contents or cursor truthiness, decide when listing ends.""" + cursors: list[str | None] = [] + + async def tools(ctx: ServerRequestContext, params: types.PaginatedRequestParams | None) -> types.ListToolsResult: + cursor = params.cursor if params else None + cursors.append(cursor) + if cursor is None: + return types.ListToolsResult(tools=[], next_cursor="") + if cursor == "": + return types.ListToolsResult( + tools=[types.Tool(name="found", input_schema={"type": "object"})], next_cursor="last" + ) + assert cursor == "last" + return types.ListToolsResult(tools=[]) + + group = ClientSessionGroup() + with anyio.fail_after(5): + async with Client(Server("empty-pages", on_list_tools=tools)) as client: + await group.connect_with_session(types.Implementation(name="empty-pages", version="1"), client.session) + assert set(group.tools) == {"found"} + await group.disconnect_from_server(client.session) + assert group.tools == {} + assert cursors == [None, "", "last"]