Create ClientSessionGroup for managing multiple session connections.

This abstraction concurrently manages multiple MCP session connections.
Tools, resources, and prompts are aggregated across servers. Servers may
be connected to or disconnected from at any point after initialization.

This abstractions can handle naming collisions using a custom
user-provided hook.

Author: mkeid

src/mcp/__init__.py

@@ -1,4 +1,5 @@
 from .client.session import ClientSession
+from .client.session_group import ClientSessionGroup
 from .client.stdio import StdioServerParameters, stdio_client
 from .server.session import ServerSession
 from .server.stdio import stdio_server
@@ -63,6 +64,7 @@
     "ClientRequest",
     "ClientResult",
     "ClientSession",
+    "ClientSessionGroup",
     "CreateMessageRequest",
     "CreateMessageResult",
     "ErrorData",

src/mcp/client/session_group.py

@@ -0,0 +1,207 @@
+"""
+SessionGroup concurrently manages multiple MCP session connections.
+
+Tools, resources, and prompts are aggregated across servers. Servers may
+be connected to or disconnected from at any point after initialization.
+
+This abstractions can handle naming collisions using a custom user-provided
+hook.
+"""
+
+import contextlib
+from collections.abc import Callable
+from typing import Any, TypeAlias
+
+from pydantic import BaseModel
+
+import mcp
+from mcp import types
+from mcp.client.stdio import StdioServerParameters
+from mcp.shared.exceptions import McpError
+
+
+class ClientSessionGroup:
+    """Client for managing connections to multiple MCP servers.
+
+    This class is responsible for encapsulating management of server connections.
+    It it aggregates tools, resources, and prompts from all connected servers.
+
+    For auxiliary handlers, such as resource subscription, this is delegated to
+    the client and can be accessed via the session. For example:
+      mcp_session_group.get_session("server_name").subscribe_to_resource(...)
+    """
+
+    class _ComponentNames(BaseModel):
+        """Used for reverse index to find components."""
+
+        prompts: set[str] = set()
+        resources: set[str] = set()
+        tools: set[str] = set()
+
+    # Standard MCP components.
+    _prompts: dict[str, types.Prompt]
+    _resources: dict[str, types.Resource]
+    _tools: dict[str, types.Tool]
+
+    # Client-server connection management.
+    _sessions: dict[mcp.ClientSession, _ComponentNames]
+    _tool_to_session: dict[str, mcp.ClientSession]
+    _exit_stack: contextlib.AsyncExitStack
+
+    # Optional fn consuming (component_name, serverInfo) for custom names.
+    # This is provide a means to mitigate naming conflicts across servers.
+    # Example: (tool_name, serverInfo) => "{result.serverInfo.name}.{tool_name}"
+    _ComponentNameHook: TypeAlias = Callable[[str, types.Implementation], str]
+    _component_name_hook: _ComponentNameHook | None
+
+    def __init__(
+        self,
+        exit_stack: contextlib.AsyncExitStack = contextlib.AsyncExitStack(),
+        component_name_hook: _ComponentNameHook | None = None,
+    ) -> None:
+        """Initializes the MCP client."""
+
+        self._tools = {}
+        self._resources = {}
+        self._prompts = {}
+
+        self._sessions = {}
+        self._tool_to_session = {}
+        self._exit_stack = exit_stack
+        self._component_name_hook = component_name_hook
+
+    @property
+    def prompts(self) -> dict[str, types.Prompt]:
+        """Returns the prompts as a dictionary of names to prompts."""
+        return self._prompts
+
+    @property
+    def resources(self) -> dict[str, types.Resource]:
+        """Returns the resources as a dictionary of names to resources."""
+        return self._resources
+
+    @property
+    def tools(self) -> dict[str, types.Tool]:
+        """Returns the tools as a dictionary of names to tools."""
+        return self._tools
+
+    async def call_tool(self, name: str, args: dict[str, Any]) -> types.CallToolResult:
+        """Executes a tool given its name and arguments."""
+        session = self._tool_to_session[name]
+        return await session.call_tool(name, args)
+
+    def disconnect_from_server(self, session: mcp.ClientSession) -> None:
+        """Disconnects from a single MCP server."""
+
+        if session not in self._sessions:
+            raise McpError(
+                types.ErrorData(
+                    code=types.INVALID_PARAMS,
+                    message="Provided session is not being managed.",
+                )
+            )
+        component_names = self._sessions[session]
+
+        # Remove prompts associated with the session.
+        for name in component_names.prompts:
+            del self._prompts[name]
+
+        # Remove resources associated with the session.
+        for name in component_names.resources:
+            del self._resources[name]
+
+        # Remove tools associated with the session.
+        for name in component_names.tools:
+            del self._tools[name]
+
+        del self._sessions[session]
+
+    async def connect_to_server(
+        self,
+        server_params: StdioServerParameters,
+    ) -> mcp.ClientSession:
+        """Connects to a single MCP server."""
+
+        # Establish server connection and create session.
+        server_info, session = await self._establish_session(server_params)
+
+        # Create a reverse index so we can find all prompts, resources, and
+        # tools belonging to this session. Used for removing components from
+        # the session group via self.disconnect_from_server.
+        component_names = self._ComponentNames()
+
+        # Temporary components dicts. We do not want to modify the aggregate
+        # lists in case of an intermediate failure.
+        prompts_temp: dict[str, types.Prompt] = {}
+        resources_temp: dict[str, types.Resource] = {}
+        tools_temp: dict[str, types.Tool] = {}
+        tool_to_session_temp: dict[str, mcp.ClientSession] = {}
+
+        # Query the server for its prompts and aggregate to list.
+        prompts = (await session.list_prompts()).prompts
+        for prompt in prompts:
+            name = self._component_name(prompt.name, server_info)
+            if name in self._prompts:
+                raise McpError(
+                    types.ErrorData(
+                        code=types.INVALID_PARAMS,
+                        message=f"{name} already exists in group prompts.",
+                    )
+                )
+            prompts_temp[name] = prompt
+            component_names.prompts.add(name)
+
+        # Query the server for its resources and aggregate to list.
+        resources = (await session.list_resources()).resources
+        for resource in resources:
+            name = self._component_name(resource.name, server_info)
+            if name in self._resources:
+                raise McpError(
+                    types.ErrorData(
+                        code=types.INVALID_PARAMS,
+                        message=f"{name} already exists in group resources.",
+                    )
+                )
+            resources_temp[name] = resource
+            component_names.resources.add(name)
+
+        # Query the server for its tools and aggregate to list.
+        tools = (await session.list_tools()).tools
+        for tool in tools:
+            name = self._component_name(tool.name, server_info)
+            if name in self._tools:
+                raise McpError(
+                    types.ErrorData(
+                        code=types.INVALID_PARAMS,
+                        message=f"{name} already exists in group tools.",
+                    )
+                )
+            tools_temp[name] = tool
+            tool_to_session_temp[name] = session
+            component_names.tools.add(name)
+
+        # Aggregate components.
+        self._sessions[session] = component_names
+        self._prompts.update(prompts_temp)
+        self._resources.update(resources_temp)
+        self._tools.update(tools_temp)
+        self._tool_to_session.update(tool_to_session_temp)
+
+        return session
+
+    async def _establish_session(
+        self, server_params: StdioServerParameters
+    ) -> tuple[types.Implementation, mcp.ClientSession]:
+        """Establish a client session to an MCP server."""
+        client = mcp.stdio_client(server_params)
+        read, write = await self._exit_stack.enter_async_context(client)
+        session = await self._exit_stack.enter_async_context(
+            mcp.ClientSession(read, write)
+        )
+        result = await session.initialize()
+        return result.serverInfo, session
+
+    def _component_name(self, name: str, server_info: types.Implementation) -> str:
+        if self._component_name_hook:
+            return self._component_name_hook(name, server_info)
+        return name