from __future__ import annotations import asyncio import contextlib import logging import time from contextlib import AsyncExitStack from dataclasses import dataclass, field from datetime import timedelta 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, resolve_paths from .secrets import resolve_secrets logger = logging.getLogger(__name__) # ── Budgets ────────────────────────────────────────────────────────────────── # Every wait on this path is bounded, because a server that accepted a request # and went silent is indistinguishable from a slow one — and nothing above this # layer could tell them apart. Before these existed the wait was *infinite*: # with ``read_timeout_seconds=None`` the SDK waits inside # ``anyio.fail_after(None)``, the caller's future is never resolved, and the # runner task stays wedged for the life of the process. Every later reload then # blocked on that same client (``disconnect_all`` is sequential), so one silent # server made ``reload_tools`` look permanently broken. # # Two layers, deliberately ordered: the SDK-side read timeout fires first and # cleans up *inside* the runner, so the client stays usable; the caller-side # bound is the backstop, covering also a caller queued behind a stuck operation # (and a server that stops reading its requests, where the SDK never gets to its # read wait). It aborts the runner, which is rebuilt on the next command. _SESSION_READ_TIMEOUT_S = 30.0 """Per-request read timeout handed to ``ClientSession``: initialize + tools/list.""" _CALL_TOOL_READ_TIMEOUT_S = 240.0 """Per-request override for tools/call — a remote tool may legitimately take minutes.""" _CONNECT_TIMEOUT_S = 60.0 """Caller-side. Twice the SDK bound, so a silent handshake fails SDK-side first.""" _LIST_TOOLS_TIMEOUT_S = 45.0 """Caller-side, likewise above the SDK bound (see the retry in the runner).""" _CALL_TOOL_TIMEOUT_S = 300.0 _QUICK_OP_TIMEOUT_S = 15.0 """Ops that do no I/O — only the wait for their turn behind another op.""" _DISCONNECT_TIMEOUT_S = 15.0 """Transport teardown. One that will not finish is abandoned (see disconnect()).""" # httpx bounds for the HTTP transports. These are the SDK's own recommended # values (``mcp.shared._httpx_utils``), which we were not getting: we hand the # transport our own AsyncClient for its headers, and a bare ``httpx.AsyncClient`` # defaults to 5s for *everything* — so the streamable-http GET stream died every # time its server went quiet for five seconds. ``read`` must stay generous: a # streamable-http server holds the response stream open. _HTTP_TIMEOUT_S = 30.0 _HTTP_STREAM_READ_TIMEOUT_S = 300.0 class McpOperationTimeout(RuntimeError): """An MCP operation did not answer within its budget. Carries the server and op so the caller can name the culprit — on prod the log line ``mcp.op_timeout`` is the only thing that says *which* server was silent, because a server that never answers writes nothing itself. """ def __init__(self, server: str, op: str, timeout: float) -> None: super().__init__( f"MCP server {server!r} did not answer '{op}' within {timeout:.0f}s" ) self.server = server self.op = op self.timeout = timeout @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", _CONNECT_TIMEOUT_S) async def disconnect(self) -> None: """Close transport, stop the runner task, and reset state. Bounded, because ``McpManager.disconnect_all`` walks the whole pool and this is where a hang used to become *every* reload's problem: the teardown runs the MCP SDK's exit stack, which can wait on a server that stopped answering. A teardown that will not finish in :data:`_DISCONNECT_TIMEOUT_S` is abandoned — the runner is cancelled and its transport dies with the task. """ 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 wait around it is bounded. try: await asyncio.wait_for(asyncio.shield(runner), _DISCONNECT_TIMEOUT_S) except TimeoutError: logger.warning( "mcp.disconnect_timeout server=%s after=%.0fs — abandoning the runner", self.name, _DISCONNECT_TIMEOUT_S, ) self._abort_runner() except asyncio.CancelledError: # The caller is being cancelled (shutdown, or an outer deadline # such as the budget on the MCP reload phase). Re-raising is # what lets that deadline actually fire — and the runner, still # shielded and possibly stuck in this same teardown, is # cancelled on the way out. self._abort_runner() raise 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( # type: ignore[return-value] "list_tools", _LIST_TOOLS_TIMEOUT_S ) 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", _CALL_TOOL_TIMEOUT_S, 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", _QUICK_OP_TIMEOUT_S) # ── 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()) def _abort_runner(self) -> None: """Cancel the runner task and drop our handles to it. The runner closes the transport in its own ``finally``, so cancelling it is what actually frees the stream. Clearing the handles here — rather than letting the dying task do it — is what keeps the next command from queueing behind a runner that is already unwinding, and the mirrored state has to be cleared here too: the dying runner's own teardown is ignored (it is no longer the current one), so nothing else would. """ runner = self._runner_task self._runner_task = None self._queue = None self._connected = False self._instructions = None if runner is not None and not runner.done(): runner.cancel() async def _send(self, op: str, timeout: float, **payload: Any) -> Any: """Enqueue a command and await its result, for at most *timeout* seconds. *timeout* is required so no future op can quietly reintroduce an unbounded wait; the module constants above say what each op is worth. """ 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)) try: return await asyncio.wait_for(fut, timeout) except TimeoutError: # The runner is somewhere inside this op and will not come back: # abort it, or every later command queues behind the same silence. logger.warning( "mcp.op_timeout server=%s op=%s after=%.0fs", self.name, op, timeout ) self._abort_runner() raise McpOperationTimeout(self.name, op, timeout) from None except asyncio.CancelledError: # The caller gave up (outer deadline, stop, shutdown). Nothing will # consume this result, and a runner still waiting on a silent server # would stay wedged for the life of the process. fut.cancel() self._abort_runner() raise 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 me = asyncio.current_task() exit_stack = AsyncExitStack() session: ClientSession | None = None connected: bool = False instructions: str | None = None backoff: float = 1.0 last_attempt: float | None = None def publish_state() -> None: """Mirror this runner's state — unless this runner was replaced. A runner that a timed-out operation aborted keeps unwinding its own teardown for a moment, and its dying state must not be written over the fresh runner a later command has already started. """ if self._runner_task is me: self._connected = connected self._instructions = instructions 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 publish_state() async def open_transport() -> None: nonlocal session, instructions # Config values are made concrete here and nowhere else: project- # relative paths become absolute, and `${VAR}` credentials are read # from the environment. Keeping both at transport-open means the # stored config — and anything that serialises it back to disk — # never holds this machine's paths or a live secret. cfg = resolve_paths(resolve_secrets(self.config, server=self.name)) 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, timeout=_HTTP_TIMEOUT_S, sse_read_timeout=_HTTP_STREAM_READ_TIMEOUT_S, ) ) 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 {}, # Never the httpx default: 5s on the GET stream read # tears the transport down on an idle server. timeout=httpx.Timeout( _HTTP_TIMEOUT_S, read=_HTTP_STREAM_READ_TIMEOUT_S ), ) ) 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, # Without this the SDK waits forever on every request. read_timeout_seconds=timedelta(seconds=_SESSION_READ_TIMEOUT_S), ) ) 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 publish_state() 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 publish_state() 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 {}, read_timeout_seconds=timedelta( seconds=_CALL_TOOL_READ_TIMEOUT_S ), ) except Exception: await close_all() await ensure_connected() result = await session.call_tool( # type: ignore[union-attr] tool_name, arguments or {}, read_timeout_seconds=timedelta( seconds=_CALL_TOOL_READ_TIMEOUT_S ), ) 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: task = asyncio.current_task() if task is None or task.cancelling(): # Runner is really being cancelled (app teardown). # Release the in-flight caller, then let close_all run # in finally. self._cancel(cmd.fut) raise # A failed transport connect inside the MCP SDK's anyio # cancel scopes can surface as CancelledError without an # actual task.cancel() (seen with streamable_http against a # dead endpoint). Treat it as an ordinary connect failure: # one unreachable MCP server must not take down the whole # startup (McpManager.load_all only catches Exception). publish_state() self._fail(cmd.fut, RuntimeError( f"MCP server {self.name!r} connect failed (anyio CancelledError)" )) except Exception as exc: # On a connect/ensure failure, mirror the disconnected # state before surfacing the error to the caller. publish_state() self._fail(cmd.fut, exc) finally: await close_all() # Runner is gone — fail anything still queued so callers waiting # in _send don't hang forever. while not queue.empty(): try: cmd = queue.get_nowait() except Exception: break self._fail(cmd.fut, RuntimeError(f"MCP runner for {self.name!r} exited")) @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()