初始化提交
This commit is contained in:
chenxin committed
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1,231 @@
|
||||
"""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.modules.tool_runtime.definition 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
|
||||
Reference in new issue
Block a user