from __future__ import annotations import asyncio import contextlib import logging import time from contextlib import AsyncExitStack from dataclasses import dataclass, field from typing import Any import httpx from mcp import ClientSession from mcp.client.sse import sse_client from mcp.client.stdio import StdioServerParameters, stdio_client from mcp.client.streamable_http import streamable_http_client from mcp.types import Tool from .config import McpServerConfig logger = logging.getLogger(__name__) @dataclass class _Cmd: """A command enqueued to the client's runner task.""" op: str fut: asyncio.Future | None payload: dict[str, Any] = field(default_factory=dict) class McpClient: """Lightweight wrapper around the official Python MCP client SDK. Manages a single server connection (stdio, SSE, or streamable-http), exposes ``list_tools()`` and ``call_tool()``, and handles lifecycle (connect / disconnect). **Threading model.** The MCP client SDK (``stdio_client`` / ``sse_client`` / ``streamable_http_client`` and ``ClientSession``) is built on anyio task groups, whose cancel scopes require that ``__aenter__`` and ``__aexit__`` run in the *same* asyncio task. To guarantee that, every client owns a single long-running **runner task** that holds the ``AsyncExitStack`` and performs all transport enter/exit (connect, disconnect, reconnect) as well as every ``list_tools`` / ``call_tool``. The public methods just enqueue a command onto a queue and await its result Future, so callers can invoke them from any task (lifespan, health-check, request handler) without crossing cancel-scope task boundaries. Reconnect uses exponential backoff (base 1s, max 30s, ±20% jitter) so that a flapping server does not hammer the transport. """ def __init__(self, name: str, config: McpServerConfig) -> None: self.name = name self.config = config # Mirrored snapshot of the runner's state, updated from inside the # runner task. asyncio is single-threaded, so these are safe to read # synchronously from any task via the properties below. self._connected: bool = False self._instructions: str | None = None # Runner task + command queue. Lazily created on first use. self._runner_task: asyncio.Task | None = None self._queue: asyncio.Queue[_Cmd] | None = None # ── Public, sync state ──────────────────────────────────────────────── @property def connected(self) -> bool: return self._connected @property def instructions(self) -> str | None: """Server-provided instructions from MCP initialize handshake.""" return self._instructions # ── Public, async API (enqueue + await) ─────────────────────────────── async def connect(self) -> None: """Open transport, initialise session. Idempotent when already connected.""" await self._send("connect") async def disconnect(self) -> None: """Close transport, stop the runner task, and reset state.""" runner = self._runner_task queue = self._queue # Ask the runner to stop. It closes the transport *inside its own # task* — the anyio cancel scope is exited where it was entered. if queue is not None and runner is not None and not runner.done(): with contextlib.suppress(Exception): queue.put_nowait(_Cmd("disconnect", None, {})) if runner is not None and not runner.done(): # Shield so a lifespan cancellation does not interrupt the # runner's own transport teardown; the runner finishes in its # own task regardless. try: await asyncio.shield(runner) except BaseException: pass self._runner_task = None self._queue = None self._connected = False self._instructions = None async def list_tools(self) -> list[Tool]: """Return the tools exposed by the remote MCP server.""" return await self._send("list_tools") # type: ignore[return-value] async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None = None) -> tuple[str, bool]: """Execute a remote tool and return its output as a string. Returns (output_text, is_error). ``is_error`` comes from the MCP ``CallToolResult.isError`` field so the caller can set ``success=False``. Text content is concatenated; images are reported as a placeholder. """ return await self._send( # type: ignore[return-value] "call_tool", tool_name=tool_name, arguments=arguments ) async def mark_disconnected(self) -> None: """Force the client into a disconnected state without closing transport. Used by the health-check loop when a server drops silently: the next operation will reconnect (with backoff). """ await self._send("mark_disconnected") # ── Runner plumbing ─────────────────────────────────────────────────── def _ensure_runner(self) -> None: """Lazily start the runner task + queue if not already running.""" if self._runner_task is not None and not self._runner_task.done(): return loop = asyncio.get_running_loop() self._queue = asyncio.Queue() self._runner_task = loop.create_task(self._runner()) async def _send(self, op: str, **payload: Any) -> Any: """Enqueue a command and await its result.""" self._ensure_runner() loop = asyncio.get_running_loop() fut: asyncio.Future = loop.create_future() assert self._queue is not None await self._queue.put(_Cmd(op, fut, payload)) return await fut async def _runner(self) -> None: """Owns the AsyncExitStack; processes commands until ``disconnect``. All transport ``enter_async_context`` / ``aclose`` calls happen here, in this single task, so anyio cancel scopes are always entered and exited by the same task. """ queue = self._queue assert queue is not None exit_stack = AsyncExitStack() session: ClientSession | None = None connected: bool = False instructions: str | None = None backoff: float = 1.0 last_attempt: float | None = None async def close_all() -> None: nonlocal session, connected, instructions try: await exit_stack.aclose() except SystemExit: raise except BaseException: # anyio raises RuntimeError / BaseExceptionGroup when an # async generator is closed; we are tearing down anyway. pass finally: session = None connected = False instructions = None self._connected = False self._instructions = None async def open_transport() -> None: nonlocal session, instructions cfg = self.config try: if cfg.is_stdio: if not cfg.command: raise ValueError("stdio transport requires 'command'") params = StdioServerParameters( command=cfg.command, args=cfg.args or [], env=cfg.env, cwd=cfg.cwd, ) transport = await exit_stack.enter_async_context(stdio_client(params)) elif cfg.is_sse: if not cfg.url: raise ValueError("sse transport requires 'url'") transport = await exit_stack.enter_async_context( sse_client(cfg.url, headers=cfg.headers) ) elif cfg.is_streamable_http: if not cfg.url: raise ValueError("streamable_http transport requires 'url'") http_client = await exit_stack.enter_async_context( httpx.AsyncClient(headers=cfg.headers or {}) ) transport = await exit_stack.enter_async_context( streamable_http_client(cfg.url, http_client=http_client) ) # streamable_http_client returns (read, write, get_session_id). transport = transport[:2] else: raise ValueError(f"unknown transport: {cfg.transport}") read_stream, write_stream = transport session = await exit_stack.enter_async_context( ClientSession(read_stream, write_stream) ) init_result = await session.initialize() instructions = getattr(init_result, "instructions", None) except BaseException: # Roll back any half-entered contexts before propagating. await close_all() raise def check_backoff() -> bool: nonlocal last_attempt if last_attempt is None: return True elapsed = time.monotonic() - last_attempt jitter = backoff * 0.2 * (2 * (time.monotonic() % 1) - 1) return elapsed >= (backoff + jitter) async def ensure_connected() -> None: nonlocal connected, backoff, last_attempt, instructions if connected and session is not None: return if not check_backoff(): remaining = backoff - (time.monotonic() - (last_attempt or 0)) logger.warning( "MCP server %r reconnect blocked by backoff (%.1fs remaining)", self.name, max(0, remaining), ) raise RuntimeError( f"MCP server {self.name!r} is disconnected and reconnect is throttled" ) last_attempt = time.monotonic() logger.warning("MCP server %r disconnected, reconnecting...", self.name) await close_all() await open_transport() # raises on failure -> caller retries later connected = True backoff = 1.0 last_attempt = None self._connected = True self._instructions = instructions logger.info("MCP server %r connected (%s)", self.name, self.config.transport) try: while True: cmd = await queue.get() try: if cmd.op == "connect": await ensure_connected() self._resolve(cmd.fut, None) elif cmd.op == "disconnect": break # finally -> close_all elif cmd.op == "mark_disconnected": connected = False self._connected = False self._resolve(cmd.fut, None) elif cmd.op == "list_tools": await ensure_connected() try: result = await session.list_tools() # type: ignore[union-attr] except Exception: await close_all() await ensure_connected() result = await session.list_tools() # type: ignore[union-attr] self._resolve(cmd.fut, list(result.tools)) elif cmd.op == "call_tool": tool_name = cmd.payload["tool_name"] arguments = cmd.payload.get("arguments") await ensure_connected() try: result = await session.call_tool( # type: ignore[union-attr] tool_name, arguments or {} ) except Exception: await close_all() await ensure_connected() result = await session.call_tool( # type: ignore[union-attr] tool_name, arguments or {} ) parts: list[str] = [] for item in result.content: if item.type == "text": parts.append(item.text) elif item.type == "image": parts.append(f"[image: {item.mimeType} ({len(item.data)} bytes)]") else: parts.append(f"[{item.type}]") is_error = getattr(result, "isError", False) self._resolve(cmd.fut, ("\n".join(parts), is_error)) else: self._fail(cmd.fut, ValueError(f"unknown op: {cmd.op}")) except asyncio.CancelledError: # Runner is being cancelled (e.g. app teardown). Release # the in-flight caller, then let close_all run in finally. self._cancel(cmd.fut) raise except Exception as exc: # On a connect/ensure failure, mirror the disconnected # state before surfacing the error to the caller. self._connected = connected self._instructions = instructions self._fail(cmd.fut, exc) finally: await close_all() @staticmethod def _resolve(fut: asyncio.Future | None, value: Any) -> None: if fut is not None and not fut.done(): fut.set_result(value) @staticmethod def _fail(fut: asyncio.Future | None, exc: BaseException) -> None: if fut is not None and not fut.done(): fut.set_exception(exc) @staticmethod def _cancel(fut: asyncio.Future | None) -> None: if fut is not None and not fut.done(): fut.cancel()