hunjing/backend/app/integrations/mcp/transports/stdio.py

"""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