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()
