runtime: MCP stdio client and per-process connection manager

src/plyngent/runtime/__init__.py

@@ -1,5 +1,11 @@
+from plyngent.config.models import McpConfig as McpConfig
+
 from .client_factory import InvalidHttpTimeoutError as InvalidHttpTimeoutError
 from .client_factory import ProviderNotSupportedError as ProviderNotSupportedError
 from .client_factory import create_client as create_client
 from .client_factory import normalize_http_timeout as normalize_http_timeout
 from .client_factory import provider_to_openai_config as provider_to_openai_config
+from .mcp_client import McpConnectionError as McpConnectionError
+from .mcp_client import McpManager as McpManager
+from .mcp_client import McpServerConnection as McpServerConnection
+from .mcp_client import McpTool as McpTool

src/plyngent/runtime/mcp_client.py

@@ -0,0 +1,361 @@
+"""MCP client: Model Context Protocol over stdio (newline-delimited JSON-RPC).
+
+One :class:`McpServerConnection` owns one spawned server subprocess and its
+initialize handshake. :class:`McpManager` owns all configured connections for
+the process: start, tool listing, tool calls, reconnect, shutdown.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import importlib.metadata
+import json
+import os
+from collections import deque
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, cast
+
+if TYPE_CHECKING:
+    from plyngent.config.models import McpConfig, McpServerConfig
+
+MCP_PROTOCOL_VERSION = "2025-06-18"
+MCP_CLIENT_NAME = "plyngent"
+DEFAULT_MCP_TIMEOUT = 30.0
+_STDERR_TAIL_LINES = 32
+_CLOSE_GRACE_SECONDS = 2.0
+_MAX_TOOL_PAGES = 1024
+
+type Json = dict[str, object]
+
+
+def _json_dict(raw: object) -> Json:
+    """Narrow an untyped JSON payload to a string-keyed object mapping."""
+    return cast("Json", raw) if isinstance(raw, dict) else {}
+
+
+def _json_list(raw: object) -> list[object]:
+    """Narrow an untyped JSON payload to a list of objects."""
+    return cast("list[object]", raw) if isinstance(raw, list) else []
+
+
+def _client_version() -> str:
+    try:
+        return importlib.metadata.version("plyngent")
+    except importlib.metadata.PackageNotFoundError:
+        return "0.0.0"
+
+
+class McpConnectionError(RuntimeError):
+    """Server process failed to spawn, handshake, or answer in time."""
+
+
+@dataclass(frozen=True, slots=True)
+class McpTool:
+    """One tool advertised by a server (``tools/list`` entry)."""
+
+    name: str
+    description: str
+    input_schema: Json
+
+    @classmethod
+    def from_raw(cls, raw: object) -> McpTool:
+        entry = _json_dict(raw)
+        return cls(
+            name=str(entry.get("name", "")),
+            description=str(entry.get("description", "") or ""),
+            input_schema=_json_dict(entry.get("inputSchema")) or {"type": "object"},
+        )
+
+
+def _format_content(result: Json) -> str:
+    """Flatten a ``tools/call`` result payload into model-facing text."""
+    parts: list[str] = []
+    for item in _json_list(result.get("content")):
+        entry = _json_dict(item)
+        kind = str(entry.get("type", ""))
+        if kind == "text":
+            text = entry.get("text")
+            if isinstance(text, str):
+                parts.append(text)
+        elif kind == "resource":
+            parts.append(f"[resource] {json.dumps(entry, ensure_ascii=False)}")
+        else:
+            mime = entry.get("mimeType", "?")
+            parts.append(f"[{kind or 'unknown'}: {mime}]")
+    return "\n".join(parts)
+
+
+class McpServerConnection:
+    """One server subprocess with an active MCP session (stdio transport)."""
+
+    name: str
+    error: str | None
+    tools: list[McpTool]
+    tools_stale: bool
+
+    _config: McpServerConfig
+    _proc: asyncio.subprocess.Process | None
+    _reader_task: asyncio.Task[None] | None
+    _pending: dict[int, asyncio.Future[Json]]
+    _next_id: int
+    _stderr_tail: deque[str]
+
+    def __init__(self, name: str, config: McpServerConfig) -> None:
+        self.name = name
+        self._config = config
+        self._proc = None
+        self._reader_task = None
+        self._pending = {}
+        self._next_id = 1
+        self._stderr_tail = deque(maxlen=_STDERR_TAIL_LINES)
+        self.error = None
+        self.tools = []
+        self.tools_stale = False
+
+    @property
+    def status(self) -> str:
+        if self.error is not None:
+            return f"error: {self.error}"
+        if self._proc is not None:
+            return "connected (tools changed; /mcp reconnect)" if self.tools_stale else "connected"
+        return "not started"
+
+    @property
+    def connected(self) -> bool:
+        return self._proc is not None and self.error is None
+
+    def stderr_tail(self) -> list[str]:
+        return list(self._stderr_tail)
+
+    async def start(self) -> None:
+        """Spawn the server process and run the initialize handshake."""
+        if self._proc is not None:
+            return
+        try:
+            self._proc = await asyncio.create_subprocess_exec(
+                self._config.command,
+                *self._config.args,
+                stdin=asyncio.subprocess.PIPE,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                env=os.environ | self._config.env,
+                cwd=self._config.cwd or None,
+            )
+        except (OSError, ValueError) as exc:
+            self.error = f"spawn {self._config.command!r}: {exc}"
+            self._proc = None
+            return
+        self._reader_task = asyncio.create_task(self._read_loop())
+        try:
+            await self._initialize()
+            self.tools = await self.list_tools()
+        except (McpConnectionError, TimeoutError) as exc:
+            self.error = f"initialize: {exc}"
+            await self.aclose()
+
+    async def aclose(self) -> None:
+        """Terminate the server subprocess and fail pending requests."""
+        proc = self._proc
+        self._proc = None
+        if self._reader_task is not None:
+            _ = self._reader_task.cancel()
+            self._reader_task = None
+        if proc is not None and proc.returncode is None:
+            _ = proc.terminate()
+            try:
+                _ = await asyncio.wait_for(proc.wait(), _CLOSE_GRACE_SECONDS)
+            except TimeoutError:
+                proc.kill()
+                _ = await proc.wait()
+        for future in self._pending.values():
+            if not future.done():
+                future.set_exception(McpConnectionError("connection closed"))
+        self._pending.clear()
+
+    async def _read_loop(self) -> None:
+        proc = self._proc
+        if proc is None or proc.stdout is None or proc.stderr is None:
+            return
+        stderr_task = asyncio.create_task(self._drain_stderr(proc.stderr))
+        try:
+            while True:
+                line = await proc.stdout.readline()
+                if not line:
+                    break
+                self._dispatch(line)
+        finally:
+            _ = stderr_task.cancel()
+            for future in self._pending.values():
+                if not future.done():
+                    future.set_exception(McpConnectionError("server closed the stream"))
+            self._pending.clear()
+
+    async def _drain_stderr(self, stderr: asyncio.StreamReader) -> None:
+        while True:
+            line = await stderr.readline()
+            if not line:
+                return
+            self._stderr_tail.append(line.decode(errors="replace").rstrip())
+
+    def _dispatch(self, line: bytes) -> None:
+        try:
+            payload = _json_dict(cast("object", json.loads(line)))
+        except json.JSONDecodeError:
+            return
+        message_id = payload.get("id")
+        if message_id is not None:
+            if isinstance(message_id, bool) or not isinstance(message_id, (int, str)):
+                return
+            request_id = int(message_id)
+            future = self._pending.pop(request_id, None)
+            if future is None or future.done():
+                return
+            error = payload.get("error")
+            if error is not None:
+                err = _json_dict(error)
+                detail = str(err.get("message", "unknown error")) if err else str(error)
+                future.set_exception(McpConnectionError(detail))
+            else:
+                future.set_result(_json_dict(payload.get("result")))
+            return
+        # Notifications carry no id; only tool-list changes matter to the host.
+        if payload.get("method") == "notifications/tools/list_changed":
+            self.tools_stale = True
+
+    def _send(self, payload: Json) -> None:
+        proc = self._proc
+        stdin = proc.stdin if proc is not None else None
+        if stdin is None:
+            msg = "server process is not running"
+            raise McpConnectionError(msg)
+        _ = stdin.write(json.dumps(payload, ensure_ascii=False).encode() + b"\n")
+
+    async def _request(self, method: str, params: Json | None = None) -> Json:
+        timeout = self._config.timeout or DEFAULT_MCP_TIMEOUT
+        future: asyncio.Future[Json] = asyncio.get_running_loop().create_future()
+        request_id = self._next_id
+        self._next_id += 1
+        self._pending[request_id] = future
+        payload: Json = {"jsonrpc": "2.0", "id": request_id, "method": method}
+        if params is not None:
+            payload["params"] = params
+        try:
+            self._send(payload)
+        except McpConnectionError:
+            _ = self._pending.pop(request_id, None)
+            raise
+        try:
+            return await asyncio.wait_for(future, timeout)
+        except TimeoutError:
+            _ = self._pending.pop(request_id, None)
+            msg = f"{method} timed out after {timeout}s"
+            raise McpConnectionError(msg) from None
+
+    async def _notify(self, method: str, params: Json | None = None) -> None:
+        payload: Json = {"jsonrpc": "2.0", "method": method}
+        if params is not None:
+            payload["params"] = params
+        self._send(payload)
+
+    async def _initialize(self) -> None:
+        _ = await self._request(
+            "initialize",
+            {
+                "protocolVersion": MCP_PROTOCOL_VERSION,
+                "capabilities": {},
+                "clientInfo": {"name": MCP_CLIENT_NAME, "version": _client_version()},
+            },
+        )
+        await self._notify("notifications/initialized")
+
+    async def list_tools(self) -> list[McpTool]:
+        """Fetch the full tool list, following ``nextCursor`` pagination."""
+        tools: list[McpTool] = []
+        cursor: str | None = None
+        for _ in range(_MAX_TOOL_PAGES):
+            params: Json = {} if cursor is None else {"cursor": cursor}
+            result = await self._request("tools/list", params)
+            tools.extend(McpTool.from_raw(item) for item in _json_list(result.get("tools")))
+            next_cursor = result.get("nextCursor")
+            if not isinstance(next_cursor, str) or not next_cursor:
+                return tools
+            cursor = next_cursor
+        return tools
+
+    async def call_tool(self, name: str, arguments: Json) -> str:
+        """Run ``tools/call`` and return the flattened text result.
+
+        Tool-level failures (``isError``) come back as ``error: …`` text, not
+        exceptions: the protocol round-trip itself succeeded.
+        """
+        result = await self._request("tools/call", {"name": name, "arguments": arguments})
+        text = _format_content(result)
+        if result.get("isError") is True:
+            return f"error: {text or 'tool reported failure without content'}"
+        return text
+
+
+class McpManager:
+    """All configured MCP server connections for this process."""
+
+    _config: McpConfig
+    _connections: dict[str, McpServerConnection]
+
+    def __init__(self, config: McpConfig) -> None:
+        self._config = config
+        self._connections = {}
+
+    def enabled_servers(self) -> dict[str, McpServerConfig]:
+        """Configured servers minus the ``disable`` list, in config order."""
+        disabled = {name.strip() for name in self._config.disable}
+        return {name: cfg for name, cfg in self._config.servers.items() if name not in disabled}
+
+    async def ensure_started(self) -> None:
+        """Connect every enabled server that has no live connection yet."""
+        for name, cfg in self.enabled_servers().items():
+            connection = self._connections.get(name)
+            if connection is not None and connection.connected:
+                continue
+            if connection is not None:
+                await connection.aclose()
+            connection = McpServerConnection(name, cfg)
+            self._connections[name] = connection
+            await connection.start()
+
+    async def restart(self, config: McpConfig) -> None:
+        """Drop every connection and adopt a new config (e.g. /mcp reconnect)."""
+        await self.aclose()
+        self._config = config
+
+    async def aclose(self) -> None:
+        for connection in self._connections.values():
+            await connection.aclose()
+        self._connections.clear()
+
+    def statuses(self) -> list[tuple[str, str, int]]:
+        """(server, status, tool_count) rows for diagnostics surfaces."""
+        rows: list[tuple[str, str, int]] = []
+        for name in self.enabled_servers():
+            connection = self._connections.get(name)
+            if connection is None:
+                rows.append((name, "not started", 0))
+            else:
+                rows.append((name, connection.status, len(connection.tools)))
+        return rows
+
+    def connections(self) -> list[McpServerConnection]:
+        return list(self._connections.values())
+
+    async def refresh_tools(self) -> None:
+        """Re-run tools/list on every connected server."""
+        for connection in self._connections.values():
+            if connection.connected:
+                connection.tools = await connection.list_tools()
+                connection.tools_stale = False
+
+    async def call(self, server: str, tool: str, arguments: Json) -> str:
+        connection = self._connections.get(server)
+        if connection is None or not connection.connected:
+            msg = f"server {server!r} is not connected"
+            raise McpConnectionError(msg)
+        return await connection.call_tool(tool, arguments)

tests/test_runtime/mcp_fake_server.py

@@ -0,0 +1,96 @@
+"""Minimal MCP stdio server used by the client tests (echo + sleep tools)."""
+
+from __future__ import annotations
+
+import json
+import sys
+
+
+def respond(request_id: int, result: object) -> None:
+    sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": request_id, "result": result}) + "\n")
+    sys.stdout.flush()
+
+
+def main() -> None:
+    for line in sys.stdin:
+        try:
+            payload = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        method = payload.get("method", "")
+        request_id = payload.get("id")
+        if request_id is None:
+            continue  # notification (notifications/initialized)
+        if method == "initialize":
+            respond(
+                int(request_id),
+                {
+                    "protocolVersion": "2025-06-18",
+                    "capabilities": {"tools": {}},
+                    "serverInfo": {"name": "fake-mcp", "version": "0.1.0"},
+                },
+            )
+        elif method == "tools/list":
+            respond(
+                int(request_id),
+                {
+                    "tools": [
+                        {
+                            "name": "echo",
+                            "description": "Echo the given text back.",
+                            "inputSchema": {
+                                "type": "object",
+                                "properties": {"text": {"type": "string"}},
+                                "required": ["text"],
+                            },
+                        },
+                        {
+                            "name": "fail",
+                            "description": "Always reports a tool error.",
+                            "inputSchema": {"type": "object", "properties": {}},
+                        },
+                        {
+                            "name": "slow",
+                            "description": "Sleeps past any client timeout.",
+                            "inputSchema": {"type": "object", "properties": {}},
+                        },
+                    ]
+                },
+            )
+        elif method == "tools/call":
+            name = (payload.get("params") or {}).get("name")
+            if name == "echo":
+                text = str((payload.get("params") or {}).get("arguments", {}).get("text", ""))
+                respond(
+                    int(request_id),
+                    {
+                        "content": [{"type": "text", "text": f"echo: {text}"}],
+                        "isError": False,
+                    },
+                )
+            elif name == "fail":
+                respond(
+                    int(request_id),
+                    {"content": [{"type": "text", "text": "boom"}], "isError": True},
+                )
+            else:
+                import time
+
+                time.sleep(5)
+                respond(int(request_id), {"content": [], "isError": False})
+        else:
+            sys.stdout.write(
+                json.dumps(
+                    {
+                        "jsonrpc": "2.0",
+                        "id": request_id,
+                        "error": {"code": -32601, "message": f"unknown method {method}"},
+                    }
+                )
+                + "\n"
+            )
+            sys.stdout.flush()
+
+
+if __name__ == "__main__":
+    main()

tests/test_runtime/test_mcp_client.py

@@ -0,0 +1,106 @@
+"""MCP stdio client + manager against the bundled fake server subprocess."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+
+from plyngent.config.models import McpConfig, McpServerConfig
+from plyngent.runtime.mcp_client import (
+    McpConnectionError,
+    McpManager,
+    McpServerConnection,
+)
+
+_FAKE_SERVER = Path(__file__).parent / "mcp_fake_server.py"
+
+
+def _server_config(**overrides: object) -> McpServerConfig:
+    base: dict[str, object] = {
+        "command": sys.executable,
+        "args": [str(_FAKE_SERVER)],
+        "timeout": 5.0,
+    }
+    return McpServerConfig(**{**base, **overrides})  # type: ignore[arg-type]
+
+
+async def test_connection_handshake_lists_and_calls_tools() -> None:
+    connection = McpServerConnection("fake", _server_config())
+    await connection.start()
+    try:
+        assert connection.connected
+        assert connection.error is None
+        assert [tool.name for tool in connection.tools] == ["echo", "fail", "slow"]
+        echo = connection.tools[0]
+        assert echo.description == "Echo the given text back."
+        assert echo.input_schema["type"] == "object"
+
+        result = await connection.call_tool("echo", {"text": "hello"})
+        assert result == "echo: hello"
+    finally:
+        await connection.aclose()
+
+
+async def test_tool_error_returns_error_text() -> None:
+    connection = McpServerConnection("fake", _server_config())
+    await connection.start()
+    try:
+        result = await connection.call_tool("fail", {})
+        assert result == "error: boom"
+    finally:
+        await connection.aclose()
+
+
+async def test_request_timeout_raises() -> None:
+    connection = McpServerConnection("fake", _server_config(timeout=0.2))
+    await connection.start()
+    try:
+        with pytest.raises(McpConnectionError, match="timed out"):
+            await connection.call_tool("slow", {})
+    finally:
+        await connection.aclose()
+
+
+async def test_spawn_failure_records_error() -> None:
+    connection = McpServerConnection("bad", _server_config(command="/nonexistent/plyngent-mcp-binary"))
+    await connection.start()
+    assert not connection.connected
+    assert connection.error is not None
+    assert "spawn" in connection.status
+    await connection.aclose()
+
+
+async def test_close_terminates_subprocess() -> None:
+    connection = McpServerConnection("fake", _server_config())
+    await connection.start()
+    proc = connection._proc
+    assert proc is not None
+    await connection.aclose()
+    assert connection._proc is None
+    assert proc.returncode is not None
+
+
+async def test_manager_start_status_and_call() -> None:
+    manager = McpManager(McpConfig(servers={"fake": _server_config()}, disable=[]))
+    await manager.ensure_started()
+    try:
+        assert [(name, status, count) for name, status, count in manager.statuses()] == [("fake", "connected", 3)]
+        assert await manager.call("fake", "echo", {"text": "hi"}) == "echo: hi"
+    finally:
+        await manager.aclose()
+    assert manager.statuses() == [("fake", "not started", 0)]
+
+
+async def test_manager_disable_and_restart() -> None:
+    manager = McpManager(McpConfig(servers={"fake": _server_config()}, disable=["fake"]))
+    await manager.ensure_started()
+    assert manager.connections() == []
+    with pytest.raises(McpConnectionError, match="not connected"):
+        await manager.call("fake", "echo", {})
+
+    await manager.restart(McpConfig(servers={"fake": _server_config()}, disable=[]))
+    await manager.ensure_started()
+    assert manager.statuses() == [("fake", "connected", 3)]
+    await manager.aclose()