fix(client): drain paginated session-group catalogs

Author: namtran1812

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.

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.

tests/client/test_session_group.py

@@ -1,18 +1,21 @@
 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,
     SseServerParameters,
     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"]