""" Mind CLI — WebSocket Tunnel 客户端。 连接到 Cloud 端 mindcli_bridge，接收工具调用指令并在本地执行。 采用 Browser-Donated JWT 认证：浏览器授权 CLI，CLI 不需独立认证。 """ import asyncio import json import logging import os import time from typing import Any logger = logging.getLogger("mindcli.tunnel") # 连接状态 DISCONNECTED = "disconnected" CONNECTING = "connecting" CONNECTED = "connected" class TunnelClient: """ CLI → Cloud WebSocket 隧道。 生命周期： 1. 浏览器 POST /tunnel/activate → 提供 JWT + tunnelUrl 2. TunnelClient.connect() → WebSocket 握手 + 能力协商 3. 消息循环：接收 tool_call → ManagedMCP 执行 → 返回结果 4. 心跳维持 30s / 断线指数退避重连 """ def __init__(self): self._status = DISCONNECTED self._ws = None self._jwt: str | None = None self._tunnel_url: str | None = None self._user_id: str | None = None self._managed_mcp = None self._reconnect_task: asyncio.Task | None = None self._heartbeat_task: asyncio.Task | None = None self._message_task: asyncio.Task | None = None # 重连参数 self._reconnect_delay = 1.0 self._max_reconnect_delay = 30.0 self._reconnect_attempts = 0 # 回调：通知 health.py 更新状态 self._on_status_change = None @property def status(self) -> str: return self._status @property def user_id(self) -> str | None: return self._user_id def set_status_callback(self, callback) -> None: """设置状态变更回调（由 health.py 注册）。""" self._on_status_change = callback def _set_status(self, status: str) -> None: self._status = status if self._on_status_change: self._on_status_change(status) async def activate(self, jwt: str, tunnel_url: str) -> dict: """ 浏览器授权激活 Tunnel。 Args: jwt: MindPass JWT（浏览器提供） tunnel_url: Cloud Tunnel WebSocket URL Returns: {"ok": True, "status": "connecting"} """ self._jwt = jwt self._tunnel_url = tunnel_url # 取消旧连接 await self.disconnect() # 后台启动连接 self._reconnect_task = asyncio.create_task(self._connect_loop()) return {"ok": True, "status": "connecting"} async def disconnect(self) -> None: """断开 Tunnel 连接。""" # 取消所有后台任务 for task in [self._reconnect_task, self._heartbeat_task, self._message_task]: if task and not task.done(): task.cancel() try: await task except asyncio.CancelledError: pass if self._ws: try: await self._ws.close() except Exception: pass self._ws = None self._set_status(DISCONNECTED) self._reconnect_attempts = 0 self._reconnect_delay = 1.0 logger.info("[Tunnel] 已断开") async def _connect_loop(self) -> None: """连接循环：握手 → 消息循环 → 断线重连。""" while True: try: self._set_status(CONNECTING) await self._connect() # 连接成功，重置重连参数 self._reconnect_delay = 1.0 self._reconnect_attempts = 0 # 进入消息循环（阻塞直到断线） await self._run() except asyncio.CancelledError: return except Exception as e: logger.warning("[Tunnel] 连接异常: %s", e) # 断线，指数退避重连 self._set_status(DISCONNECTED) self._reconnect_attempts += 1 delay = min(self._reconnect_delay, self._max_reconnect_delay) logger.info("[Tunnel] %ds 后重连（第 %d 次）...", delay, self._reconnect_attempts) await asyncio.sleep(delay) self._reconnect_delay = min(self._reconnect_delay * 2, self._max_reconnect_delay) async def _connect(self) -> None: """WebSocket 握手 + 能力协商。""" try: import websockets except ImportError: logger.error("[Tunnel] websockets 未安装。运行: pip install websockets") raise headers = {"Authorization": f"Bearer {self._jwt}"} self._ws = await websockets.connect( self._tunnel_url, additional_headers=headers, ping_interval=30, ping_timeout=10, close_timeout=5, ) # 等待 Cloud 确认 raw = await asyncio.wait_for(self._ws.recv(), timeout=10) msg = json.loads(raw) if msg.get("type") != "connected": raise ConnectionError(f"握手失败: {msg}") self._user_id = msg.get("userId") logger.info("[Tunnel] 已连接，userId=%s", self._user_id) # 发送能力报告 from mindcli.capability import scan_capabilities cap = scan_capabilities() await self._ws.send(json.dumps({ "type": "capability_report", **cap, })) # 接收审批结果 raw = await asyncio.wait_for(self._ws.recv(), timeout=10) approval = json.loads(raw) if approval.get("type") == "approved_tools": approved = approval.get("tools", []) # 初始化 ManagedMCP from mindcli.managed_mcp import ManagedMCP self._managed_mcp = ManagedMCP(approved_tools=approved) logger.info("[Tunnel] 审批通过工具: %s", approved) # 更新 health 状态 from mindcli.health import set_tunnel_status set_tunnel_status("connected", len(approved) if approval.get("type") == "approved_tools" else 0) self._set_status(CONNECTED) async def _run(self) -> None: """主消息循环：接收 Cloud 指令 → 本地执行 → 返回结果。""" try: async for raw in self._ws: msg = json.loads(raw) if msg.get("jsonrpc") == "2.0" and msg.get("method") == "tool_call": await self._handle_tool_call(msg) elif msg.get("type") == "approved_tools": # 热更新白名单 if self._managed_mcp: self._managed_mcp.update_approved(msg.get("tools", [])) elif msg.get("type") == "ping": await self._ws.send(json.dumps({"type": "pong"})) else: logger.debug("[Tunnel] 未知消息: %s", msg.get("type")) except Exception as e: logger.warning("[Tunnel] 消息循环异常: %s", e) raise async def _handle_tool_call(self, msg: dict) -> None: """处理工具调用请求。""" call_id = msg.get("id", "unknown") params = msg.get("params", {}) tool_name = params.get("tool", "") tool_args = params.get("args", {}) logger.info("[Tunnel] 工具调用: %s (id=%s)", tool_name, call_id) if not self._managed_mcp: result = {"error": "ManagedMCP not initialized"} else: result = await self._managed_mcp.execute(tool_name, tool_args) # 截断过大的输出（防止 WS 阻塞） output = result.get("output", "") if isinstance(output, str) and len(output) > 50000: result["output"] = output[:50000] + f"\n\n... [截断：原始 {len(output)} 字符]" result["truncated"] = True response = { "jsonrpc": "2.0", "id": call_id, "result": result, } await self._ws.send(json.dumps(response)) # ── 全局单例 ────────────────────────────────────────── _tunnel_client = TunnelClient() def get_tunnel_client() -> TunnelClient: """获取全局 TunnelClient 实例。""" return _tunnel_client