From 7d574e1933fe7d1fd65f1eea99e73567f50742ff Mon Sep 17 00:00:00 2001 From: Basel Arw <100733493+baselarw@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:27:27 +0300 Subject: [PATCH] feat(session-group): add read_resource and get_prompt routing ClientSessionGroup could aggregate resources and prompts and call_tool against them, but offered no routed way to actually read a resource or get a prompt: callers had to track the owning session themselves. Add read_resource(name, ...) and get_prompt(name, arguments, ...) mirroring call_tool: they resolve the aggregate key (honoring component_name_hook) to the owning session via new _resource_to_session / _prompt_to_session reverse indexes and forward the resource's wire URI / prompt's wire name. Both carry the same allow_input_required overloads as ClientSession. Reverse indexes are cleaned up on disconnect_from_server. Adds unit + in-memory end-to-end tests and documents the methods in the session-groups guide. --- docs/client/session-groups.md | 16 +++- src/mcp/client/session_group.py | 122 ++++++++++++++++++++++++++ tests/client/test_session_group.py | 136 +++++++++++++++++++++++++++++ 3 files changed, 273 insertions(+), 1 deletion(-) diff --git a/docs/client/session-groups.md b/docs/client/session-groups.md index de43e74890..c4773e4977 100644 --- a/docs/client/session-groups.md +++ b/docs/client/session-groups.md @@ -60,6 +60,20 @@ Run it again. `print(sorted(group.tools))` now shows both: The hook runs on **every** name from **every** server, not only on conflicts: there is no prefix-on-collision mode. Pick one scheme and let it apply everywhere. +## Reading resources and prompts + +`call_tool` is not the only routed call. `group.read_resource(name)` and `group.get_prompt(name, arguments)` work the same way: you pass the aggregate key (the one in `group.resources` / `group.prompts`, prefixed if you use a hook), and the group finds the owning session and forwards the call with the resource's real URI or the prompt's real name. + +```python +# `Library.hours` is the group key; `library://hours` goes on the wire. +result = await group.read_resource("Library.hours") + +# `arguments` are forwarded to the owning server unchanged. +prompt = await group.get_prompt("Greeter.greet", {"name": "Ada"}) +``` + +Both raise `KeyError` if the key isn't an aggregated resource/prompt, and both accept the same `allow_input_required` flag as `ClientSession`, so a server that needs input mid-call behaves identically through the group. + ## Adding and removing servers `connect_to_server` returns the `ClientSession` it opened. Keep it if you ever want that server gone: `await group.disconnect_from_server(session)` removes its tools, resources, and prompts from the group. @@ -74,7 +88,7 @@ If you already hold a connected `ClientSession` (`Client.session` is one), hand * `ClientSessionGroup` holds many server connections and merges their tools, resources, and prompts into one `dict` each. * `connect_to_server(params)` per server. It takes transport parameters, never the URL or `Transport` a `Client` takes. -* `group.call_tool(name, arguments)` routes to the owning server for you. +* `group.call_tool(name, arguments)` routes to the owning server for you; `group.read_resource(name)` and `group.get_prompt(name, arguments)` route the same way. * Names must be unique across the whole group; two servers with a `search` tool cannot coexist on their own. * `component_name_hook=` rewrites every registered name. The dict key changes, the wire name does not. * `connect_with_session` adds a session you already hold; `disconnect_from_server` removes one. diff --git a/src/mcp/client/session_group.py b/src/mcp/client/session_group.py index a544cecbe8..c6450eb7e8 100644 --- a/src/mcp/client/session_group.py +++ b/src/mcp/client/session_group.py @@ -116,6 +116,8 @@ class _ComponentNames(BaseModel): # Client-server connection management. _sessions: dict[mcp.ClientSession, _ComponentNames] _tool_to_session: dict[str, mcp.ClientSession] + _resource_to_session: dict[str, mcp.ClientSession] + _prompt_to_session: dict[str, mcp.ClientSession] _exit_stack: contextlib.AsyncExitStack _session_exit_stacks: dict[mcp.ClientSession, contextlib.AsyncExitStack] @@ -138,6 +140,8 @@ def __init__( self._sessions = {} self._tool_to_session = {} + self._resource_to_session = {} + self._prompt_to_session = {} if exit_stack is None: self._exit_stack = contextlib.AsyncExitStack() self._owns_exit_stack = True @@ -249,6 +253,112 @@ async def call_tool( allow_input_required=allow_input_required, ) + @overload + async def read_resource( + self, + name: str, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: types.RequestParamsMeta | None = None, + allow_input_required: Literal[False] = False, + ) -> types.ReadResourceResult: ... + + @overload + async def read_resource( + self, + name: str, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: types.RequestParamsMeta | None = None, + allow_input_required: bool, + ) -> types.ReadResourceResult | types.InputRequiredResult: ... + + async def read_resource( + self, + name: str, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: types.RequestParamsMeta | None = None, + allow_input_required: bool = False, + ) -> types.ReadResourceResult | types.InputRequiredResult: + """Reads an aggregated resource, routing to the server that owns it. + + ``name`` is the aggregate key (i.e. the value used in ``resources``, + which is affected by ``component_name_hook``), not necessarily the + resource's wire URI. + + Raises: + KeyError: If ``name`` is not an aggregated resource. + RuntimeError: If the server returns an ``InputRequiredResult`` and + ``allow_input_required`` is ``False``. + """ + session = self._resource_to_session[name] + return await session.read_resource( + self.resources[name].uri, + input_responses=input_responses, + request_state=request_state, + meta=meta, + allow_input_required=allow_input_required, + ) + + @overload + async def get_prompt( + self, + name: str, + arguments: dict[str, str] | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: types.RequestParamsMeta | None = None, + allow_input_required: Literal[False] = False, + ) -> types.GetPromptResult: ... + + @overload + async def get_prompt( + self, + name: str, + arguments: dict[str, str] | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: types.RequestParamsMeta | None = None, + allow_input_required: bool, + ) -> types.GetPromptResult | types.InputRequiredResult: ... + + async def get_prompt( + self, + name: str, + arguments: dict[str, str] | None = None, + *, + input_responses: types.InputResponses | None = None, + request_state: str | None = None, + meta: types.RequestParamsMeta | None = None, + allow_input_required: bool = False, + ) -> types.GetPromptResult | types.InputRequiredResult: + """Gets an aggregated prompt, routing to the server that owns it. + + ``name`` is the aggregate key (i.e. the value used in ``prompts``, + which is affected by ``component_name_hook``), not necessarily the + prompt's wire name. + + Raises: + KeyError: If ``name`` is not an aggregated prompt. + RuntimeError: If the server returns an ``InputRequiredResult`` and + ``allow_input_required`` is ``False``. + """ + session = self._prompt_to_session[name] + return await session.get_prompt( + self.prompts[name].name, + arguments, + input_responses=input_responses, + request_state=request_state, + meta=meta, + allow_input_required=allow_input_required, + ) + async def disconnect_from_server(self, session: mcp.ClientSession) -> None: """Disconnects from a single MCP server.""" @@ -272,6 +382,12 @@ async def disconnect_from_server(self, session: mcp.ClientSession) -> None: for name in component_names.resources: if name in self._resources: # pragma: no branch del self._resources[name] + if name in self._resource_to_session: # pragma: no branch + del self._resource_to_session[name] + # Remove prompts' reverse index for this session. + for name in component_names.prompts: + if name in self._prompt_to_session: # pragma: no branch + del self._prompt_to_session[name] # Remove tools associated with the session. for name in component_names.tools: if name in self._tools: # pragma: no branch @@ -382,6 +498,8 @@ async def _aggregate_components(self, server_info: types.Implementation, session resources_temp: dict[str, types.Resource] = {} tools_temp: dict[str, types.Tool] = {} tool_to_session_temp: dict[str, mcp.ClientSession] = {} + resource_to_session_temp: dict[str, mcp.ClientSession] = {} + prompt_to_session_temp: dict[str, mcp.ClientSession] = {} # Query the server for its prompts and aggregate to list. try: @@ -389,6 +507,7 @@ async def _aggregate_components(self, server_info: types.Implementation, session for prompt in prompts: name = self._component_name(prompt.name, server_info) prompts_temp[name] = prompt + prompt_to_session_temp[name] = session component_names.prompts.add(name) except MCPError as err: # pragma: no cover logging.warning(f"Could not fetch prompts: {err}") @@ -399,6 +518,7 @@ async def _aggregate_components(self, server_info: types.Implementation, session for resource in resources: name = self._component_name(resource.name, server_info) resources_temp[name] = resource + resource_to_session_temp[name] = session component_names.resources.add(name) except MCPError as err: # pragma: no cover logging.warning(f"Could not fetch resources: {err}") @@ -442,6 +562,8 @@ async def _aggregate_components(self, server_info: types.Implementation, session self._resources.update(resources_temp) self._tools.update(tools_temp) self._tool_to_session.update(tool_to_session_temp) + self._resource_to_session.update(resource_to_session_temp) + self._prompt_to_session.update(prompt_to_session_temp) def _component_name(self, name: str, server_info: types.Implementation) -> str: if self._component_name_hook: diff --git a/tests/client/test_session_group.py b/tests/client/test_session_group.py index b75d22b7a0..630b8a8878 100644 --- a/tests/client/test_session_group.py +++ b/tests/client/test_session_group.py @@ -6,6 +6,7 @@ import pytest import mcp +from mcp import Client from mcp.client.session_group import ( ClientSessionGroup, ClientSessionParameters, @@ -13,6 +14,7 @@ StreamableHttpParameters, ) from mcp.client.stdio import StdioServerParameters +from mcp.server import MCPServer from mcp.shared.exceptions import MCPError @@ -402,3 +404,137 @@ 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 +async def test_read_resource_routes_to_owning_session(): + """read_resource resolves the aggregate key to the owning session and passes the wire URI.""" + mock_session = mock.AsyncMock(spec=mcp.ClientSession) + resource = types.Resource(name="hours", uri="library://hours") + expected = types.ReadResourceResult(contents=[]) + mock_session.read_resource.return_value = expected + + group = ClientSessionGroup() + group._resources = {"Library.hours": resource} + group._resource_to_session = {"Library.hours": mock_session} + + result = await group.read_resource("Library.hours") + + assert result is expected + mock_session.read_resource.assert_awaited_once_with( + "library://hours", + input_responses=None, + request_state=None, + meta=None, + allow_input_required=False, + ) + + +@pytest.mark.anyio +async def test_read_resource_unknown_name_raises_key_error(): + group = ClientSessionGroup() + with pytest.raises(KeyError): + await group.read_resource("missing") + + +@pytest.mark.anyio +async def test_get_prompt_routes_to_owning_session(): + """get_prompt resolves the aggregate key to the owning session and passes the wire name.""" + mock_session = mock.AsyncMock(spec=mcp.ClientSession) + prompt = types.Prompt(name="greet") + expected = types.GetPromptResult(messages=[]) + mock_session.get_prompt.return_value = expected + + group = ClientSessionGroup() + group._prompts = {"Web.greet": prompt} + group._prompt_to_session = {"Web.greet": mock_session} + + result = await group.get_prompt("Web.greet", {"name": "Ada"}) + + assert result is expected + mock_session.get_prompt.assert_awaited_once_with( + "greet", + {"name": "Ada"}, + input_responses=None, + request_state=None, + meta=None, + allow_input_required=False, + ) + + +@pytest.mark.anyio +async def test_get_prompt_unknown_name_raises_key_error(): + group = ClientSessionGroup() + with pytest.raises(KeyError): + await group.get_prompt("missing") + + +@pytest.mark.anyio +async def test_disconnect_clears_resource_and_prompt_reverse_index(): + """disconnect_from_server drops the session's resource/prompt routing entries.""" + session = mock.Mock(spec=mcp.ClientSession) + group = ClientSessionGroup() + group._resources = {"res1": mock.Mock(spec=types.Resource)} + group._prompts = {"prm1": mock.Mock(spec=types.Prompt)} + group._resource_to_session = {"res1": session} + group._prompt_to_session = {"prm1": session} + group._sessions = { + session: ClientSessionGroup._ComponentNames( + prompts={"prm1"}, + resources={"res1"}, + tools=set(), + ) + } + + await group.disconnect_from_server(session) + + assert "res1" not in group._resource_to_session + assert "prm1" not in group._prompt_to_session + + +def _server_info(client: Client) -> types.Implementation: + assert client.server_info is not None + return client.server_info + + +@pytest.mark.anyio +async def test_read_resource_end_to_end_routes_through_the_owning_server(): + """The group reads an aggregated resource against a real in-memory session.""" + server = MCPServer("Library") + + @server.resource("library://hours") + def hours() -> str: + return "Mon-Fri 09:00-17:00" + + async with Client(server) as client: + group = ClientSessionGroup() + await group.connect_with_session(_server_info(client), client.session) + + (name,) = group.resources + result = await group.read_resource(name) + + assert isinstance(result, types.ReadResourceResult) + (content,) = result.contents + assert isinstance(content, types.TextResourceContents) + assert content.text == "Mon-Fri 09:00-17:00" + + +@pytest.mark.anyio +async def test_get_prompt_end_to_end_routes_through_the_owning_server(): + """The group gets an aggregated prompt against a real in-memory session.""" + server = MCPServer("Greeter") + + @server.prompt() + def greet(name: str) -> str: + return f"Hello, {name}!" + + async with Client(server) as client: + group = ClientSessionGroup() + await group.connect_with_session(_server_info(client), client.session) + + result = await group.get_prompt("greet", {"name": "Ada"}) + + assert isinstance(result, types.GetPromptResult) + (message,) = result.messages + assert isinstance(message.content, types.TextContent) + assert message.content.text == "Hello, Ada!"