cad-agent/mcp_bridge/server.py

"""stdio ↔ TCP 桥的核心：读线程 + 单工作线程队列。

线程模型：
  - reader：逐行读 stdin。initialize / ping 在**读线程即时应答**，因此不会被
    慢工具调用阻塞；tools/list、tools/call 入队交给 worker。
  - worker：**单线程串行**转发到 CAD 基座，保证工具调用顺序（基座本身也是单线程执行）。

两个线程写 stdout 时共用一把锁，报文永不交叠。stdout 只承载协议报文，
所有诊断信息一律走 stderr。
"""

import json
import os
import queue
import sys
import threading

from dwg_review import config
from dwg_review.mcp_client import McpClient, McpError

from . import protocol


class StdioBridge:
    """把 stdio 上的 MCP 报文转发到 CAD 基座的 TCP 服务。"""

    def __init__(self, host=None, port=None, timeout=None, verbose=False):
        self.host = host or config.HOST
        self.port = port or config.PORT
        self.timeout = timeout if timeout is not None else config.CALL_TIMEOUT
        self.verbose = verbose

        self._client = McpClient(host=self.host, port=self.port, timeout=self.timeout)
        self._queue = queue.Queue()
        self._out_lock = threading.Lock()
        self._stop = object()  # 投毒丸哨兵

    # ------------------------------------------------------------------
    #  生命周期
    # ------------------------------------------------------------------
    def run(self):
        """阻塞运行，直到客户端关闭 stdin。"""
        reader = threading.Thread(target=self._read_loop, name="bridge-reader", daemon=True)
        worker = threading.Thread(target=self._work_loop, name="bridge-worker", daemon=True)
        worker.start()
        reader.start()
        try:
            reader.join()
        except KeyboardInterrupt:
            self._log("收到中断信号")
        self._queue.put(self._stop)
        worker.join()

    # ------------------------------------------------------------------
    #  reader
    # ------------------------------------------------------------------
    def _read_loop(self):
        stream = sys.stdin.buffer
        while True:
            try:
                line = stream.readline()
            except (OSError, ValueError):
                break
            if not line:
                break  # EOF：客户端已退出
            line = line.strip()
            if line:
                self._handle_line(line)

    def _handle_line(self, line):
        if len(line) > config.MAX_FRAME_BYTES:
            self._send(protocol.err(None, -32700, f"报文过大（>{config.MAX_FRAME_BYTES} 字节）"))
            return

        try:
            msg = json.loads(line.decode("utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError):
            self._send(protocol.err(None, -32700, "Parse error"))
            return

        if not isinstance(msg, dict):
            self._send(protocol.err(None, -32600, "Invalid Request：报文必须是 JSON 对象"))
            return

        msg_id = msg.get("id")
        if msg.get("jsonrpc") != "2.0":
            self._send(protocol.err(msg_id, -32600, "Invalid Request：jsonrpc 必须为 \"2.0\""))
            return

        method = msg.get("method")
        if not isinstance(method, str) or not method:
            self._send(protocol.err(msg_id, -32600, "Invalid Request：method 缺失或非字符串"))
            return

        has_id = "id" in msg
        params = msg.get("params")

        kind = protocol.classify(method)
        if kind == protocol.LOCAL:
            result = protocol.initialize_result(params) if method == "initialize" else {}
            if has_id:
                self._send(protocol.ok(msg_id, result))
        elif kind == protocol.FORWARD:
            if has_id:
                self._queue.put((msg_id, method, params))
            else:
                self._log(f"忽略无 id 的 {method} 请求")
        elif kind == protocol.NOTIFY:
            # 通知按规范不回包；若客户端误带 id，回 -32601 以免它一直等待
            if has_id:
                self._send(protocol.err(msg_id, -32601, f"Method not found: {method}"))
        else:
            if has_id:
                self._send(protocol.err(msg_id, -32601, f"Method not found: {method}"))

    # ------------------------------------------------------------------
    #  worker
    # ------------------------------------------------------------------
    def _work_loop(self):
        while True:
            item = self._queue.get()
            if item is self._stop:
                break

            msg_id, method, params = item
            try:
                # retries=0：桥只发一次，避免"发完即超时"被重试而重复执行写类工具
                resp = self._client.raw_rpc(method, params, timeout=self.timeout, retries=0)
            except McpError as exc:
                self._send(protocol.err(
                    msg_id, -32603,
                    f"无法连接 CAD 基座 {self.host}:{self.port}：{exc.message}"))
                continue
            except Exception as exc:  # noqa: BLE001
                self._send(protocol.err(msg_id, -32603, f"桥内部错误：{exc}"))
                continue

            if not isinstance(resp, dict) or not ({"result", "error"} & set(resp)):
                self._send(protocol.err(msg_id, -32603, "CAD 基座返回了非法响应"))
                continue

            # 忠实透传 result / 顶层 error，只把 id 改写回客户端的原始 id
            resp["jsonrpc"] = "2.0"
            resp["id"] = msg_id
            self._send(resp)

    # ------------------------------------------------------------------
    #  IO
    # ------------------------------------------------------------------
    def _send(self, obj):
        data = json.dumps(obj, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
        try:
            with self._out_lock:
                sys.stdout.buffer.write(data + b"\n")
                sys.stdout.buffer.flush()
        except (BrokenPipeError, ValueError):
            # 客户端已关闭管道，桥无事可做，直接退出（os._exit 避免解释器收尾时再次 flush 报错）
            self._warn("stdout 已关闭，桥退出")
            os._exit(0)

    def _log(self, message):
        """诊断信息（仅 --verbose）。"""
        if self.verbose:
            print(f"[mcp_bridge] {message}", file=sys.stderr, flush=True)

    def _warn(self, message):
        """警告，始终输出到 stderr（stderr 不属于协议通道）。"""
        print(f"[mcp_bridge] {message}", file=sys.stderr, flush=True)