From 4922b71e2044e7429c9e63f12b08be134d23d44c Mon Sep 17 00:00:00 2001 From: Abhay Tiwari Date: Thu, 15 Oct 2026 19:42:32 +0530 Subject: [PATCH] Refresh session components on re-registration in ClientSessionGroup (#3713) - Evict previous components registered by a session before aggregating new catalog - Exclude session's own components from duplicate collision checks - Guard exit stack cleanup against KeyError for caller-owned sessions - Update docs/client/session-groups.md to document connect_with_session refresh behavior - Add comprehensive tests asserting full internal and public invariants --- docs/client/session-groups.md | 4 +- src/mcp/client/session_group.py | 34 +++- tests/client/test_session_group.py | 282 +++++++++++++++++++++++++++++ 3 files changed, 311 insertions(+), 9 deletions(-) diff --git a/docs/client/session-groups.md b/docs/client/session-groups.md index de43e74890..a4307f328e 100644 --- a/docs/client/session-groups.md +++ b/docs/client/session-groups.md @@ -66,6 +66,8 @@ Run it again. `print(sorted(group.tools))` now shows both: 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. +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. + ## The classic handshake `ClientSessionGroup` is built on `ClientSession`, not on `Client`. Each `connect_to_server` runs the classic `initialize` handshake. It never sends the `server/discover` probe described in **[Protocol versions](../protocol-versions.md)**. Every MCP server understands that handshake, so this costs you compatibility with nothing; it only means a group takes the older, slower path to a server that could do better. @@ -77,6 +79,6 @@ If you already hold a connected `ClientSession` (`Client.session` is one), hand * `group.call_tool(name, arguments)` routes to the owning server for you. * Names must be unique across the whole group; two servers with a `search` tool cannot coexist on their own. * `component_name_hook=` rewrites every registered name. The dict key changes, the wire name does not. -* `connect_with_session` adds a session you already hold; `disconnect_from_server` removes one. +* `connect_with_session` adds a session you already hold (or refreshes an existing one); `disconnect_from_server` removes one. The handshake a group speaks (and the faster one a `Client` prefers) is the subject of **[Protocol versions](../protocol-versions.md)**. diff --git a/src/mcp/client/session_group.py b/src/mcp/client/session_group.py index 13d8928ed4..08c50bf011 100644 --- a/src/mcp/client/session_group.py +++ b/src/mcp/client/session_group.py @@ -175,7 +175,7 @@ async def __aexit__( @property def sessions(self) -> list[mcp.ClientSession]: """Returns the list of sessions being managed.""" - return list(self._sessions.keys()) # pragma: no cover + return list(self._sessions.keys()) @property def prompts(self) -> dict[str, types.Prompt]: @@ -379,8 +379,9 @@ async def _aggregate_components(self, server_info: types.Implementation, session # 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. + # Temporary component dictionaries to stage changes. We query and validate + # all components before mutating group state so that any failure leaves + # existing registrations intact. prompts_temp: dict[str, types.Prompt] = {} resources_temp: dict[str, types.Resource] = {} tools_temp: dict[str, types.Tool] = {} @@ -420,25 +421,42 @@ async def _aggregate_components(self, server_info: types.Implementation, session # 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] # pragma: no cover + if session in self._session_exit_stacks: + del self._session_exit_stacks[session] - # Check for duplicates. - matching_prompts = prompts_temp.keys() & self._prompts.keys() + # Check for duplicates, ignoring components already registered by this session. + own_components = self._sessions.get(session) + other_prompts = self._prompts.keys() - own_components.prompts if own_components else self._prompts.keys() + matching_prompts = prompts_temp.keys() & other_prompts if matching_prompts: raise MCPError( # pragma: no cover code=types.INVALID_PARAMS, message=f"{matching_prompts} already exist in group prompts.", ) - matching_resources = resources_temp.keys() & self._resources.keys() + other_resources = ( + self._resources.keys() - own_components.resources if own_components else self._resources.keys() + ) + matching_resources = resources_temp.keys() & other_resources if matching_resources: raise MCPError( # pragma: no cover code=types.INVALID_PARAMS, message=f"{matching_resources} already exist in group resources.", ) - matching_tools = tools_temp.keys() & self._tools.keys() + other_tools = self._tools.keys() - own_components.tools if own_components else self._tools.keys() + matching_tools = tools_temp.keys() & other_tools if matching_tools: raise MCPError(code=types.INVALID_PARAMS, message=f"{matching_tools} already exist in group tools.") + # Clean up any previously registered components for this session before aggregating new ones. + if own_components: + for name in own_components.prompts: + self._prompts.pop(name, None) + for name in own_components.resources: + self._resources.pop(name, None) + for name in own_components.tools: + self._tools.pop(name, None) + self._tool_to_session.pop(name, None) + # Aggregate components. self._sessions[session] = component_names self._prompts.update(prompts_temp) diff --git a/tests/client/test_session_group.py b/tests/client/test_session_group.py index 922e1eebcb..af7b50a982 100644 --- a/tests/client/test_session_group.py +++ b/tests/client/test_session_group.py @@ -405,3 +405,285 @@ async def test_client_session_group_establish_session_parameterized( # 3. Assert returned values assert returned_server_info is mock_initialize_result.server_info assert returned_session is mock_entered_session + + +@pytest.mark.anyio +async def test_disconnect_removes_all_tools_after_same_session_reregistration() -> None: + """Verify disconnecting after re-registering a session removes all tools.""" + server_info = types.Implementation(name="TestServer", version="1.0") + mock_session = mock.AsyncMock(spec=mcp.ClientSession) + + old_tool = types.Tool(name="old", input_schema={}) + new_tool = types.Tool(name="new", input_schema={}) + + mock_session.list_tools.return_value = types.ListToolsResult(tools=[old_tool]) + mock_session.list_resources.return_value = types.ListResourcesResult(resources=[]) + mock_session.list_prompts.return_value = types.ListPromptsResult(prompts=[]) + + group = ClientSessionGroup() + await group.connect_with_session(server_info, mock_session) + assert group.tools == {"old": old_tool} + assert group._tool_to_session == {"old": mock_session} + assert group._sessions[mock_session].tools == {"old"} + assert group.sessions == [mock_session] + + # Re-register same session with changed catalog (disjoint tools) + mock_session.list_tools.return_value = types.ListToolsResult(tools=[new_tool]) + await group.connect_with_session(server_info, mock_session) + assert group.tools == {"new": new_tool} + assert "old" not in group.tools + assert group._tool_to_session == {"new": mock_session} + assert "old" not in group._tool_to_session + assert group._sessions[mock_session].tools == {"new"} + assert group.sessions == [mock_session] + + # Disconnect removes all tools belonging to the session + await group.disconnect_from_server(mock_session) + assert group.sessions == [] + assert group.tools == {} + assert group._tool_to_session == {} + assert mock_session not in group._sessions + + +@pytest.mark.anyio +async def test_connect_with_session_collision_with_other_session_raises() -> None: + """Verify that re-registering raises MCPError if a component collides with another session.""" + server_info_1 = types.Implementation(name="Server1", version="1.0") + session_1 = mock.AsyncMock(spec=mcp.ClientSession) + tool_shared = types.Tool(name="shared_tool", input_schema={}) + session_1.list_tools.return_value = types.ListToolsResult(tools=[tool_shared]) + session_1.list_resources.return_value = types.ListResourcesResult(resources=[]) + session_1.list_prompts.return_value = types.ListPromptsResult(prompts=[]) + + server_info_2 = types.Implementation(name="Server2", version="1.0") + session_2 = mock.AsyncMock(spec=mcp.ClientSession) + tool_other = types.Tool(name="other_tool", input_schema={}) + session_2.list_tools.return_value = types.ListToolsResult(tools=[tool_other]) + session_2.list_resources.return_value = types.ListResourcesResult(resources=[]) + session_2.list_prompts.return_value = types.ListPromptsResult(prompts=[]) + + group = ClientSessionGroup() + await group.connect_with_session(server_info_1, session_1) + await group.connect_with_session(server_info_2, session_2) + + # Initial invariant checks + assert group.tools == {"shared_tool": tool_shared, "other_tool": tool_other} + assert group._tool_to_session == {"shared_tool": session_1, "other_tool": session_2} + assert group._sessions[session_1].tools == {"shared_tool"} + assert group._sessions[session_2].tools == {"other_tool"} + + # Attempt to re-register session 2 with a tool that collides with session 1 + session_2.list_tools.return_value = types.ListToolsResult(tools=[tool_shared]) + with pytest.raises(MCPError) as excinfo: + await group.connect_with_session(server_info_2, session_2) + + assert excinfo.value.error.code == types.INVALID_PARAMS + assert "shared_tool" in excinfo.value.error.message + + # Assert full invariant: session 2's previous state must be completely preserved + assert group.tools == {"shared_tool": tool_shared, "other_tool": tool_other} + assert group._tool_to_session == {"shared_tool": session_1, "other_tool": session_2} + assert group._sessions[session_1].tools == {"shared_tool"} + assert group._sessions[session_2].tools == {"other_tool"} + + # Ensure colliding tool was NOT re-routed to session 2 + session_1.call_tool.return_value = types.CallToolResult( + content=[types.TextContent(type="text", text="server_1_ok")] + ) + call_result = await group.call_tool("shared_tool", {}) + assert isinstance(call_result.content[0], types.TextContent) + assert call_result.content[0].text == "server_1_ok" + session_1.call_tool.assert_awaited_once_with( + "shared_tool", + arguments={}, + read_timeout_seconds=None, + progress_callback=None, + input_responses=None, + request_state=None, + meta=None, + allow_input_required=False, + ) + session_2.call_tool.assert_not_called() + + +@pytest.mark.anyio +async def test_connect_with_session_reregister_refreshes_all_components() -> None: + """Verify that re-registering refreshes prompts, resources, and tools without false collision.""" + server_info = types.Implementation(name="TestServer", version="1.0") + session = mock.AsyncMock(spec=mcp.ClientSession) + + tool_1 = types.Tool(name="tool_1", input_schema={}) + res_1 = types.Resource(uri="file:///res1", name="res_1") + prompt_1 = types.Prompt(name="prompt_1") + + session.list_tools.return_value = types.ListToolsResult(tools=[tool_1]) + session.list_resources.return_value = types.ListResourcesResult(resources=[res_1]) + session.list_prompts.return_value = types.ListPromptsResult(prompts=[prompt_1]) + + group = ClientSessionGroup() + await group.connect_with_session(server_info, session) + + # Invariant checks for initial state + assert group.tools == {"tool_1": tool_1} + assert group.resources == {"res_1": res_1} + assert group.prompts == {"prompt_1": prompt_1} + assert group._tool_to_session == {"tool_1": session} + assert group._sessions[session].tools == {"tool_1"} + assert group._sessions[session].resources == {"res_1"} + assert group._sessions[session].prompts == {"prompt_1"} + + # Re-register with updated components + tool_2 = types.Tool(name="tool_2", input_schema={}) + res_2 = types.Resource(uri="file:///res2", name="res_2") + prompt_2 = types.Prompt(name="prompt_2") + + session.list_tools.return_value = types.ListToolsResult(tools=[tool_2]) + session.list_resources.return_value = types.ListResourcesResult(resources=[res_2]) + session.list_prompts.return_value = types.ListPromptsResult(prompts=[prompt_2]) + + await group.connect_with_session(server_info, session) + + # Invariant checks after refresh + assert group.tools == {"tool_2": tool_2} + assert group.resources == {"res_2": res_2} + assert group.prompts == {"prompt_2": prompt_2} + assert "tool_1" not in group.tools + assert "res_1" not in group.resources + assert "prompt_1" not in group.prompts + assert group._tool_to_session == {"tool_2": session} + assert "tool_1" not in group._tool_to_session + assert group._sessions[session].tools == {"tool_2"} + assert group._sessions[session].resources == {"res_2"} + assert group._sessions[session].prompts == {"prompt_2"} + + # Verify call_tool routing via _tool_to_session + session.call_tool.return_value = types.CallToolResult(content=[types.TextContent(type="text", text="tool_2_ok")]) + call_res = await group.call_tool("tool_2", {"arg": "val"}) + assert isinstance(call_res.content[0], types.TextContent) + assert call_res.content[0].text == "tool_2_ok" + session.call_tool.assert_awaited_once_with( + "tool_2", + arguments={"arg": "val"}, + read_timeout_seconds=None, + progress_callback=None, + input_responses=None, + request_state=None, + meta=None, + allow_input_required=False, + ) + + # Calling evicted tool fails + with pytest.raises(KeyError): + await group.call_tool("tool_1", {}) + + +@pytest.mark.anyio +async def test_connect_with_session_reregister_with_name_hook() -> None: + """Verify re-registration respects component_name_hook and cleans up hooked names.""" + + def name_hook(name: str, server_info: types.Implementation) -> str: + return f"{server_info.name}.{name}" + + server_info = types.Implementation(name="MyServer", version="1.0") + session = mock.AsyncMock(spec=mcp.ClientSession) + + tool_a = types.Tool(name="alpha", input_schema={}) + session.list_tools.return_value = types.ListToolsResult(tools=[tool_a]) + session.list_resources.return_value = types.ListResourcesResult(resources=[]) + session.list_prompts.return_value = types.ListPromptsResult(prompts=[]) + + group = ClientSessionGroup(component_name_hook=name_hook) + await group.connect_with_session(server_info, session) + assert group.tools == {"MyServer.alpha": tool_a} + assert group._tool_to_session == {"MyServer.alpha": session} + assert group._sessions[session].tools == {"MyServer.alpha"} + + tool_b = types.Tool(name="beta", input_schema={}) + session.list_tools.return_value = types.ListToolsResult(tools=[tool_b]) + + await group.connect_with_session(server_info, session) + assert group.tools == {"MyServer.beta": tool_b} + assert "MyServer.alpha" not in group.tools + assert group._tool_to_session == {"MyServer.beta": session} + assert "MyServer.alpha" not in group._tool_to_session + assert group._sessions[session].tools == {"MyServer.beta"} + + session.call_tool.return_value = types.CallToolResult(content=[types.TextContent(type="text", text="beta_ok")]) + call_res = await group.call_tool("MyServer.beta", {}) + assert isinstance(call_res.content[0], types.TextContent) + assert call_res.content[0].text == "beta_ok" + session.call_tool.assert_awaited_once_with( + "beta", + arguments={}, + read_timeout_seconds=None, + progress_callback=None, + input_responses=None, + request_state=None, + meta=None, + allow_input_required=False, + ) + + with pytest.raises(KeyError): + await group.call_tool("MyServer.alpha", {}) + + await group.disconnect_from_server(session) + assert group.tools == {} + assert group._tool_to_session == {} + assert group.sessions == [] + + +@pytest.mark.anyio +async def test_connect_with_session_empty_catalog() -> None: + """Verify caller-owned session with empty catalog registers and refreshes without KeyError.""" + server_info = types.Implementation(name="EmptyServer", version="1.0") + session = mock.AsyncMock(spec=mcp.ClientSession) + session.list_tools.return_value = types.ListToolsResult(tools=[]) + session.list_resources.return_value = types.ListResourcesResult(resources=[]) + session.list_prompts.return_value = types.ListPromptsResult(prompts=[]) + + group = ClientSessionGroup() + # Connect with empty catalog — must not raise KeyError (:422-423 guard) + await group.connect_with_session(server_info, session) + assert group.tools == {} + assert group.resources == {} + assert group.prompts == {} + assert group.sessions == [session] + assert session not in group._session_exit_stacks + + # Refresh with a tool added + sample_tool = types.Tool(name="sample", input_schema={}) + session.list_tools.return_value = types.ListToolsResult(tools=[sample_tool]) + await group.connect_with_session(server_info, session) + assert group.tools == {"sample": sample_tool} + + # Refresh back to empty catalog — must not raise KeyError + session.list_tools.return_value = types.ListToolsResult(tools=[]) + await group.connect_with_session(server_info, session) + assert group.tools == {} + assert "sample" not in group.tools + assert group.sessions == [session] + + await group.disconnect_from_server(session) + assert group.sessions == [] + + +@pytest.mark.anyio +async def test_connect_to_server_empty_catalog(mock_exit_stack: contextlib.AsyncExitStack) -> None: + """Verify group-owned connect_to_server with empty catalog cleans up exit stack without error.""" + mock_server_info = mock.Mock(spec=types.Implementation) + mock_server_info.name = "EmptyServer" + mock_session = mock.AsyncMock(spec=mcp.ClientSession) + mock_session.list_tools.return_value = types.ListToolsResult(tools=[]) + mock_session.list_resources.return_value = types.ListResourcesResult(resources=[]) + mock_session.list_prompts.return_value = types.ListPromptsResult(prompts=[]) + + async def fake_establish(*args: object, **kwargs: object) -> tuple[types.Implementation, mock.AsyncMock]: + group._session_exit_stacks[mock_session] = mock.AsyncMock() + return mock_server_info, mock_session + + group = ClientSessionGroup(exit_stack=mock_exit_stack) + with mock.patch.object(group, "_establish_session", side_effect=fake_establish): + await group.connect_to_server(StdioServerParameters(command="test")) + + assert mock_session not in group._session_exit_stacks + assert group.tools == {}