MindOS_CLI/mindcli/pipelines/tunnel_session.py

"""
Mind CLI — WebSocket Tunnel 会话管线（无状态）。

连接到 Cloud 端 mindcli_bridge，接收工具调用指令并在本地执行。
采用 Browser-Donated JWT 认证：浏览器授权 CLI，CLI 不需独立认证。

无状态：不持有进程级单例、不反向 import 调用方模块（health.py）。
状态变更通过 on_status 回调通知调用方，工具调用通过 on_dispatch 回调派发。
生命周期由调用方（health.py）管理。
"""

import asyncio
import json
import logging
from typing import Callable

logger = logging.getLogger("mindcli.pipelines.tunnel_session")

# 连接状态
DISCONNECTED = "disconnected"
CONNECTING = "connecting"
CONNECTED = "connected"


async def connect(
    url: str,
    jwt: str,
    on_dispatch: Callable[[dict], None] | None = None,
    on_status: Callable[[str, int], None] | None = None,
) -> "TunnelHandle":
    """
    建立 Tunnel 连接，返回句柄。

    Args:
        url: Cloud Tunnel WebSocket URL
        jwt: MindPass JWT（浏览器提供）
        on_dispatch: 工具调用派发回调（可选，默认走内部 ToolProxy）
        on_status: 状态变更回调 (status_str, tool_count)

    Returns:
        TunnelHandle（已在后台启动连接循环）
    """
    handle = TunnelHandle(url, jwt, on_dispatch, on_status)
    handle._start_connect_loop()
    return handle


class TunnelHandle:
    """
    CLI → Cloud WebSocket 隧道句柄。

    非单例——生命周期由调用方管理。

    生命周期：
    1. 调用方调 connect() → 后台启动 _connect_loop
    2. 握手 + 能力协商 → Cloud 下发 approved_tools → 创建 ToolProxy
    3. 消息循环：接收 tool_call → ToolProxy 执行 → 返回结果
    4. 心跳维持 30s / 断线指数退避重连
    """

    def __init__(
        self,
        url: str,
        jwt: str,
        on_dispatch: Callable[[dict], None] | None = None,
        on_status: Callable[[str, int], None] | None = None,
    ):
        self._tunnel_url = url
        self._jwt = jwt
        self._on_dispatch = on_dispatch
        self._on_status = on_status

        self._status = DISCONNECTED
        self._ws = None
        self._user_id: str | None = None
        self._tool_proxy = None  # ToolProxy 实例，握手成功后创建

        # 后台任务
        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

    @property
    def status(self) -> str:
        return self._status

    @property
    def user_id(self) -> str | None:
        return self._user_id

    def _set_status(self, status: str, tools: int = 0) -> None:
        self._status = status
        if self._on_status:
            self._on_status(status, tools)

    def _start_connect_loop(self) -> None:
        """在当前 event loop 中后台启动连接循环。"""
        self._reconnect_task = asyncio.create_task(self._connect_loop())

    async def activate(self, jwt: str, tunnel_url: str) -> dict:
        """
        更新 JWT + URL 并重新连接（由调用方在浏览器重新授权时调用）。

        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] 已断开")

    def update_approved(self, tools: list[str]) -> None:
        """热更新工具白名单（Cloud 下发时调用）。"""
        if self._tool_proxy:
            self._tool_proxy.update_approved(tools)

    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)
        approved: list[str] = []
        if approval.get("type") == "approved_tools":
            approved = approval.get("tools", [])
            # 初始化 ToolProxy（从 pipelines.tool_proxy 导入，无状态）
            from mindcli.pipelines.tool_proxy import ToolProxy
            self._tool_proxy = ToolProxy(approved_tools=approved)
            logger.info("[Tunnel] 审批通过工具: %s", approved)

        # ★ 通过回调通知调用方状态（不反向 import health.py）
        self._set_status(CONNECTED, len(approved))

    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._tool_proxy:
                        self._tool_proxy.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 self._on_dispatch:
            # 调用方自定义派发
            result = self._on_dispatch(msg)
        elif not self._tool_proxy:
            result = {"error": "ToolProxy not initialized"}
        else:
            result = await self._tool_proxy.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))