Multi-session client support in mcp.client #321

import asyncio
from dataclasses import dataclass
from contextlib import AsyncExitStack, AbstractAsyncContextManager

from mcp import ClientSession, StdioServerParameters, types
from mcp.client.stdio import stdio_client

@dataclass(frozen=True)
class MCPClient(AbstractAsyncContextManager):
    _exit_stack: AsyncExitStack
    _sessions_by_name: dict[str, ClientSession]

    @classmethod
    async def start(cls, servers_params_by_name: dict[str, StdioServerParameters]) -> "MCPClient":
        async def get_session(exit_stack: AsyncExitStack, server_params: StdioServerParameters) -> ClientSession:
            read, write = await exit_stack.enter_async_context(stdio_client(server_params))
            session = await exit_stack.enter_async_context(ClientSession(read, write))
            await session.initialize()
            return session

        async with AsyncExitStack() as exit_stack:
            sessions = await asyncio.gather(
                *(get_session(exit_stack, server_params) for server_params in servers_params_by_name.values())
            )
            return cls(exit_stack.pop_all(), {
                name: session
                for name, session in zip(servers_params_by_name.keys(), sessions, strict=True)
            })

        assert False  # @TODO: Why is this necessary?