From 37c1a21c4e856a6d9fa6bdd05bfbdcd68b6a9aad Mon Sep 17 00:00:00 2001 From: Mo Date: Thu, 22 May 2025 15:56:37 +0000 Subject: [PATCH] Fix `ClientSessionGroup` cleanup logic The order of steps responsible for closing exist stacks is reversed. This is causing error when the `ClientSessionGroup` exits its context. This change also includes some small doc improvements. --- src/mcp/client/session_group.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/mcp/client/session_group.py b/src/mcp/client/session_group.py index c23f2523e..a430533b3 100644 --- a/src/mcp/client/session_group.py +++ b/src/mcp/client/session_group.py @@ -77,10 +77,10 @@ class ClientSessionGroup: the client and can be accessed via the session. Example Usage: - name_fn = lambda name, server_info: f"{(server_info.name)}-{name}" + 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) + await group.connect_to_server(server_param) ... """ @@ -145,14 +145,15 @@ async def __aexit__( ) -> bool | None: """Closes session exit stacks and main exit stack upon completion.""" + # Only close the main exit stack if we created it + if self._owns_exit_stack: + await self._exit_stack.aclose() + # 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]: