sundynix-agentix/sundynix-mcp-py/src/sundynix_mcp_py/mcp_gateway.py

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


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] = {}

    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,
        }

    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:
        """解析 ToolCall → 路由到工具 → Respond ToolResult。"""
        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 _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)