Fix ClientSessionGroup cleanup logic (#787)

This commit is contained in:
Mo
2025-05-23 03:57:46 -04:00
committed by GitHub
parent e80c0150e1
commit 0e2a34189c

View File

@@ -77,10 +77,10 @@ class ClientSessionGroup:
the client and can be accessed via the session. the client and can be accessed via the session.
Example Usage: 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: async with ClientSessionGroup(component_name_hook=name_fn) as group:
for server_params in server_params: for server_params in server_params:
group.connect_to_server(server_param) await group.connect_to_server(server_param)
... ...
""" """
@@ -145,14 +145,15 @@ class ClientSessionGroup:
) -> bool | None: ) -> bool | None:
"""Closes session exit stacks and main exit stack upon completion.""" """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. # Concurrently close session stacks.
async with anyio.create_task_group() as tg: async with anyio.create_task_group() as tg:
for exit_stack in self._session_exit_stacks.values(): for exit_stack in self._session_exit_stacks.values():
tg.start_soon(exit_stack.aclose) 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 @property
def sessions(self) -> list[mcp.ClientSession]: def sessions(self) -> list[mcp.ClientSession]: