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

@@ -11,9 +11,11 @@
 import contextlib
 from collections.abc import Callable
 from datetime import timedelta
+from types import TracebackType
 from typing import Any, TypeAlias
 
 from pydantic import BaseModel
+from typing_extensions import Self
 
 import mcp
 from mcp import types
@@ -72,6 +74,14 @@ class ClientSessionGroup:
     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(...)
+
+    Example Usage:
+        name_fn = lambda name, server_info: f"{(server_info.name)}.{name}"
+        with async 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 +100,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 +110,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 +121,32 @@ def __init__(
 
         self._sessions = {}
         self._tool_to_session = {}
-        self._exit_stack = exit_stack
+        self._main_exit_stack = exit_stack or contextlib.AsyncExitStack()
+        self._session_exit_stacks = {}
         self._component_name_hook = component_name_hook
 
+    async def __aenter__(self) -> Self:
+        # If ClientSessionGroup itself is managing the lifecycle of _main_exit_stack
+        # (i.e., it created it), it should enter it.
+        # If _main_exit_stack was passed in, it's assumed the caller manages
+        # its entry/exit.
+        # For simplicity and consistency with how AsyncExitStack is often used when
+        # provided as a dependency, we might not need to enter it here if it's
+        # managed externally. However, if this class is the primary owner, entering it
+        # ensures its 'aclose' is called even if passed in. Let's assume the
+        # passed-in stack is already entered by the caller if needed.
+        # For now, we just return self as the main stack's lifecycle is tied to aclose.
+        return self
+
+    async def __aexit__(
+        self,
+        _exc_type: type[BaseException] | None,
+        _exc_val: BaseException | None,
+        _exc_tb: TracebackType | None,
+    ) -> bool | None:
+        await self._main_exit_stack.aclose()
+        return None  # Do not suppress exceptio
+
     @property
     def prompts(self) -> dict[str, types.Prompt]:
         """Returns the prompts as a dictionary of names to prompts."""
@@ -133,31 +167,42 @@ async def call_tool(self, name: str, args: dict[str, Any]) -> types.CallToolResu
         session = self._tool_to_session[name]
         return await session.call_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_to_server(
         self,
@@ -168,102 +213,130 @@ async def connect_to_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.",
+        try:
+            # 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.",
+                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.",
+                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
+                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
+        except Exception:
+            # If any error occurs during component fetching or registration
+            # after session establishment, clean up the newly established
+            # session's resources.
+            if session in self._session_exit_stacks:
+                session_stack_to_close = self._session_exit_stacks.pop(session)
+                await session_stack_to_close.aclose()
+                # The session_stack_to_close was also entered into
+                # self._exit_stack. Its .aclose() being called here means when
+                # self._exit_stack later tries to close it, it will be a
+                # no-op, which is fine.
+            raise
 
     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_specific_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 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
+            session = await self._exit_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_specific_stack
+            # session_specific_stack itself becomes a resource managed by the
+            # main _exit_stack.
+            await self._exit_stack.enter_async_context(session_specific_stack)
+
+            return result.serverInfo, session
+        except Exception:
+            # If anything during this setup fails, ensure the session-specific
+            # stack is closed.
+            await session_specific_stack.aclose()
+            raise
 
     def _component_name(self, name: str, server_info: types.Implementation) -> str:
         if self._component_name_hook:

tests/client/test_session_group.py

@@ -151,7 +151,7 @@ def name_hook(name: str, server_info: types.Implementation) -> str:
         assert group.tools[expected_tool_name] == mock_tool
         assert group._tool_to_session[expected_tool_name] == mock_session
 
-    def test_disconnect_from_server(self):  # No mock arguments needed
+    async def test_disconnect_from_server(self):  # No mock arguments needed
         """Test disconnecting from a server."""
         # --- Test Setup ---
         group = ClientSessionGroup()
@@ -205,7 +205,7 @@ def test_disconnect_from_server(self):  # No mock arguments needed
         assert "prm1" in group._prompts
 
         # --- Test Execution ---
-        group.disconnect_from_server(mock_session)
+        await group.disconnect_from_server(mock_session)
 
         # --- Assertions ---
         assert mock_session not in group._sessions
@@ -263,12 +263,12 @@ async def test_connect_to_server_duplicate_tool_raises_error(self, mock_exit_sta
         )  # Ensure it's the original mock
 
     # No patching needed here
-    def test_disconnect_non_existent_server(self):
+    async def test_disconnect_non_existent_server(self):
         """Test disconnecting a server that isn't connected."""
         session = mock.Mock(spec=mcp.ClientSession)
         group = ClientSessionGroup()
         with pytest.raises(McpError):
-            group.disconnect_from_server(session)
+            await group.disconnect_from_server(session)
 
     @pytest.mark.parametrize(
         "server_params_instance, client_type_name, patch_target_for_client_func",