fast-agent-mcp 0.0.7__py3-none-any.whl

mcp_agent/mcp/mcp_connection_manager.py

@@ -0,0 +1,330 @@
+"""
+Manages the lifecycle of multiple MCP server connections.
+"""
+
+from datetime import timedelta
+import asyncio
+from typing import (
+    AsyncGenerator,
+    Callable,
+    Dict,
+    Optional,
+    TYPE_CHECKING,
+)
+
+from anyio import Event, create_task_group, Lock
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+
+from mcp import ClientSession
+from mcp.client.stdio import (
+    StdioServerParameters,
+    get_default_environment,
+)
+from mcp.client.sse import sse_client
+from mcp.types import JSONRPCMessage
+
+from mcp_agent.config import MCPServerSettings
+from mcp_agent.logging.logger import get_logger
+from mcp_agent.mcp.stdio import stdio_client_with_rich_stderr
+from mcp_agent.context_dependent import ContextDependent
+
+if TYPE_CHECKING:
+    from mcp_agent.mcp_server_registry import InitHookCallable, ServerRegistry
+    from mcp_agent.context import Context
+
+logger = get_logger(__name__)
+
+
+class ServerConnection:
+    """
+    Represents a long-lived MCP server connection, including:
+    - The ClientSession to the server
+    - The transport streams (via stdio/sse, etc.)
+    """
+
+    def __init__(
+        self,
+        server_name: str,
+        server_config: MCPServerSettings,
+        transport_context_factory: Callable[
+            [],
+            AsyncGenerator[
+                tuple[
+                    MemoryObjectReceiveStream[JSONRPCMessage | Exception],
+                    MemoryObjectSendStream[JSONRPCMessage],
+                ],
+                None,
+            ],
+        ],
+        client_session_factory: Callable[
+            [MemoryObjectReceiveStream, MemoryObjectSendStream, timedelta | None],
+            ClientSession,
+        ],
+        init_hook: Optional["InitHookCallable"] = None,
+    ):
+        self.server_name = server_name
+        self.server_config = server_config
+        self.session: ClientSession | None = None
+        self._client_session_factory = client_session_factory
+        self._init_hook = init_hook
+        self._transport_context_factory = transport_context_factory
+        # Signal that session is fully up and initialized
+        self._initialized_event = Event()
+
+        # Signal we want to shut down
+        self._shutdown_event = Event()
+
+    def request_shutdown(self) -> None:
+        """
+        Request the server to shut down. Signals the server lifecycle task to exit.
+        """
+        self._shutdown_event.set()
+
+    async def wait_for_shutdown_request(self) -> None:
+        """
+        Wait until the shutdown event is set.
+        """
+        await self._shutdown_event.wait()
+
+    async def initialize_session(self) -> None:
+        """
+        Initializes the server connection and session.
+        Must be called within an async context.
+        """
+
+        await self.session.initialize()
+
+        # If there's an init hook, run it
+        if self._init_hook:
+            logger.info(f"{self.server_name}: Executing init hook.")
+            self._init_hook(self.session, self.server_config.auth)
+
+        # Now the session is ready for use
+        self._initialized_event.set()
+
+    async def wait_for_initialized(self) -> None:
+        """
+        Wait until the session is fully initialized.
+        """
+        await self._initialized_event.wait()
+
+    def create_session(
+        self,
+        read_stream: MemoryObjectReceiveStream,
+        send_stream: MemoryObjectSendStream,
+    ) -> ClientSession:
+        """
+        Create a new session instance for this server connection.
+        """
+
+        read_timeout = (
+            timedelta(seconds=self.server_config.read_timeout_seconds)
+            if self.server_config.read_timeout_seconds
+            else None
+        )
+
+        session = self._client_session_factory(read_stream, send_stream, read_timeout)
+
+        # Make the server config available to the session for initialization
+        if hasattr(session, "server_config"):
+            session.server_config = self.server_config
+
+        self.session = session
+
+        return session
+
+
+async def _server_lifecycle_task(server_conn: ServerConnection) -> None:
+    """
+    Manage the lifecycle of a single server connection.
+    Runs inside the MCPConnectionManager's shared TaskGroup.
+    """
+    server_name = server_conn.server_name
+    try:
+        transport_context = server_conn._transport_context_factory()
+
+        async with transport_context as (read_stream, write_stream):
+            server_conn.create_session(read_stream, write_stream)
+
+            async with server_conn.session:
+                await server_conn.initialize_session()
+
+                await server_conn.wait_for_shutdown_request()
+
+    except Exception as exc:
+        logger.error(
+            f"{server_name}: Lifecycle task encountered an error: {exc}", exc_info=True
+        )
+        # If there's an error, we should also set the event so that
+        # 'get_server' won't hang
+        server_conn._initialized_event.set()
+        raise
+
+
+class MCPConnectionManager(ContextDependent):
+    """
+    Manages the lifecycle of multiple MCP server connections.
+    Integrates with the application context system for proper resource management.
+    """
+
+    def __init__(
+        self, server_registry: "ServerRegistry", context: Optional["Context"] = None
+    ):
+        super().__init__(context=context)
+        self.server_registry = server_registry
+        self.running_servers: Dict[str, ServerConnection] = {}
+        self._lock = Lock()
+
+    async def __aenter__(self):
+        current_task = asyncio.current_task()
+
+        # Get or create task group from context
+        if not hasattr(self.context, "_connection_task_group"):
+            self.context._connection_task_group = create_task_group()
+            self.context._connection_task_group_context = current_task.get_name()
+            await self.context._connection_task_group.__aenter__()
+
+        self._tg = self.context._connection_task_group
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Ensure clean shutdown of all connections before exiting."""
+        current_task = asyncio.current_task()
+
+        try:
+            # First request all servers to shutdown
+            await self.disconnect_all()
+
+            # Only clean up task group if we're in the original context
+            if (
+                hasattr(self.context, "_connection_task_group")
+                and current_task.get_name()
+                == self.context._connection_task_group_context
+            ):
+                await self.context._connection_task_group.__aexit__(
+                    exc_type, exc_val, exc_tb
+                )
+                delattr(self.context, "_connection_task_group")
+                delattr(self.context, "_connection_task_group_context")
+        except Exception as e:
+            logger.error(f"Error during connection manager shutdown: {e}")
+
+    async def launch_server(
+        self,
+        server_name: str,
+        client_session_factory: Callable[
+            [MemoryObjectReceiveStream, MemoryObjectSendStream, timedelta | None],
+            ClientSession,
+        ],
+        init_hook: Optional["InitHookCallable"] = None,
+    ) -> ServerConnection:
+        """
+        Connect to a server and return a RunningServer instance that will persist
+        until explicitly disconnected.
+        """
+        if not self._tg:
+            raise RuntimeError(
+                "MCPConnectionManager must be used inside an async context (i.e. 'async with' or after __aenter__)."
+            )
+
+        config = self.server_registry.registry.get(server_name)
+        if not config:
+            raise ValueError(f"Server '{server_name}' not found in registry.")
+
+        logger.debug(
+            f"{server_name}: Found server configuration=", data=config.model_dump()
+        )
+
+        def transport_context_factory():
+            if config.transport == "stdio":
+                server_params = StdioServerParameters(
+                    command=config.command,
+                    args=config.args,
+                    env={**get_default_environment(), **(config.env or {})},
+                )
+                # Create stdio client config with redirected stderr
+                return stdio_client_with_rich_stderr(server_params)
+            elif config.transport == "sse":
+                return sse_client(config.url)
+            else:
+                raise ValueError(f"Unsupported transport: {config.transport}")
+
+        server_conn = ServerConnection(
+            server_name=server_name,
+            server_config=config,
+            transport_context_factory=transport_context_factory,
+            client_session_factory=client_session_factory,
+            init_hook=init_hook or self.server_registry.init_hooks.get(server_name),
+        )
+
+        async with self._lock:
+            # Check if already running
+            if server_name in self.running_servers:
+                return self.running_servers[server_name]
+
+            self.running_servers[server_name] = server_conn
+            self._tg.start_soon(_server_lifecycle_task, server_conn)
+
+        logger.info(f"{server_name}: Up and running with a persistent connection!")
+        return server_conn
+
+    async def get_server(
+        self,
+        server_name: str,
+        client_session_factory: Callable,
+        init_hook: Optional["InitHookCallable"] = None,
+    ) -> ServerConnection:
+        """
+        Get a running server instance, launching it if needed.
+        """
+        # Get the server connection if it's already running
+        async with self._lock:
+            server_conn = self.running_servers.get(server_name)
+            if server_conn:
+                return server_conn
+
+        # Launch the connection
+        server_conn = await self.launch_server(
+            server_name=server_name,
+            client_session_factory=client_session_factory,
+            init_hook=init_hook,
+        )
+
+        # Wait until it's fully initialized, or an error occurs
+        await server_conn.wait_for_initialized()
+
+        # If the session is still None, it means the lifecycle task crashed
+        if not server_conn or not server_conn.session:
+            raise RuntimeError(
+                f"{server_name}: Failed to initialize server; check logs for errors."
+            )
+        return server_conn
+
+    async def disconnect_server(self, server_name: str) -> None:
+        """
+        Disconnect a specific server if it's running under this connection manager.
+        """
+        logger.info(f"{server_name}: Disconnecting persistent connection to server...")
+
+        async with self._lock:
+            server_conn = self.running_servers.pop(server_name, None)
+        if server_conn:
+            server_conn.request_shutdown()
+            logger.info(
+                f"{server_name}: Shutdown signal sent (lifecycle task will exit)."
+            )
+        else:
+            logger.info(
+                f"{server_name}: No persistent connection found. Skipping server shutdown"
+            )
+
+    async def disconnect_all(self) -> None:
+        """Disconnect all servers that are running under this connection manager."""
+        async with self._lock:
+            if not self.running_servers:
+                return
+
+            for name, conn in self.running_servers.items():
+                conn.request_shutdown()
+
+            self.running_servers.clear()

mcp_agent/mcp/stdio.py

@@ -0,0 +1,104 @@
+"""
+Custom implementation of stdio_client that handles stderr through rich console.
+"""
+
+from contextlib import asynccontextmanager
+import subprocess
+import anyio
+from anyio.streams.text import TextReceiveStream
+from mcp.client.stdio import StdioServerParameters, get_default_environment
+import mcp.types as types
+from mcp_agent.logging.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+@asynccontextmanager
+async def stdio_client_with_rich_stderr(server: StdioServerParameters):
+    """
+    Modified version of stdio_client that captures stderr and routes it through our rich console.
+    Follows the original pattern closely for reliability.
+
+    Args:
+        server: The server parameters for the stdio connection
+    """
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    # Open process with stderr piped for capture
+    process = await anyio.open_process(
+        [server.command, *server.args],
+        env=server.env if server.env is not None else get_default_environment(),
+        stderr=subprocess.PIPE,
+    )
+
+    if process.pid:
+        logger.debug(f"Started process '{server.command}' with PID: {process.pid}")
+
+    if process.returncode is not None:
+        logger.debug(f"return code (early){process.returncode}")
+        raise RuntimeError(
+            f"Process terminated immediately with code {process.returncode}"
+        )
+
+    async def stdout_reader():
+        assert process.stdout, "Opened process is missing stdout"
+        try:
+            async with read_stream_writer:
+                buffer = ""
+                async for chunk in TextReceiveStream(
+                    process.stdout,
+                    encoding=server.encoding,
+                    errors=server.encoding_error_handler,
+                ):
+                    lines = (buffer + chunk).split("\n")
+                    buffer = lines.pop()
+
+                    for line in lines:
+                        if not line:
+                            continue
+                        try:
+                            message = types.JSONRPCMessage.model_validate_json(line)
+                        except Exception as exc:
+                            await read_stream_writer.send(exc)
+                            continue
+
+                        await read_stream_writer.send(message)
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    async def stderr_reader():
+        assert process.stderr, "Opened process is missing stderr"
+        try:
+            async for chunk in TextReceiveStream(
+                process.stderr,
+                encoding=server.encoding,
+                errors=server.encoding_error_handler,
+            ):
+                if chunk.strip():
+                    # Let the logging system handle the formatting consistently
+                    logger.event("info", "mcpserver.stderr", chunk.rstrip(), None, {})
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    async def stdin_writer():
+        assert process.stdin, "Opened process is missing stdin"
+        try:
+            async with write_stream_reader:
+                async for message in write_stream_reader:
+                    json = message.model_dump_json(by_alias=True, exclude_none=True)
+                    await process.stdin.send(
+                        (json + "\n").encode(
+                            encoding=server.encoding,
+                            errors=server.encoding_error_handler,
+                        )
+                    )
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    # Use context managers to handle cleanup automatically
+    async with anyio.create_task_group() as tg, process:
+        tg.start_soon(stdout_reader)
+        tg.start_soon(stdin_writer)
+        tg.start_soon(stderr_reader)
+        yield read_stream, write_stream