Skip to content

Commit 2c27bef

Browse files
committed
Address review feedback on session refresh and exit stack cleanup
- Preserve existing components when MCPError occurs during refresh - Retain exit stack for group-owned sessions with empty catalogs so disconnect cleans them up - Retain original public-Client repro alongside mock-based tests - Document failed-fetch refresh preservation in session-groups docs
1 parent 4922b71 commit 2c27bef

3 files changed

Lines changed: 201 additions & 17 deletions

File tree

‎docs/client/session-groups.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ Run it again. `print(sorted(group.tools))` now shows both:
6666

6767
If you already hold a connected `ClientSession` (`Client.session` is one), hand it to `await group.connect_with_session(server_info, session)` instead of opening a new transport. It aggregates the same way. The group never closes a session it didn't open. `server_info` names the server for component prefixes; on a 2026-era connection `client.server_info` can be `None` (identity is optional), so pass your own `Implementation(name=..., version=...)` in that case.
6868

69-
Calling `connect_with_session` with an already-connected session refreshes that session's aggregated components to match its current catalog. Any components the server no longer exposes are removed from the group, and newly added ones are registered, without triggering duplicate-name collisions against the session's own previous registrations. If any component collides with a *different* session in the group, an `MCPError` is raised and the group's existing state is preserved.
69+
Calling `connect_with_session` with an already-connected session refreshes that session's aggregated components to match its current catalog. Any components the server no longer exposes are removed from the group, and newly added ones are registered, without triggering duplicate-name collisions against the session's own previous registrations. If querying a component type fails with `MCPError` during refresh, existing registrations for that component type remain intact. If any component collides with a *different* session in the group, an `MCPError` is raised and the group's existing state is preserved.
7070

7171
## The classic handshake
7272

‎src/mcp/client/session_group.py‎

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -283,8 +283,8 @@ async def disconnect_from_server(self, session: mcp.ClientSession) -> None:
283283

284284
# Clean up the session's resources via its dedicated exit stack
285285
if session_known_for_stack:
286-
session_stack_to_close = self._session_exit_stacks.pop(session) # pragma: no cover
287-
await session_stack_to_close.aclose() # pragma: no cover
286+
session_stack_to_close = self._session_exit_stacks.pop(session)
287+
await session_stack_to_close.aclose()
288288

289289
async def connect_with_session(
290290
self, server_info: types.Implementation, session: mcp.ClientSession
@@ -378,10 +378,14 @@ async def _aggregate_components(self, server_info: types.Implementation, session
378378
# tools belonging to this session. Used for removing components from
379379
# the session group via self.disconnect_from_server.
380380
component_names = self._ComponentNames()
381+
own_components = self._sessions.get(session)
381382

382383
# Temporary component dictionaries to stage changes. We query and validate
383384
# all components before mutating group state so that any failure leaves
384-
# existing registrations intact.
385+
# existing registrations intact. If querying a component type fails with
386+
# MCPError during a refresh, existing registrations for that component type
387+
# are preserved rather than evicted (a failed fetch is not evidence that the
388+
# server removed those components).
385389
prompts_temp: dict[str, types.Prompt] = {}
386390
resources_temp: dict[str, types.Resource] = {}
387391
tools_temp: dict[str, types.Tool] = {}
@@ -394,8 +398,12 @@ async def _aggregate_components(self, server_info: types.Implementation, session
394398
name = self._component_name(prompt.name, server_info)
395399
prompts_temp[name] = prompt
396400
component_names.prompts.add(name)
397-
except MCPError as err: # pragma: no cover
401+
except MCPError as err:
398402
logging.warning(f"Could not fetch prompts: {err}")
403+
if own_components:
404+
for name in own_components.prompts:
405+
prompts_temp[name] = self._prompts[name]
406+
component_names.prompts.add(name)
399407

400408
# Query the server for its resources and aggregate to list.
401409
try:
@@ -404,8 +412,12 @@ async def _aggregate_components(self, server_info: types.Implementation, session
404412
name = self._component_name(resource.name, server_info)
405413
resources_temp[name] = resource
406414
component_names.resources.add(name)
407-
except MCPError as err: # pragma: no cover
415+
except MCPError as err:
408416
logging.warning(f"Could not fetch resources: {err}")
417+
if own_components:
418+
for name in own_components.resources:
419+
resources_temp[name] = self._resources[name]
420+
component_names.resources.add(name)
409421

410422
# Query the server for its tools and aggregate to list.
411423
try:
@@ -415,17 +427,15 @@ async def _aggregate_components(self, server_info: types.Implementation, session
415427
tools_temp[name] = tool
416428
tool_to_session_temp[name] = session
417429
component_names.tools.add(name)
418-
except MCPError as err: # pragma: no cover
430+
except MCPError as err:
419431
logging.warning(f"Could not fetch tools: {err}")
420-
421-
# Clean up exit stack for session if we couldn't retrieve anything
422-
# from the server.
423-
if not any((prompts_temp, resources_temp, tools_temp)):
424-
if session in self._session_exit_stacks:
425-
del self._session_exit_stacks[session]
432+
if own_components:
433+
for name in own_components.tools:
434+
tools_temp[name] = self._tools[name]
435+
tool_to_session_temp[name] = session
436+
component_names.tools.add(name)
426437

427438
# Check for duplicates, ignoring components already registered by this session.
428-
own_components = self._sessions.get(session)
429439
other_prompts = self._prompts.keys() - own_components.prompts if own_components else self._prompts.keys()
430440
matching_prompts = prompts_temp.keys() & other_prompts
431441
if matching_prompts:

‎tests/client/test_session_group.py‎

Lines changed: 177 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import contextlib
22
from unittest import mock
33

4+
import anyio
45
import httpx2
56
import mcp_types as types
67
import pytest
@@ -13,6 +14,7 @@
1314
StreamableHttpParameters,
1415
)
1516
from mcp.client.stdio import StdioServerParameters
17+
from mcp.server import Server, ServerRequestContext
1618
from mcp.shared.exceptions import MCPError
1719

1820

@@ -408,8 +410,50 @@ async def test_client_session_group_establish_session_parameterized(
408410

409411

410412
@pytest.mark.anyio
413+
@pytest.mark.filterwarnings("ignore:Passing 'msg' argument.*:DeprecationWarning")
411414
async def test_disconnect_removes_all_tools_after_same_session_reregistration() -> None:
412-
"""Verify disconnecting after re-registering a session removes all tools."""
415+
"""Reproduce the SDK cleanup gap after a caller re-registers a session with changed tools."""
416+
catalog = [types.Tool(name="old", input_schema={"type": "object"})]
417+
418+
async def list_tools(
419+
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
420+
) -> types.ListToolsResult:
421+
return types.ListToolsResult(tools=list(catalog))
422+
423+
async def list_resources(
424+
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
425+
) -> types.ListResourcesResult:
426+
return types.ListResourcesResult(resources=[])
427+
428+
async def list_prompts(
429+
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
430+
) -> types.ListPromptsResult:
431+
return types.ListPromptsResult(prompts=[])
432+
433+
server = Server(
434+
"mutable-catalog",
435+
on_list_tools=list_tools,
436+
on_list_resources=list_resources,
437+
on_list_prompts=list_prompts,
438+
)
439+
group = ClientSessionGroup()
440+
with anyio.fail_after(5):
441+
async with mcp.Client(server) as client:
442+
server_info = types.Implementation(name="mutable-catalog", version="1")
443+
await group.connect_with_session(server_info, client.session)
444+
445+
catalog[:] = [types.Tool(name="new", input_schema={"type": "object"})]
446+
await group.connect_with_session(server_info, client.session)
447+
448+
await group.disconnect_from_server(client.session)
449+
450+
assert group.sessions == []
451+
assert group.tools == {}
452+
453+
454+
@pytest.mark.anyio
455+
async def test_disconnect_removes_all_tools_after_same_session_reregistration_mock() -> None:
456+
"""Verify disconnecting after re-registering a session removes all tools (mock-based)."""
413457
server_info = types.Implementation(name="TestServer", version="1.0")
414458
mock_session = mock.AsyncMock(spec=mcp.ClientSession)
415459

@@ -669,21 +713,151 @@ async def test_connect_with_session_empty_catalog() -> None:
669713

670714
@pytest.mark.anyio
671715
async def test_connect_to_server_empty_catalog(mock_exit_stack: contextlib.AsyncExitStack) -> None:
672-
"""Verify group-owned connect_to_server with empty catalog cleans up exit stack without error."""
716+
"""Verify group-owned connect_to_server with empty catalog retains exit stack and cleans up on disconnect."""
673717
mock_server_info = mock.Mock(spec=types.Implementation)
674718
mock_server_info.name = "EmptyServer"
675719
mock_session = mock.AsyncMock(spec=mcp.ClientSession)
676720
mock_session.list_tools.return_value = types.ListToolsResult(tools=[])
677721
mock_session.list_resources.return_value = types.ListResourcesResult(resources=[])
678722
mock_session.list_prompts.return_value = types.ListPromptsResult(prompts=[])
679723

724+
session_stack = mock.AsyncMock(spec=contextlib.AsyncExitStack)
725+
680726
async def fake_establish(*args: object, **kwargs: object) -> tuple[types.Implementation, mock.AsyncMock]:
681-
group._session_exit_stacks[mock_session] = mock.AsyncMock()
727+
group._session_exit_stacks[mock_session] = session_stack
682728
return mock_server_info, mock_session
683729

684730
group = ClientSessionGroup(exit_stack=mock_exit_stack)
685731
with mock.patch.object(group, "_establish_session", side_effect=fake_establish):
686732
await group.connect_to_server(StdioServerParameters(command="test"))
687733

734+
assert mock_session in group._session_exit_stacks
735+
assert group.tools == {}
736+
assert group.sessions == [mock_session]
737+
738+
await group.disconnect_from_server(mock_session)
739+
session_stack.aclose.assert_awaited_once()
688740
assert mock_session not in group._session_exit_stacks
741+
assert group.sessions == []
742+
743+
744+
@pytest.mark.anyio
745+
async def test_group_owned_session_refreshed_to_empty_and_disconnected_cleans_up(
746+
mock_exit_stack: contextlib.AsyncExitStack,
747+
) -> None:
748+
"""Verify group-owned session refreshed to empty catalog retains exit stack and closes it on disconnect."""
749+
mock_server_info = mock.Mock(spec=types.Implementation)
750+
mock_server_info.name = "TestServer"
751+
mock_session = mock.AsyncMock(spec=mcp.ClientSession)
752+
initial_tool = types.Tool(name="initial_tool", input_schema={})
753+
mock_session.list_tools.return_value = types.ListToolsResult(tools=[initial_tool])
754+
mock_session.list_resources.return_value = types.ListResourcesResult(resources=[])
755+
mock_session.list_prompts.return_value = types.ListPromptsResult(prompts=[])
756+
757+
session_stack = mock.AsyncMock(spec=contextlib.AsyncExitStack)
758+
759+
async def fake_establish(*args: object, **kwargs: object) -> tuple[types.Implementation, mock.AsyncMock]:
760+
group._session_exit_stacks[mock_session] = session_stack
761+
return mock_server_info, mock_session
762+
763+
group = ClientSessionGroup(exit_stack=mock_exit_stack)
764+
with mock.patch.object(group, "_establish_session", side_effect=fake_establish):
765+
await group.connect_to_server(StdioServerParameters(command="test"))
766+
767+
assert mock_session in group._session_exit_stacks
768+
assert group.tools == {"initial_tool": initial_tool}
769+
assert group.sessions == [mock_session]
770+
771+
# Refresh to empty catalog
772+
mock_session.list_tools.return_value = types.ListToolsResult(tools=[])
773+
await group.connect_with_session(mock_server_info, mock_session)
774+
689775
assert group.tools == {}
776+
assert group.sessions == [mock_session]
777+
# Exit stack must NOT be deleted without closing
778+
assert mock_session in group._session_exit_stacks
779+
session_stack.aclose.assert_not_called()
780+
781+
# Disconnect must close the exit stack
782+
await group.disconnect_from_server(mock_session)
783+
session_stack.aclose.assert_awaited_once()
784+
assert mock_session not in group._session_exit_stacks
785+
assert group.sessions == []
786+
787+
788+
@pytest.mark.anyio
789+
async def test_connect_with_session_failed_refresh_preserves_existing_components() -> None:
790+
"""Verify that MCPError when fetching components during refresh preserves existing registrations."""
791+
server_info = types.Implementation(name="TestServer", version="1.0")
792+
session = mock.AsyncMock(spec=mcp.ClientSession)
793+
794+
initial_prompt = types.Prompt(name="orig_prompt")
795+
initial_resource = types.Resource(uri="file:///orig_resource", name="orig_res")
796+
initial_tool = types.Tool(name="orig_tool", input_schema={})
797+
798+
session.list_prompts.return_value = types.ListPromptsResult(prompts=[initial_prompt])
799+
session.list_resources.return_value = types.ListResourcesResult(resources=[initial_resource])
800+
session.list_tools.return_value = types.ListToolsResult(tools=[initial_tool])
801+
802+
group = ClientSessionGroup()
803+
await group.connect_with_session(server_info, session)
804+
805+
assert group.prompts == {"orig_prompt": initial_prompt}
806+
assert group.resources == {"orig_res": initial_resource}
807+
assert group.tools == {"orig_tool": initial_tool}
808+
assert group._tool_to_session == {"orig_tool": session}
809+
assert group.sessions == [session]
810+
811+
# Simulate MCPError on all component listings during refresh
812+
session.list_prompts.side_effect = MCPError(code=types.INTERNAL_ERROR, message="prompts failed")
813+
session.list_resources.side_effect = MCPError(code=types.INTERNAL_ERROR, message="resources failed")
814+
session.list_tools.side_effect = MCPError(code=types.INTERNAL_ERROR, message="tools failed")
815+
816+
await group.connect_with_session(server_info, session)
817+
818+
# Invariants: existing components are preserved, not evicted
819+
assert group.prompts == {"orig_prompt": initial_prompt}
820+
assert group.resources == {"orig_res": initial_resource}
821+
assert group.tools == {"orig_tool": initial_tool}
822+
assert group._tool_to_session == {"orig_tool": session}
823+
assert group._sessions[session].prompts == {"orig_prompt"}
824+
assert group._sessions[session].resources == {"orig_res"}
825+
assert group._sessions[session].tools == {"orig_tool"}
826+
assert group.sessions == [session]
827+
828+
# Preserved tools remain callable and routed
829+
session.call_tool.return_value = types.CallToolResult(content=[types.TextContent(type="text", text="ok")])
830+
call_res = await group.call_tool("orig_tool", {})
831+
assert isinstance(call_res.content[0], types.TextContent)
832+
assert call_res.content[0].text == "ok"
833+
834+
# Disconnect cleans up preserved components completely
835+
await group.disconnect_from_server(session)
836+
assert group.prompts == {}
837+
assert group.resources == {}
838+
assert group.tools == {}
839+
assert group._tool_to_session == {}
840+
assert group.sessions == []
841+
842+
843+
@pytest.mark.anyio
844+
async def test_connect_with_session_initial_fetch_mcperror_handled() -> None:
845+
"""Verify that MCPError on initial component queries is handled gracefully without error."""
846+
server_info = types.Implementation(name="ErrorServer", version="1.0")
847+
session = mock.AsyncMock(spec=mcp.ClientSession)
848+
849+
session.list_prompts.side_effect = MCPError(code=types.INTERNAL_ERROR, message="no prompts")
850+
session.list_resources.side_effect = MCPError(code=types.INTERNAL_ERROR, message="no resources")
851+
session.list_tools.side_effect = MCPError(code=types.INTERNAL_ERROR, message="no tools")
852+
853+
group = ClientSessionGroup()
854+
await group.connect_with_session(server_info, session)
855+
856+
assert group.prompts == {}
857+
assert group.resources == {}
858+
assert group.tools == {}
859+
assert group._tool_to_session == {}
860+
assert group.sessions == [session]
861+
862+
await group.disconnect_from_server(session)
863+
assert group.sessions == []

0 commit comments

Comments
 (0)