Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion src/mcp/client/session_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -298,7 +298,13 @@ async def connect_to_server(
) -> mcp.ClientSession:
"""Connects to a single MCP server."""
server_info, session = await self._establish_session(server_params, session_params or ClientSessionParameters())
return await self.connect_with_session(server_info, session)
try:
return await self.connect_with_session(server_info, session)
except BaseException:
session_stack = self._session_exit_stacks.pop(session, None)
if session_stack is not None:
await session_stack.aclose()
raise

async def _establish_session(
self,
Expand Down
9 changes: 8 additions & 1 deletion tests/client/test_session_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -266,13 +266,18 @@ async def test_client_session_group_connect_to_server_duplicate_tool_raises_erro
# Keep other lists empty for simplicity
mock_session_new.list_resources.return_value = mock.AsyncMock(resources=[])
mock_session_new.list_prompts.return_value = mock.AsyncMock(prompts=[])
mock_stack_new = mock.AsyncMock(spec=contextlib.AsyncExitStack)

async def _establish(*_args: object, **_kwargs: object):
group._session_exit_stacks[mock_session_new] = mock_stack_new
return mock_server_info_new, mock_session_new

# --- Test Execution and Assertion ---
with pytest.raises(MCPError) as excinfo:
with mock.patch.object(
group,
"_establish_session",
return_value=(mock_server_info_new, mock_session_new),
side_effect=_establish,
):
await group.connect_to_server(StdioServerParameters(command="test"))

Expand All @@ -284,6 +289,8 @@ async def test_client_session_group_connect_to_server_duplicate_tool_raises_erro
# Verify the duplicate tool was *not* added again (state should be unchanged)
assert len(group._tools) == 1 # Should still only have the original
assert group._tools[existing_tool_name] is not duplicate_tool # Ensure it's the original mock
mock_stack_new.aclose.assert_awaited()
assert mock_session_new not in group._session_exit_stacks


@pytest.mark.anyio
Expand Down
Loading