Blizzard
53669437d5
垂直扩：单个 mcp 实例内每个工具调用分发到独立 goroutine/task 并发处理， 不再受 NATS 单订阅回调串行所限；配合队列组多副本即「单实例并发 × 副本数」横向扩。 - bus.ServeTool：回调立即起 goroutine 并返回（不阻塞 NATS 投递），信号量限并发 （MCP_TOOL_CONCURRENCY，默认 16），handler panic → 回错误结果避免调用方干等超时。 - mcp-py mcp_gateway：_on_call 改 asyncio.create_task 派发 + Semaphore 限并发（同默认 16）。 - 测试 TestConcurrentToolServe：slow 工具阻塞时 fast 工具仍返回（旧串行下会超时）。 验证：单测通过；mcp-go/mcp-py live 重启工具就绪，工具往返正常（memory/history/kb_search 均正常响应，无 panic）。全模块 build+vet+test 全绿。 注：下游共享后端（Milvus/Neo4j/PG/嵌入端点）是横向扩的最终上限，届时扩这些而非 mcp 实例。 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
184 lines
7.5 KiB
Python
184 lines
7.5 KiB
Python
"""MCP 协议网关：注册算法型工具并经 NATS 分发调用。
|
||
|
||
与 Go 侧 sundynix-shared/bus 的 ServeTool 同契约（core NATS request-reply）：
|
||
请求体 {"tool": str, "args": {...}, "task_id": str}
|
||
应答体 {"ok": bool, "content": str, "error": str}
|
||
订阅 sundynix.tools.py.>，队列组 mcp-py-workers（多副本负载均衡）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import os
|
||
|
||
import nats
|
||
|
||
from .interpreter import CodeInterpreter
|
||
from .mineru import MultimodalParser
|
||
from .sandbox import SecureSandbox
|
||
|
||
log = logging.getLogger("mcp_py")
|
||
|
||
# 与 contract.SubjectToolsPyAll / QueueToolsPy 保持一致。
|
||
SUBJECT_PY_ALL = "sundynix.tools.py.>"
|
||
QUEUE_PY = "mcp-py-workers"
|
||
|
||
# 工具元信息（与 mcp-go 注册表同形）：list_tools 上报给管理端展示 + dispatcher 动态构建自主 agent 工具集。
|
||
# agent=是否暴露给自主 agent；params=模型可填参数；inject=服务端注入参数（不暴露给模型）。
|
||
TOOL_META = {
|
||
"echo": {"cn": "回显", "desc": "原样返回入参（调试用）"},
|
||
"run_code": {
|
||
"cn": "代码执行",
|
||
"desc": "在隔离沙箱里执行 Python 代码（静态守卫 + Docker，256m/10s）。用于计算、数据处理、逻辑验证。",
|
||
"agent": True,
|
||
"params": [{"name": "code", "type": "string", "desc": "要执行的 Python 代码", "required": True}],
|
||
},
|
||
"parse_document": {"cn": "文档解析", "desc": "文件 → 纯文本（MinerU / PaddleOCR）"},
|
||
"secure_sandbox": {"cn": "安全沙箱", "desc": "更严资源档（128m/5s）的隔离执行，用于高风险代码"},
|
||
}
|
||
|
||
|
||
class McpGateway:
|
||
def __init__(self) -> None:
|
||
self.sandbox = SecureSandbox() # gVisor / KataVM + Static Code Guard
|
||
self.parser = MultimodalParser() # MinerU / PaddleOCR
|
||
self.interpreter = CodeInterpreter() # Docker 隔离沙箱
|
||
self._nc: nats.NATS | None = None
|
||
self._tools: dict[str, callable] = {}
|
||
# 单实例工具并发上限（垂直扩）；队列组多副本再叠加水平扩。
|
||
self._sem = asyncio.Semaphore(int(os.getenv("MCP_TOOL_CONCURRENCY", "16")))
|
||
|
||
async def register_tools(self) -> None:
|
||
"""注册工具名 → 处理协程。"""
|
||
self._tools = {
|
||
"echo": self._echo,
|
||
"run_code": self._run_code,
|
||
"parse_document": self._parse_document,
|
||
"secure_sandbox": self._secure_sandbox,
|
||
"list_tools": self._list_tools,
|
||
}
|
||
|
||
async def serve(self, url: str | None = None) -> None:
|
||
url = url or os.getenv("NATS_URL", "nats://localhost:4222")
|
||
# 容忍服务先于 NATS 启动：无限重连 + 等待间隔。
|
||
self._nc = await nats.connect(
|
||
url,
|
||
allow_reconnect=True,
|
||
max_reconnect_attempts=-1,
|
||
reconnect_time_wait=1,
|
||
connect_timeout=5,
|
||
)
|
||
log.info("[mcp_py] connected %s", url)
|
||
await self._nc.subscribe(SUBJECT_PY_ALL, queue=QUEUE_PY, cb=self._on_call)
|
||
log.info(
|
||
"[mcp_py] tools ready on %s (queue=%s): %s",
|
||
SUBJECT_PY_ALL, QUEUE_PY, ", ".join(self._tools),
|
||
)
|
||
await asyncio.Event().wait()
|
||
|
||
async def _on_call(self, msg) -> None:
|
||
"""立即派发到独立 task 并发处理，不阻塞订阅消息循环（nats-py 单订阅回调本是串行 await）。"""
|
||
asyncio.create_task(self._handle(msg))
|
||
|
||
async def _handle(self, msg) -> None:
|
||
"""解析 ToolCall → 路由到工具 → Respond ToolResult（限并发）。"""
|
||
async with self._sem:
|
||
try:
|
||
req = json.loads(msg.data)
|
||
except Exception as e: # noqa: BLE001
|
||
await self._reply(msg, ok=False, error=f"bad tool call: {e}")
|
||
return
|
||
tool = req.get("tool", "")
|
||
task_id = req.get("task_id", "")
|
||
args = req.get("args") or {}
|
||
log.info("[mcp_py] tool=%s task=%s args=%s", tool, task_id, args)
|
||
fn = self._tools.get(tool)
|
||
if fn is None:
|
||
await self._reply(msg, ok=False, error=f"unknown tool: {tool}")
|
||
return
|
||
try:
|
||
content = await fn(args)
|
||
await self._reply(msg, ok=True, content=content)
|
||
except Exception as e: # noqa: BLE001
|
||
await self._reply(msg, ok=False, error=str(e))
|
||
|
||
@staticmethod
|
||
async def _reply(msg, *, ok: bool, content: str = "", error: str = "") -> None:
|
||
payload = json.dumps({"ok": ok, "content": content, "error": error})
|
||
await msg.respond(payload.encode())
|
||
|
||
# ---- 工具实现（算法层目前为桩，但调用链路做真）----
|
||
|
||
async def _echo(self, args: dict) -> str:
|
||
return str(args.get("text", ""))
|
||
|
||
async def _list_tools(self, args: dict) -> str:
|
||
"""自省：上报业务工具清单（名称/中文名/作用 + agent 暴露元信息），供管理端展示 + dispatcher 动态发现。"""
|
||
tools = [
|
||
{
|
||
"name": n,
|
||
"cn": m["cn"],
|
||
"desc": m["desc"],
|
||
"agent_exposed": m.get("agent", False),
|
||
"agent_name": m.get("agent_name", ""),
|
||
"params": m.get("params", []),
|
||
"inject": m.get("inject", []),
|
||
}
|
||
for n, m in TOOL_META.items()
|
||
if n in self._tools # 仅上报真正注册的业务工具（list_tools 自身不计入）
|
||
]
|
||
return json.dumps({"service": "mcp-py", "tools": tools})
|
||
|
||
async def _run_code(self, args: dict) -> str:
|
||
"""静态守卫 → Docker 隔离执行（标准档 256m/0.5cpu/10s）。"""
|
||
code = str(args.get("code", ""))
|
||
if not code.strip():
|
||
return "（空代码）"
|
||
ok, reason = self.sandbox.static_guard(code)
|
||
if not ok:
|
||
return f"🚫 静态守卫拒绝：{reason}"
|
||
return _fmt_result(await self.interpreter.execute(code))
|
||
|
||
async def _parse_document(self, args: dict) -> str:
|
||
"""文件 → 纯文本。content_b64=文件内容(base64)，filename 决定解析器。"""
|
||
import base64
|
||
|
||
from . import parsers
|
||
|
||
filename = str(args.get("filename", ""))
|
||
content_b64 = str(args.get("content_b64", ""))
|
||
if not content_b64:
|
||
return str(args.get("text", ""))
|
||
data = base64.b64decode(content_b64)
|
||
# 解析是 CPU 密集，丢到线程池避免阻塞事件循环。
|
||
return await asyncio.to_thread(parsers.parse, filename, data)
|
||
|
||
async def _secure_sandbox(self, args: dict) -> str:
|
||
"""同 run_code 的守卫 + 隔离，更紧资源档（128m/5s），用于高风险代码。"""
|
||
code = str(args.get("code", ""))
|
||
if not code.strip():
|
||
return "（空代码）"
|
||
ok, reason = self.sandbox.static_guard(code)
|
||
if not ok:
|
||
return f"🚫 静态守卫拒绝：{reason}"
|
||
return "[secure] " + _fmt_result(await self.interpreter.execute(code, mem="128m", timeout=5))
|
||
|
||
async def close(self) -> None:
|
||
if self._nc is not None:
|
||
await self._nc.drain()
|
||
|
||
|
||
def _fmt_result(r: dict) -> str:
|
||
"""把执行结果 {ok,stdout,stderr,exit,degraded} 整理成可读文本。"""
|
||
if r.get("degraded"):
|
||
return f"⚠️ {r.get('stderr', 'Docker 不可用')}"
|
||
parts = []
|
||
if r.get("stdout", "").strip():
|
||
parts.append("stdout:\n" + r["stdout"].strip())
|
||
if r.get("stderr", "").strip():
|
||
parts.append("stderr:\n" + r["stderr"].strip())
|
||
parts.append(f"exit={r.get('exit')}")
|
||
return "\n".join(parts)
|