diff --git a/navi/core/container.py b/navi/core/container.py index 81390fa..fd87f46 100644 --- a/navi/core/container.py +++ b/navi/core/container.py @@ -67,6 +67,13 @@ except BaseException: pass if self.mcp_manager is not None: + # Stop the health-check loop first so it cannot enqueue commands + # onto a client's runner while disconnect_all is tearing it down + # (which would race the runner's transport teardown). + try: + await self.mcp_manager.stop_health_check() + except BaseException: + pass try: await self.mcp_manager.disconnect_all() except BaseException: diff --git a/navi/mcp/client.py b/navi/mcp/client.py index 9807441..feabb82 100644 --- a/navi/mcp/client.py +++ b/navi/mcp/client.py @@ -1,12 +1,13 @@ 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 anyio import httpx from mcp import ClientSession from mcp.client.sse import sse_client @@ -19,13 +20,34 @@ 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 or SSE), exposes - ``list_tools()`` and ``call_tool()``, and handles lifecycle + 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. """ @@ -33,172 +55,59 @@ def __init__(self, name: str, config: McpServerConfig) -> None: self.name = name self.config = config - self._session: ClientSession | None = None - self._exit_stack = AsyncExitStack() - self._connected = False + + # 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 - # Exponential backoff state for reconnect - self._last_reconnect_attempt: float | None = None - self._reconnect_backoff: float = 1.0 - self._max_reconnect_backoff: float = 30.0 - self._reconnect_jitter: float = 0.2 + # 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 and self._session is not None + 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, and store it.""" - if self._connected: - return - - try: - if self.config.is_stdio: - if not self.config.command: - raise ValueError("stdio transport requires 'command'") - params = StdioServerParameters( - command=self.config.command, - args=self.config.args or [], - env=self.config.env, - cwd=self.config.cwd, - ) - transport = await self._exit_stack.enter_async_context( - stdio_client(params) - ) - elif self.config.is_sse: - if not self.config.url: - raise ValueError("sse transport requires 'url'") - transport = await self._exit_stack.enter_async_context( - sse_client( - self.config.url, - headers=self.config.headers, - ) - ) - elif self.config.is_streamable_http: - if not self.config.url: - raise ValueError("streamable_http transport requires 'url'") - http_client = await self._exit_stack.enter_async_context( - httpx.AsyncClient(headers=self.config.headers or {}) - ) - transport = await self._exit_stack.enter_async_context( - streamable_http_client(self.config.url, http_client=http_client) - ) - # streamable_http_client returns (read, write, get_session_id). - # We only need read and write for ClientSession. - transport = transport[:2] - else: - raise ValueError(f"unknown transport: {self.config.transport}") - - read_stream, write_stream = transport - session = await self._exit_stack.enter_async_context( - ClientSession(read_stream, write_stream) - ) - init_result = await session.initialize() - self._instructions = init_result.instructions if hasattr(init_result, "instructions") else None - self._session = session - self._connected = True - # Reset backoff on successful connect - self._reconnect_backoff = 1.0 - self._last_reconnect_attempt = None - logger.info( - "MCP server %r connected (%s)", - self.name, - self.config.transport, - ) - except Exception: - await self._cleanup() - raise + """Open transport, initialise session. Idempotent when already connected.""" + await self._send("connect") async def disconnect(self) -> None: - """Close transport and reset state.""" - if not self._connected: - return - try: - # Do NOT use asyncio.wait_for here — it spawns a new task, - # and MCP transports (stdio, sse, streamable_http) use anyio - # task groups that require __aexit__ in the same task as __aenter__. - await self._cleanup() - except asyncio.CancelledError: - # Graceful shutdown during app teardown. - pass - - 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. - """ + """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 - - async def _cleanup(self) -> None: - try: - await self._exit_stack.aclose() - except SystemExit: - raise - except BaseException: - # Graceful shutdown noise from MCP transport task groups — - # anyio raises RuntimeError / BaseExceptionGroup when an async - # generator is closed from a different task than it was created in. - pass - finally: - self._session = None - self._instructions = None - self._connected = False - self._exit_stack = AsyncExitStack() - - def _check_backoff(self) -> bool: - """Return True if enough time has passed since the last reconnect attempt.""" - if self._last_reconnect_attempt is None: - return True - elapsed = time.monotonic() - self._last_reconnect_attempt - # Add jitter to prevent thundering herd - jitter = self._reconnect_backoff * self._reconnect_jitter * (2 * (time.monotonic() % 1) - 1) - return elapsed >= (self._reconnect_backoff + jitter) - - async def _ensure_connected(self) -> None: - """Reconnect if the underlying transport is dead, respecting backoff.""" - if self._connected and self._session is not None: - return - - if not self._check_backoff(): - remaining = self._reconnect_backoff - (time.monotonic() - (self._last_reconnect_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" - ) - - self._last_reconnect_attempt = time.monotonic() - logger.warning("MCP server %r disconnected, reconnecting...", self.name) - try: - await self._cleanup() - await self.connect() - except Exception: - # Double the backoff for the next attempt (capped at max) - self._reconnect_backoff = min( - self._reconnect_backoff * 2, - self._max_reconnect_backoff, - ) - raise + self._instructions = None async def list_tools(self) -> list[Tool]: """Return the tools exposed by the remote MCP server.""" - await self._ensure_connected() - try: - result = await self._session.list_tools() - except Exception: - await self._cleanup() - await self._ensure_connected() - result = await self._session.list_tools() - return list(result.tools) + 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. @@ -208,22 +117,222 @@ Text content is concatenated; images are reported as a placeholder. """ - await self._ensure_connected() + 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: - result = await self._session.call_tool(tool_name, arguments or {}) - except Exception: - await self._cleanup() - await self._ensure_connected() - result = await self._session.call_tool(tool_name, arguments or {}) + 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() - 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}]") + @staticmethod + def _resolve(fut: asyncio.Future | None, value: Any) -> None: + if fut is not None and not fut.done(): + fut.set_result(value) - is_error = getattr(result, "isError", False) - return "\n".join(parts), is_error + @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() \ No newline at end of file diff --git a/navi/mcp/manager.py b/navi/mcp/manager.py index 60ba80f..80641eb 100644 --- a/navi/mcp/manager.py +++ b/navi/mcp/manager.py @@ -220,7 +220,7 @@ await client.list_tools() except Exception: logger.warning("MCP server %r dropped during health check", name) - client.mark_disconnected() + await client.mark_disconnected() self._connected_status[name] = False await get_event_bus().publish( McpStatusUpdate(server_name=name, status="disconnected") diff --git a/tests/integration/test_mcp_integration.py b/tests/integration/test_mcp_integration.py index e19d76e..ac4398c 100644 --- a/tests/integration/test_mcp_integration.py +++ b/tests/integration/test_mcp_integration.py @@ -1,5 +1,6 @@ """Integration test for navi/mcp/ against a real stdio MCP server.""" +import asyncio import sys import pytest @@ -32,3 +33,40 @@ await client.disconnect() assert not client.connected + + +@pytest.mark.anyio +async def test_disconnect_from_different_task_than_connect(): + """Regression: connect in one task, disconnect in another. + + The MCP client SDK uses anyio task groups whose cancel scopes require + ``__aexit__`` in the same task as ``__aenter__``. Before the dedicated + runner-task refactor this raised ``RuntimeError: Attempted to exit cancel + scope in a different task than it was entered in``. The runner owns the + AsyncExitStack, so enter and exit always happen in its one task. + """ + cfg = McpServerConfig( + transport="stdio", + command=sys.executable, + args=["-m", "tests._mcp_test_server"], + ) + client = McpClient("test", cfg) + + async def _connect() -> None: + await client.connect() + + async def _use() -> None: + # A second task also exercises the session so the transport's anyio + # task group is genuinely live when teardown begins. + await client.list_tools() + + async def _disconnect() -> None: + await client.disconnect() + + await asyncio.ensure_future(_connect()) + assert client.connected + await asyncio.ensure_future(_use()) + # Teardown from a *different* task than the one that connected — must not + # raise the cross-task cancel-scope RuntimeError. + await asyncio.ensure_future(_disconnect()) + assert not client.connected