cad-agent/dwg_review/mcp_client.py

"""唯一的 TCP / JSON-RPC 封装。

要点：
  - 换行分帧：请求与响应均以 '\\n' 结尾（新基座按 '\\n' 分帧；旧基座读到即止，兼容）。
  - 单飞锁：CAD 侧单线程，所有调用串行，避免并发踩踏主线程派发链。
  - 错误归一化：同时识别 JSON-RPC 顶层 error 与 MCP 的 result.isError。
"""

import json
import socket
import threading
import time

from . import config


class McpError(Exception):
    """协议级错误（JSON-RPC error：-32700/-32600/-32601/-32602 等）。"""

    def __init__(self, code, message):
        super().__init__(f"[{code}] {message}")
        self.code = code
        self.message = message


class McpToolError(Exception):
    """工具执行错误（MCP result.isError == true）。"""

    def __init__(self, message, code=None):
        super().__init__(message)
        self.code = code
        self.message = message


def result_texts(resp):
    """取出 result.content 中所有 text 内容。"""
    result = (resp or {}).get("result") or {}
    return [i.get("text", "") for i in result.get("content", []) if i.get("type") == "text"]


def result_images(resp):
    """取出 result.content 中所有图片，返回 [(base64, mimeType), ...]。"""
    result = (resp or {}).get("result") or {}
    return [
        (i.get("data"), i.get("mimeType", "image/png"))
        for i in result.get("content", [])
        if i.get("type") == "image"
    ]


def result_json(resp, default=None):
    """取第一个 text 内容并尝试按 JSON 解析；解析失败则返回原始字符串。"""
    texts = result_texts(resp)
    if not texts:
        return default
    text = texts[0]
    try:
        return json.loads(text)
    except (json.JSONDecodeError, TypeError):
        return text


class McpClient:
    """与 CAD MCP 基座通信的唯一入口。"""

    def __init__(self, host=None, port=None, timeout=None):
        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._lock = threading.Lock()  # 单飞：CAD 单线程，必须串行
        self._next_id = 1
        self._tools_cache = None

    # ------------------------------------------------------------------
    #  底层收发
    # ------------------------------------------------------------------
    def _connect(self, timeout=None):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(self.timeout if timeout is None else timeout)
        sock.connect((self.host, self.port))
        return sock

    def _send_frame(self, payload, timeout=None):
        raw = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        if len(raw) > config.MAX_FRAME_BYTES:
            raise McpError(-32600, f"请求帧过大: {len(raw)} > {config.MAX_FRAME_BYTES} 字节")

        sock = self._connect(timeout)
        try:
            sock.sendall(raw + b"\n")
            data = b""
            while True:
                chunk = sock.recv(65536)
                if not chunk:
                    break
                data += chunk
                if b"\n" in data:
                    break
                if len(data) > config.MAX_FRAME_BYTES:
                    raise McpError(-32603, "响应帧过大")
        finally:
            sock.close()

        text = data.decode("utf-8", errors="replace").strip()
        if not text:
            raise McpError(-32700, "空响应")
        try:
            return json.loads(text)
        except json.JSONDecodeError as exc:
            raise McpError(-32700, f"响应解析失败: {exc}") from exc

    # ------------------------------------------------------------------
    #  高层 RPC
    # ------------------------------------------------------------------
    def _rpc(self, method, params=None, *, timeout=None, retries=None, is_write=False,
             normalize=True):
        retries = config.CALL_RETRIES if retries is None else retries
        if is_write:
            retries = 0  # 写类工具不重试，避免重复副作用

        with self._lock:
            last_exc = None
            for attempt in range(retries + 1):
                req_id = self._next_id
                self._next_id += 1
                payload = {"jsonrpc": "2.0", "id": req_id, "method": method}
                if params is not None:
                    payload["params"] = params
                try:
                    resp = self._send_frame(payload, timeout)
                    return self._normalize(resp) if normalize else resp
                except (McpError, McpToolError):
                    raise  # 协议/工具错误是确定性的，不重试
                except (OSError, socket.timeout) as exc:
                    last_exc = exc
                    if attempt < retries:
                        time.sleep(0.3 * (2 ** attempt))
                        continue
                    raise McpError(-32603, f"连接 CAD 基座失败: {exc}") from exc
            raise McpError(-32603, f"连接 CAD 基座失败: {last_exc}")

    @staticmethod
    def _normalize(resp):
        if isinstance(resp, dict) and "error" in resp:
            err = resp.get("error") or {}
            raise McpError(err.get("code", -32603), err.get("message", "unknown error"))
        result = resp.get("result") if isinstance(resp, dict) else None
        if isinstance(result, dict) and result.get("isError"):
            texts = result_texts(resp)
            raise McpToolError(texts[0] if texts else "tool error")
        return resp

    # ------------------------------------------------------------------
    #  对外 API
    # ------------------------------------------------------------------
    def call_tool(self, name, arguments=None, *, timeout=None, retries=None, is_write=False):
        """调用一个工具；成功返回归一化后的响应（含 result.content）。"""
        params = {"name": name, "arguments": arguments or {}}
        return self._rpc("tools/call", params, timeout=timeout, retries=retries, is_write=is_write)

    def raw_rpc(self, method, params=None, *, timeout=None, retries=None):
        """转发式调用：返回基座**原始响应 dict**，不做 error / isError 归一化。

        供 stdio 桥使用——桥需要把基座的 `result` 与顶层 `error` 原样透传给客户端，
        因此不能走 `_normalize`（它会把两者都变成异常）。
        """
        return self._rpc(method, params, timeout=timeout, retries=retries, normalize=False)

    def list_tools(self, *, use_cache=True):
        """返回 tools/list 的 tools 数组。"""
        if use_cache and self._tools_cache is not None:
            return self._tools_cache
        resp = self._rpc("tools/list")
        tools = (resp.get("result") or {}).get("tools", [])
        self._tools_cache = tools
        return tools

    def is_available(self, timeout=1.5):
        """探测基座是否在监听（仅建立连接，不发请求）。"""
        try:
            with self._lock:
                sock = self._connect(timeout)
                sock.close()
            return True
        except OSError:
            return False