Files
MoviePilot/app/agent/mcp.py
T
2026-07-08 08:44:33 +08:00

601 lines
24 KiB
Python

"""Agent 外部 MCP 客户端与配置管理。"""
from __future__ import annotations
import asyncio
import hashlib
import json
import os
import re
import uuid
from dataclasses import dataclass
from typing import Any, Optional
from urllib.parse import urljoin
from app.db.systemconfig_oper import SystemConfigOper
from app.log import logger
from app.schemas.agent import (
AgentMcpServerConfig,
AgentMcpServerTestResult,
AgentMcpServerToolInfo,
)
from app.schemas.types import SystemConfigKey
from app.utils.http import AsyncRequestUtils
MCP_PROTOCOL_VERSION = "2025-11-25"
MCP_CLIENT_NAME = "MoviePilot Agent"
DEFAULT_MCP_TIMEOUT = 30
MCP_TOOL_NAME_PATTERN = re.compile(r"[^a-zA-Z0-9_]+")
@dataclass(frozen=True)
class AgentMcpToolSpec:
"""已发现的外部 MCP 工具定义。"""
server: AgentMcpServerConfig
name: str
agent_tool_name: str
description: str
input_schema: dict[str, Any]
def _normalize_identifier(value: str, fallback: str = "mcp") -> str:
"""把服务器或工具名称转换为 Agent 工具可用的标识片段。"""
normalized = MCP_TOOL_NAME_PATTERN.sub("_", str(value or "").strip())
normalized = re.sub(r"_+", "_", normalized).strip("_").lower()
if not normalized:
normalized = fallback
if normalized[0].isdigit():
normalized = f"{fallback}_{normalized}"
return normalized[:64]
def _normalize_timeout(value: Any) -> int:
"""规范化 MCP 连接和调用超时时间。"""
try:
timeout = int(value or DEFAULT_MCP_TIMEOUT)
except (TypeError, ValueError):
timeout = DEFAULT_MCP_TIMEOUT
return min(max(timeout, 1), 600)
def _normalize_string_dict(value: Any) -> dict[str, str]:
"""规范化请求头和环境变量字典,移除空键。"""
if not isinstance(value, dict):
return {}
normalized: dict[str, str] = {}
for key, item in value.items():
normalized_key = str(key or "").strip()
if not normalized_key:
continue
normalized[normalized_key] = str(item or "")
return normalized
def _normalize_input_schema(value: Any) -> dict[str, Any]:
"""规范化 MCP 工具参数 Schema,保证至少是 object schema。"""
if not isinstance(value, dict):
return {"type": "object", "properties": {}, "required": []}
schema = dict(value)
schema.setdefault("type", "object")
schema.setdefault("properties", {})
schema.setdefault("required", [])
return schema
def _build_agent_tool_name(server: AgentMcpServerConfig, tool_name: str) -> str:
"""构造注入 Agent 的外部 MCP 工具名。"""
prefix = server.tool_prefix or f"mcp_{server.name or server.id}"
normalized_prefix = _normalize_identifier(prefix, fallback="mcp")
normalized_tool_name = _normalize_identifier(tool_name, fallback="tool")
if normalized_tool_name.startswith(f"{normalized_prefix}_"):
return normalized_tool_name
return f"{normalized_prefix}_{normalized_tool_name}"[:128]
def _jsonrpc_message(method: str, params: Optional[dict[str, Any]] = None, *, request_id: Optional[str] = None) -> dict:
"""构造 JSON-RPC 2.0 消息。"""
payload = {"jsonrpc": "2.0", "method": method}
if request_id is not None:
payload["id"] = request_id
if params is not None:
payload["params"] = params
return payload
def _raise_for_jsonrpc_error(payload: Any) -> None:
"""检查 JSON-RPC 响应错误并转换为运行时异常。"""
if isinstance(payload, dict) and payload.get("error"):
error = payload["error"]
if isinstance(error, dict):
message = error.get("message") or error
else:
message = error
raise RuntimeError(f"MCP JSON-RPC 错误: {message}")
def _extract_jsonrpc_result(payload: Any, request_id: str) -> Any:
"""从 JSON-RPC 响应中提取 result 字段。"""
if not isinstance(payload, dict):
raise RuntimeError("MCP 响应不是有效 JSON 对象")
if payload.get("id") != request_id:
raise RuntimeError("MCP 响应 ID 与请求不匹配")
_raise_for_jsonrpc_error(payload)
return payload.get("result")
async def _iter_sse_events(response) -> Any:
"""按 SSE 事件格式迭代响应流。"""
event_name = "message"
data_lines: list[str] = []
async for raw_line in response.aiter_lines():
line = raw_line.rstrip("\r")
if not line:
if data_lines:
yield {"event": event_name, "data": "\n".join(data_lines)}
event_name = "message"
data_lines = []
continue
if line.startswith(":"):
continue
field, _, value = line.partition(":")
if value.startswith(" "):
value = value[1:]
if field == "event":
event_name = value or "message"
elif field == "data":
data_lines.append(value)
if data_lines:
yield {"event": event_name, "data": "\n".join(data_lines)}
def _parse_sse_text_response(text: str, request_id: str) -> Any:
"""从非流式 SSE 文本响应中提取匹配请求的 JSON-RPC 结果。"""
event_name = "message"
data_lines: list[str] = []
for raw_line in str(text or "").splitlines():
line = raw_line.rstrip("\r")
if not line:
if data_lines:
payload = _load_sse_json_payload(event_name, "\n".join(data_lines))
if isinstance(payload, dict) and payload.get("id") == request_id:
return _extract_jsonrpc_result(payload, request_id)
event_name = "message"
data_lines = []
continue
if line.startswith(":"):
continue
field, _, value = line.partition(":")
if value.startswith(" "):
value = value[1:]
if field == "event":
event_name = value or "message"
elif field == "data":
data_lines.append(value)
if data_lines:
payload = _load_sse_json_payload(event_name, "\n".join(data_lines))
if isinstance(payload, dict) and payload.get("id") == request_id:
return _extract_jsonrpc_result(payload, request_id)
raise RuntimeError("MCP SSE 响应中未找到匹配请求")
def _load_sse_json_payload(event_name: str, data: str) -> Optional[dict]:
"""解析 SSE data 中的 JSON-RPC 消息。"""
if event_name not in {"message", "messages"}:
return None
try:
payload = json.loads(data)
except (TypeError, ValueError):
return None
return payload if isinstance(payload, dict) else None
class _StdioMcpSession:
"""stdio MCP 会话,按一次操作生命周期启动外部进程。"""
def __init__(self, server: AgentMcpServerConfig) -> None:
self.server = server
self.process: Optional[asyncio.subprocess.Process] = None
self.stderr_task: Optional[asyncio.Task] = None
async def __aenter__(self) -> "_StdioMcpSession":
"""启动 stdio MCP 子进程。"""
if not self.server.command:
raise RuntimeError("stdio MCP 服务器缺少启动命令")
env = os.environ.copy()
env.update(self.server.env or {})
self.process = await asyncio.create_subprocess_exec(
self.server.command,
*(self.server.args or []),
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env,
)
self.stderr_task = asyncio.create_task(self._drain_stderr())
return self
async def __aexit__(self, exc_type, exc, tb) -> None:
"""结束 stdio MCP 子进程。"""
if self.stderr_task:
self.stderr_task.cancel()
if not self.process:
return
if self.process.returncode is None:
self.process.terminate()
try:
await asyncio.wait_for(self.process.wait(), timeout=2)
except asyncio.TimeoutError:
self.process.kill()
await self.process.wait()
async def _drain_stderr(self) -> None:
"""持续读取子进程 stderr,避免缓冲区阻塞。"""
if not self.process or not self.process.stderr:
return
try:
while True:
line = await self.process.stderr.readline()
if not line:
break
logger.debug(f"MCP stdio[{self.server.name}] stderr: {line.decode(errors='replace').strip()}")
except asyncio.CancelledError:
return
async def notify(self, method: str, params: Optional[dict[str, Any]] = None) -> None:
"""发送不需要响应的 JSON-RPC 通知。"""
await self._write_json(_jsonrpc_message(method, params))
async def request(self, method: str, params: Optional[dict[str, Any]] = None) -> Any:
"""发送 JSON-RPC 请求并等待响应。"""
request_id = uuid.uuid4().hex
await self._write_json(_jsonrpc_message(method, params, request_id=request_id))
while True:
payload = await self._read_json()
if payload.get("id") == request_id:
return _extract_jsonrpc_result(payload, request_id)
async def _write_json(self, payload: dict) -> None:
"""写入一行 JSON-RPC 消息。"""
if not self.process or not self.process.stdin:
raise RuntimeError("stdio MCP 进程未启动")
data = json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + "\n"
self.process.stdin.write(data.encode("utf-8"))
await self.process.stdin.drain()
async def _read_json(self) -> dict:
"""从 stdout 读取一行 JSON-RPC 消息。"""
if not self.process or not self.process.stdout:
raise RuntimeError("stdio MCP 进程未启动")
timeout = _normalize_timeout(self.server.timeout)
while True:
line = await asyncio.wait_for(self.process.stdout.readline(), timeout=timeout)
if not line:
raise RuntimeError("stdio MCP 进程已退出")
try:
payload = json.loads(line.decode("utf-8"))
except ValueError:
logger.debug(f"忽略非 JSON MCP stdout 行: {line!r}")
continue
if isinstance(payload, dict):
return payload
class _HttpMcpSession:
"""Streamable HTTP MCP 会话。"""
def __init__(self, server: AgentMcpServerConfig) -> None:
self.server = server
self.session_id: Optional[str] = None
async def __aenter__(self) -> "_HttpMcpSession":
"""进入 HTTP MCP 会话。"""
return self
async def __aexit__(self, exc_type, exc, tb) -> None:
"""退出 HTTP MCP 会话。"""
return None
async def notify(self, method: str, params: Optional[dict[str, Any]] = None) -> None:
"""发送不需要响应的 JSON-RPC 通知。"""
await self._post(_jsonrpc_message(method, params), expect_response=False)
async def request(self, method: str, params: Optional[dict[str, Any]] = None) -> Any:
"""发送 JSON-RPC 请求并等待响应。"""
request_id = uuid.uuid4().hex
return await self._post(
_jsonrpc_message(method, params, request_id=request_id),
expect_response=True,
request_id=request_id,
)
async def _post(
self,
payload: dict,
*,
expect_response: bool,
request_id: Optional[str] = None,
) -> Any:
"""向 Streamable HTTP MCP 服务发送一条 JSON-RPC 消息。"""
if not self.server.url:
raise RuntimeError("HTTP MCP 服务器缺少 URL")
headers = {
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
**(self.server.headers or {}),
}
if self.session_id:
headers["Mcp-Session-Id"] = self.session_id
response = await AsyncRequestUtils(
headers=headers,
timeout=_normalize_timeout(self.server.timeout),
content_type="application/json",
accept_type="application/json, text/event-stream",
http2=False,
).post_res(self.server.url, json=payload, raise_exception=True)
try:
if not response:
raise RuntimeError("HTTP MCP 请求无响应")
response.raise_for_status()
session_id = response.headers.get("Mcp-Session-Id")
if session_id:
self.session_id = session_id
if not expect_response:
return None
content_type = response.headers.get("content-type", "").lower()
if "text/event-stream" in content_type:
return _parse_sse_text_response(response.text, request_id or "")
data = response.json()
return _extract_jsonrpc_result(data, request_id or "")
finally:
if response is not None:
await response.aclose()
class _SseMcpSession:
"""旧版 HTTP+SSE MCP 会话。"""
def __init__(self, server: AgentMcpServerConfig) -> None:
self.server = server
self.response = None
self.endpoint: Optional[str] = None
self._stream_manager = None
self._event_iterator = None
async def __aenter__(self) -> "_SseMcpSession":
"""打开 SSE 流并读取服务端回传的 POST endpoint。"""
if not self.server.url:
raise RuntimeError("SSE MCP 服务器缺少 URL")
self._stream_manager = AsyncRequestUtils(
headers={"Accept": "text/event-stream", **(self.server.headers or {})},
timeout=_normalize_timeout(self.server.timeout),
accept_type="text/event-stream",
http2=False,
).get_stream(self.server.url, raise_exception=True)
self.response = await self._stream_manager.__aenter__()
if not self.response:
raise RuntimeError("SSE MCP 连接无响应")
self.response.raise_for_status()
self._event_iterator = _iter_sse_events(self.response).__aiter__()
self.endpoint = await self._read_endpoint()
return self
async def __aexit__(self, exc_type, exc, tb) -> None:
"""关闭 SSE 流。"""
if self._stream_manager:
await self._stream_manager.__aexit__(exc_type, exc, tb)
async def notify(self, method: str, params: Optional[dict[str, Any]] = None) -> None:
"""发送不需要响应的 JSON-RPC 通知。"""
await self._post(_jsonrpc_message(method, params))
async def request(self, method: str, params: Optional[dict[str, Any]] = None) -> Any:
"""发送 JSON-RPC 请求并等待 SSE 流上的响应。"""
request_id = uuid.uuid4().hex
await self._post(_jsonrpc_message(method, params, request_id=request_id))
timeout = _normalize_timeout(self.server.timeout)
while True:
event = await asyncio.wait_for(self._event_iterator.__anext__(), timeout=timeout)
payload = _load_sse_json_payload(event.get("event", ""), event.get("data", ""))
if isinstance(payload, dict) and payload.get("id") == request_id:
return _extract_jsonrpc_result(payload, request_id)
async def _read_endpoint(self) -> str:
"""读取 SSE endpoint 事件中的 POST 地址。"""
timeout = _normalize_timeout(self.server.timeout)
while True:
event = await asyncio.wait_for(self._event_iterator.__anext__(), timeout=timeout)
if event.get("event") != "endpoint":
continue
endpoint = str(event.get("data") or "").strip()
if not endpoint:
continue
return urljoin(self.server.url, endpoint)
async def _post(self, payload: dict) -> None:
"""向 SSE 握手返回的 endpoint 发送 JSON-RPC 消息。"""
if not self.endpoint:
raise RuntimeError("SSE MCP endpoint 未初始化")
response = await AsyncRequestUtils(
headers={
"Accept": "application/json",
"Content-Type": "application/json",
**(self.server.headers or {}),
},
timeout=_normalize_timeout(self.server.timeout),
content_type="application/json",
accept_type="application/json",
http2=False,
).post_res(self.endpoint, json=payload, raise_exception=True)
try:
if not response:
raise RuntimeError("SSE MCP POST 请求无响应")
response.raise_for_status()
finally:
if response is not None:
await response.aclose()
async def _open_mcp_session(server: AgentMcpServerConfig):
"""根据配置创建对应的 MCP 传输会话。"""
transport = "http" if server.transport == "streamable_http" else server.transport
if transport == "stdio":
return _StdioMcpSession(server)
if transport == "sse":
return _SseMcpSession(server)
if transport == "http":
return _HttpMcpSession(server)
raise RuntimeError(f"不支持的 MCP 传输协议: {server.transport}")
class AgentMcpManager:
"""管理 Agent 外部 MCP 服务器配置、工具发现和工具调用。"""
def get_servers(self) -> list[AgentMcpServerConfig]:
"""读取已保存的外部 MCP 服务器配置。"""
raw_servers = SystemConfigOper().get(SystemConfigKey.AIAgentMcpServers) or []
if not isinstance(raw_servers, list):
return []
servers: list[AgentMcpServerConfig] = []
for raw_server in raw_servers:
try:
servers.append(self.normalize_server(raw_server))
except Exception as err:
logger.warning(f"忽略无效的 Agent MCP 配置: {err}")
return servers
async def save_servers(self, servers: list[AgentMcpServerConfig]) -> bool:
"""保存外部 MCP 服务器配置。"""
normalized_servers = [self.normalize_server(server).model_dump() for server in servers]
return await SystemConfigOper().async_set(
SystemConfigKey.AIAgentMcpServers,
normalized_servers or None,
)
def normalize_server(self, value: Any) -> AgentMcpServerConfig:
"""规范化单个 MCP 服务器配置。"""
if isinstance(value, AgentMcpServerConfig):
raw_server = value.model_dump()
elif isinstance(value, dict):
raw_server = dict(value)
else:
raise ValueError("MCP 服务器配置必须是对象")
raw_server["id"] = str(raw_server.get("id") or uuid.uuid4().hex[:12]).strip()
raw_server["name"] = str(raw_server.get("name") or raw_server["id"]).strip()
raw_server["transport"] = str(raw_server.get("transport") or "stdio").strip()
raw_server["description"] = str(raw_server.get("description") or "").strip() or None
raw_server["command"] = str(raw_server.get("command") or "").strip() or None
raw_server["args"] = [str(item) for item in raw_server.get("args") or []]
raw_server["env"] = _normalize_string_dict(raw_server.get("env"))
raw_server["url"] = str(raw_server.get("url") or "").strip() or None
raw_server["headers"] = _normalize_string_dict(raw_server.get("headers"))
raw_server["timeout"] = _normalize_timeout(raw_server.get("timeout"))
raw_server["tool_prefix"] = str(raw_server.get("tool_prefix") or "").strip() or None
raw_server["require_admin"] = bool(raw_server.get("require_admin", True))
return AgentMcpServerConfig.model_validate(raw_server)
def config_signature(self) -> str:
"""生成外部 MCP 配置签名,用于 Agent 图缓存失效。"""
payload = [server.model_dump() for server in self.get_servers()]
raw_text = json.dumps(payload, ensure_ascii=False, sort_keys=True, default=str)
return hashlib.sha256(raw_text.encode("utf-8")).hexdigest()
async def initialize_session(self, session) -> None:
"""完成 MCP initialize 和 initialized 通知流程。"""
await session.request(
"initialize",
{
"protocolVersion": MCP_PROTOCOL_VERSION,
"capabilities": {},
"clientInfo": {
"name": MCP_CLIENT_NAME,
"version": "1.0.0",
},
},
)
await session.notify("notifications/initialized")
async def list_server_tools(self, server: AgentMcpServerConfig) -> list[AgentMcpToolSpec]:
"""连接单个 MCP 服务器并读取工具列表。"""
normalized_server = self.normalize_server(server)
session_manager = await _open_mcp_session(normalized_server)
async with session_manager as session:
await self.initialize_session(session)
result = await session.request("tools/list")
tools_payload = result.get("tools", []) if isinstance(result, dict) else []
tool_specs: list[AgentMcpToolSpec] = []
for item in tools_payload:
if not isinstance(item, dict) or not item.get("name"):
continue
tool_name = str(item["name"])
tool_specs.append(
AgentMcpToolSpec(
server=normalized_server,
name=tool_name,
agent_tool_name=_build_agent_tool_name(normalized_server, tool_name),
description=str(item.get("description") or ""),
input_schema=_normalize_input_schema(item.get("inputSchema")),
)
)
return tool_specs
async def list_enabled_tool_specs(self) -> list[AgentMcpToolSpec]:
"""读取所有启用 MCP 服务器暴露的工具定义。"""
tool_specs: list[AgentMcpToolSpec] = []
seen_names: set[str] = set()
for server in self.get_servers():
if not server.enabled:
continue
try:
for spec in await self.list_server_tools(server):
if spec.agent_tool_name in seen_names:
logger.warning(f"跳过重复的 MCP Agent 工具名: {spec.agent_tool_name}")
continue
tool_specs.append(spec)
seen_names.add(spec.agent_tool_name)
except Exception as err:
logger.warning(f"读取 MCP 服务器 {server.name} 工具失败: {err}")
return tool_specs
async def call_server_tool(
self,
server: AgentMcpServerConfig,
tool_name: str,
arguments: Optional[dict[str, Any]] = None,
) -> Any:
"""调用单个 MCP 服务器上的指定工具。"""
normalized_server = self.normalize_server(server)
session_manager = await _open_mcp_session(normalized_server)
async with session_manager as session:
await self.initialize_session(session)
return await session.request(
"tools/call",
{
"name": tool_name,
"arguments": arguments or {},
},
)
async def test_server(self, server: AgentMcpServerConfig) -> AgentMcpServerTestResult:
"""测试 MCP 服务器连接并返回工具列表。"""
tool_specs = await self.list_server_tools(server)
tools = [
AgentMcpServerToolInfo(
name=spec.name,
agent_tool_name=spec.agent_tool_name,
description=spec.description,
input_schema=spec.input_schema,
)
for spec in tool_specs
]
return AgentMcpServerTestResult(
success=True,
message=f"连接成功,发现 {len(tools)} 个工具",
tools=tools,
tool_count=len(tools),
)
agent_mcp_manager = AgentMcpManager()