"""stdio 传输层实现。 通过子进程 stdin/stdout 通信的标准 MCP 传输协议(本地 MCP 集成)。 JSON-RPC 2.0 请求写入 stdin,从 stdout 读取响应,stderr 记日志。 """ from __future__ import annotations import asyncio import json import logging import os import signal import time from typing import Any from app.core.config import settings from app.integrations.mcp.client import _parse_call_result from app.integrations.mcp.transports.base import TransportProtocol from app.integrations.mcp.models import ToolResult logger = logging.getLogger(__name__) class StdioTransport(TransportProtocol): """通过子进程 stdin/stdout 通信的 MCP 传输层。""" def __init__( self, server_id: str, server_name: str, command: str, args: list[str] | None = None, env: dict[str, str] | None = None, cwd: str | None = None, ) -> None: self.server_id = server_id self.server_name = server_name self._command = command self._args = args or [] self._env = env self._cwd = cwd self._process: asyncio.subprocess.Process | None = None self._request_id = 0 self._started_at: float | None = None self._lock = asyncio.Lock() async def start(self) -> None: if self._process is not None and self._process.returncode is None: return # already running resolved_env = os.environ.copy() if self._env: resolved_env.update(self._env) self._process = await asyncio.create_subprocess_exec( self._command, *self._args, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=resolved_env, cwd=self._cwd, ) self._started_at = time.time() logger.info( "stdio MCP 子进程已启动", extra={ "server_id": self.server_id, "server_name": self.server_name, "pid": self._process.pid, }, ) # 后台读取 stderr,记日志 asyncio.create_task(self._read_stderr()) async def stop(self) -> None: proc = self._process if proc is None or proc.returncode is not None: self._process = None self._started_at = None return logger.info( "停止 stdio MCP 子进程", extra={ "server_id": self.server_id, "pid": proc.pid, }, ) try: proc.send_signal(signal.SIGTERM) try: await asyncio.wait_for(proc.wait(), timeout=settings.mcp_stdio_stop_timeout_seconds) except TimeoutError: proc.send_signal(signal.SIGKILL) await proc.wait() except ProcessLookupError: pass self._process = None self._started_at = None async def list_tools(self) -> list[dict[str, Any]]: result = await self._send_request("tools/list", {}) if result is None: return [] return result.get("tools") if isinstance(result.get("tools"), list) else [] async def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: result = await self._send_request( "tools/call", {"name": name, "arguments": arguments}, ) if result is None: return ToolResult( ok=False, stdout="", stderr="stdio 子进程无响应", duration_ms=0, exit_code=1, timeout=False, truncated=False, trace_id="", ) return _parse_call_result({"result": result}) async def restart(self) -> None: await self.stop() await self.start() @property def is_running(self) -> bool: return ( self._process is not None and self._process.returncode is None and self._process.stdin is not None ) async def get_status(self) -> dict[str, Any]: running = self.is_running uptime = None if running and self._started_at is not None: uptime = time.time() - self._started_at return { "server_id": self.server_id, "name": self.server_name, "transport": "stdio", "status": "running" if running else "stopped", "pid": self._process.pid if self._process else None, "uptime_seconds": uptime, } async def _send_request(self, method: str, params: dict[str, Any]) -> dict[str, Any] | None: async with self._lock: self._request_id += 1 req_id = self._request_id proc = self._process if proc is None or proc.stdin is None or proc.stdout is None: return None request = { "jsonrpc": "2.0", "id": req_id, "method": method, "params": params, } line = (json.dumps(request) + "\n").encode("utf-8") try: proc.stdin.write(line) await asyncio.wait_for( proc.stdin.drain(), timeout=settings.mcp_stdio_call_timeout_seconds, ) response_line = await asyncio.wait_for( proc.stdout.readline(), timeout=settings.mcp_stdio_call_timeout_seconds, ) if not response_line: return None data = json.loads(response_line.decode("utf-8")) except (TimeoutError, json.JSONDecodeError, OSError) as exc: logger.warning( "stdio MCP 请求失败", extra={ "server_id": self.server_id, "method": method, "error": str(exc), }, ) return None if isinstance(data, dict) and data.get("error"): logger.warning( "stdio MCP 返回 JSON-RPC 错误", extra={ "server_id": self.server_id, "method": method, "error": data["error"], }, ) return None return data.get("result") if isinstance(data, dict) else None async def _read_stderr(self) -> None: proc = self._process if proc is None or proc.stderr is None: return try: while True: line = await proc.stderr.readline() if not line: break text = line.decode("utf-8", errors="replace").rstrip() if text: logger.info( "[stdio %s] %s", self.server_name, text, extra={ "server_id": self.server_id, "server_name": self.server_name, }, ) except Exception: # noqa: BLE001 pass