Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions docs/client/session-groups.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
55 changes: 34 additions & 21 deletions src/mcp/client/session_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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."""
Expand Down Expand Up @@ -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]:
Expand Down Expand Up @@ -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.
Expand Down
201 changes: 192 additions & 9 deletions tests/client/test_session_group.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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"]
Loading