Add support for async context management to ClientSessionGroup

This changes enables context management for setting up and tearing down
async exit stacks durring server connection and disconnection respectively.

Documentation has been added to show an example use case that demonstrates
how `ClientSessionGroup` can be used with `async with`.

Author: mkeid

src/mcp/client/session_group.py

@@ -9,11 +9,15 @@
 """
 
 import contextlib
+import logging
 from collections.abc import Callable
 from datetime import timedelta
+from types import TracebackType
 from typing import Any, TypeAlias
 
+import anyio
 from pydantic import BaseModel
+from typing_extensions import Self
 
 import mcp
 from mcp import types
@@ -67,11 +71,18 @@ 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.
+    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(...)
+    the client and can be accessed via the session.
+
+    Example Usage:
+        name_fn = lambda name, server_info: f"{(server_info.name)}-{name}"
+        async with ClientSessionGroup(component_name_hook=name_fn) as group:
+            for server_params in server_params:
+                group.connect_to_server(server_param)
+            ...
+
     """
 
     class _ComponentNames(BaseModel):
@@ -90,6 +101,7 @@ class _ComponentNames(BaseModel):
     _sessions: dict[mcp.ClientSession, _ComponentNames]
     _tool_to_session: dict[str, mcp.ClientSession]
     _exit_stack: contextlib.AsyncExitStack
+    _session_exit_stacks: dict[mcp.ClientSession, contextlib.AsyncExitStack]
 
     # Optional fn consuming (component_name, serverInfo) for custom names.
     # This is provide a means to mitigate naming conflicts across servers.
@@ -99,7 +111,7 @@ class _ComponentNames(BaseModel):
 
     def __init__(
         self,
-        exit_stack: contextlib.AsyncExitStack = contextlib.AsyncExitStack(),
+        exit_stack: contextlib.AsyncExitStack | None = None,
         component_name_hook: _ComponentNameHook | None = None,
     ) -> None:
         """Initializes the MCP client."""
@@ -110,9 +122,43 @@ def __init__(
 
         self._sessions = {}
         self._tool_to_session = {}
-        self._exit_stack = exit_stack
+        if exit_stack is None:
+            self._exit_stack = contextlib.AsyncExitStack()
+            self._owns_exit_stack = True
+        else:
+            self._exit_stack = exit_stack
+            self._owns_exit_stack = False
+        self._session_exit_stacks = {}
         self._component_name_hook = component_name_hook
 
+    async def __aenter__(self) -> Self:
+        # Enter the exit stack only if we created it ourselves
+        if self._owns_exit_stack:
+            await self._exit_stack.__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        _exc_type: type[BaseException] | None,
+        _exc_val: BaseException | None,
+        _exc_tb: TracebackType | None,
+    ) -> bool | None:
+        """Closes session exit stacks and main exit stack upon completion."""
+
+        # Concurrently close session stacks.
+        async with anyio.create_task_group() as tg:
+            for exit_stack in self._session_exit_stacks.values():
+                tg.start_soon(exit_stack.aclose)
+
+        # Only close the main exit stack if we created it
+        if self._owns_exit_stack:
+            await self._exit_stack.aclose()
+
+    @property
+    def sessions(self) -> list[mcp.ClientSession]:
+        """Returns the list of sessions being managed."""
+        return list(self._sessions.keys())
+
     @property
     def prompts(self) -> dict[str, types.Prompt]:
         """Returns the prompts as a dictionary of names to prompts."""
@@ -131,42 +177,113 @@ def tools(self) -> dict[str, types.Tool]:
     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)
+        session_tool_name = self.tools[name].name
+        return await session.call_tool(session_tool_name, args)
 
-    def disconnect_from_server(self, session: mcp.ClientSession) -> None:
+    async def disconnect_from_server(self, session: mcp.ClientSession) -> None:
         """Disconnects from a single MCP server."""
 
-        if session not in self._sessions:
+        session_known_for_components = session in self._sessions
+        session_known_for_stack = session in self._session_exit_stacks
+
+        if not session_known_for_components and not session_known_for_stack:
             raise McpError(
                 types.ErrorData(
                     code=types.INVALID_PARAMS,
-                    message="Provided session is not being managed.",
+                    message="Provided session is not managed or already disconnected.",
                 )
             )
-        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]
+        if session_known_for_components:
+            component_names = self._sessions.pop(session)  # Pop from _sessions tracking
+
+            # Remove prompts associated with the session.
+            for name in component_names.prompts:
+                if name in self._prompts:
+                    del self._prompts[name]
+            # Remove resources associated with the session.
+            for name in component_names.resources:
+                if name in self._resources:
+                    del self._resources[name]
+            # Remove tools associated with the session.
+            for name in component_names.tools:
+                if name in self._tools:
+                    del self._tools[name]
+                if name in self._tool_to_session:
+                    del self._tool_to_session[name]
+
+        # Clean up the session's resources via its dedicated exit stack
+        if session_known_for_stack:
+            session_stack_to_close = self._session_exit_stacks.pop(session)
+            await session_stack_to_close.aclose()
+
+    async def connect_with_session(
+        self, server_info: types.Implementation, session: mcp.ClientSession
+    ) -> mcp.ClientSession:
+        """Connects to a single MCP server."""
+        await self._aggregate_components(server_info, session)
+        return session
 
     async def connect_to_server(
         self,
         server_params: ServerParameters,
     ) -> mcp.ClientSession:
         """Connects to a single MCP server."""
-
-        # Establish server connection and create session.
         server_info, session = await self._establish_session(server_params)
+        return await self.connect_with_session(server_info, session)
+
+    async def _establish_session(
+        self, server_params: ServerParameters
+    ) -> tuple[types.Implementation, mcp.ClientSession]:
+        """Establish a client session to an MCP server."""
+
+        session_stack = contextlib.AsyncExitStack()
+        try:
+            # Create read and write streams that facilitate io with the server.
+            if isinstance(server_params, StdioServerParameters):
+                client = mcp.stdio_client(server_params)
+                read, write = await session_stack.enter_async_context(client)
+            elif isinstance(server_params, SseServerParameters):
+                client = sse_client(
+                    url=server_params.url,
+                    headers=server_params.headers,
+                    timeout=server_params.timeout,
+                    sse_read_timeout=server_params.sse_read_timeout,
+                )
+                read, write = await session_stack.enter_async_context(client)
+            else:
+                client = streamablehttp_client(
+                    url=server_params.url,
+                    headers=server_params.headers,
+                    timeout=server_params.timeout,
+                    sse_read_timeout=server_params.sse_read_timeout,
+                    terminate_on_close=server_params.terminate_on_close,
+                )
+                read, write, _ = await session_stack.enter_async_context(client)
+
+            session = await session_stack.enter_async_context(
+                mcp.ClientSession(read, write)
+            )
+            result = await session.initialize()
+
+            # Session successfully initialized.
+            # Store its stack and register the stack with the main group stack.
+            self._session_exit_stacks[session] = session_stack
+            # session_stack itself becomes a resource managed by the
+            # main _exit_stack.
+            await self._exit_stack.enter_async_context(session_stack)
+
+            return result.serverInfo, session
+        except Exception:
+            # If anything during this setup fails, ensure the session-specific
+            # stack is closed.
+            await session_stack.aclose()
+            raise
+
+    async def _aggregate_components(
+        self, server_info: types.Implementation, session: mcp.ClientSession
+    ) -> None:
+        """Aggregates prompts, resources, and tools from a given session."""
 
         # Create a reverse index so we can find all prompts, resources, and
         # tools belonging to this session. Used for removing components from
@@ -181,47 +298,66 @@ async def connect_to_server(
         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)
+        try:
+            prompts = (await session.list_prompts()).prompts
+            for prompt in prompts:
+                name = self._component_name(prompt.name, server_info)
+                prompts_temp[name] = prompt
+                component_names.prompts.add(name)
+        except McpError as err:
+            logging.warning(f"Could not fetch prompts: {err}")
 
         # 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)
+        try:
+            resources = (await session.list_resources()).resources
+            for resource in resources:
+                name = self._component_name(resource.name, server_info)
+                resources_temp[name] = resource
+                component_names.resources.add(name)
+        except McpError as err:
+            logging.warning(f"Could not fetch resources: {err}")
 
         # 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.",
-                    )
+        try:
+            tools = (await session.list_tools()).tools
+            for tool in 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:
+            logging.warning(f"Could not fetch tools: {err}")
+
+        # Clean up exit stack for session if we couldn't retrieve anything
+        # from the server.
+        if not any((prompts_temp, resources_temp, tools_temp)):
+            del self._session_exit_stacks[session]
+
+        # Check for duplicates.
+        matching_prompts = prompts_temp.keys() & self._prompts.keys()
+        if matching_prompts:
+            raise McpError(
+                types.ErrorData(
+                    code=types.INVALID_PARAMS,
+                    message=f"{matching_prompts} already exist in group prompts.",
+                )
+            )
+        matching_resources = resources_temp.keys() & self._resources.keys()
+        if matching_resources:
+            raise McpError(
+                types.ErrorData(
+                    code=types.INVALID_PARAMS,
+                    message=f"{matching_resources} already exist in group resources.",
+                )
+            )
+        matching_tools = tools_temp.keys() & self._tools.keys()
+        if matching_tools:
+            raise McpError(
+                types.ErrorData(
+                    code=types.INVALID_PARAMS,
+                    message=f"{matching_tools} already exist 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
@@ -230,41 +366,6 @@ async def connect_to_server(
         self._tools.update(tools_temp)
         self._tool_to_session.update(tool_to_session_temp)
 
-        return session
-
-    async def _establish_session(
-        self, server_params: ServerParameters
-    ) -> tuple[types.Implementation, mcp.ClientSession]:
-        """Establish a client session to an MCP server."""
-
-        # Create read and write streams that facilitate io with the server.
-        if isinstance(server_params, StdioServerParameters):
-            client = mcp.stdio_client(server_params)
-            read, write = await self._exit_stack.enter_async_context(client)
-        elif isinstance(server_params, SseServerParameters):
-            client = sse_client(
-                url=server_params.url,
-                headers=server_params.headers,
-                timeout=server_params.timeout,
-                sse_read_timeout=server_params.sse_read_timeout,
-            )
-            read, write = await self._exit_stack.enter_async_context(client)
-        else:
-            client = streamablehttp_client(
-                url=server_params.url,
-                headers=server_params.headers,
-                timeout=server_params.timeout,
-                sse_read_timeout=server_params.sse_read_timeout,
-                terminate_on_close=server_params.terminate_on_close,
-            )
-            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)