@agentunion/kite 1.3.2 → 1.5.0

package/extensions/services/evol/relay.py

@@ -0,0 +1,682 @@
+"""
+Kernel WebSocket 中转服务
+
+功能：
+- 接受前端 WebSocket 连接
+- 处理配对流程（通过 WebSocket）
+- 为每个前端创建到 Kernel 的连接
+- 双向转发 JSON-RPC 消息
+- 管理前端模块的生命周期（注册、重连、优雅退出）
+- RPC 权限控制（白名单）
+
+零共享代码依赖 - 此文件可以独立拷贝到其他模块使用。
+"""
+
+import asyncio
+import json
+import secrets
+import time
+import uuid
+from datetime import datetime, timezone
+from typing import Optional
+
+import websockets
+from fastapi import WebSocket, WebSocketDisconnect
+
+
+class SessionInfo:
+    """前端会话信息"""
+
+    def __init__(
+        self,
+        session_token: str,
+        module_id: str,
+        kernel_ws,
+        client_ws: WebSocket,
+        frontend_token: str,
+        role: str,
+        kernel_token: str,
+    ):
+        self.session_token = session_token
+        self.module_id = module_id
+        self.kernel_ws = kernel_ws
+        self.client_ws: Optional[WebSocket] = client_ws
+        self.frontend_token = frontend_token
+        self.role = role
+        self.kernel_token = kernel_token
+        self.created_at = time.time()
+        self.last_active = time.time()
+        self.status = "active"  # active | waiting_reconnect
+
+
+class KernelRelay:
+    """Kernel WebSocket 中转服务"""
+
+    def __init__(
+        self,
+        kernel_host: str,
+        kernel_port: int,
+        kernel_token: str,
+        base_module_id: str,
+        reconnect_timeout: int,
+        permissions: dict,
+        pairing_manager,
+        evol_server,  # 新增：web server 实例，用于发送事件
+    ):
+        """
+        初始化中转服务。
+
+        Args:
+            kernel_host: Kernel 主机地址
+            kernel_port: Kernel 端口
+            kernel_token: Kernel 认证 token
+            base_module_id: 基础模块 ID（实际 module_id 会加上 session 后缀）
+            reconnect_timeout: 重连超时时间（秒）
+            permissions: 权限配置（角色 → RPC 白名单）
+            pairing_manager: 配对管理器实例
+            evol_server: web server 实例
+        """
+        self.kernel_host = kernel_host
+        self.kernel_port = kernel_port
+        self.kernel_token = kernel_token
+        self.base_module_id = base_module_id
+        self.reconnect_timeout = reconnect_timeout
+        self.permissions = permissions
+        self.pairing_manager = pairing_manager
+        self.evol_server = evol_server
+
+        # session_token → SessionInfo
+        self.sessions: dict[str, SessionInfo] = {}
+
+        # 待重连的 session（超时清理）
+        self.reconnect_timers: dict[str, asyncio.Task] = {}
+
+    async def handle_client(self, client_ws: WebSocket):
+        """
+        处理客户端连接。
+
+        流程：
+        1. 接受 WebSocket 连接
+        2. 接收认证/配对/请求配对码消息
+        3. 验证身份
+        4. 创建或恢复 Kernel 连接
+        5. 双向转发消息
+        """
+        await client_ws.accept()
+
+        try:
+            # 接收认证/配对/请求配对码消息
+            raw = await client_ws.receive_text()
+            msg = json.loads(raw)
+            msg_type = msg.get("type")
+
+            if msg_type == "request_code":
+                # 请求配对码
+                await self._handle_request_code(client_ws, msg)
+            elif msg_type == "pair":
+                # 配对流程
+                await self._handle_pair(client_ws, msg)
+            elif msg_type == "auth":
+                # 认证流程（已有 token）
+                await self._handle_auth(client_ws, msg)
+            else:
+                await client_ws.send_json({
+                    "type": "error",
+                    "message": "Invalid message type, expected 'request_code', 'pair' or 'auth'"
+                })
+                await client_ws.close(code=4000, reason="Invalid message type")
+
+        except WebSocketDisconnect:
+            pass
+        except Exception as e:
+            print(f"[relay] Client connection error: {e}")
+            try:
+                await client_ws.close(code=1011, reason="Internal error")
+            except Exception:
+                pass
+
+    async def _handle_request_code(self, client_ws: WebSocket, msg: dict):
+        """处理请求配对码"""
+        # 生成配对码
+        code = self.pairing_manager.generate_pairing_code(role="admin")
+
+        # 发送事件给 Launcher（通过 Kernel）
+        if self.evol_server and self.evol_server._ws:
+            try:
+                await self.evol_server._publish_event({
+                    "event": "pairing.status",
+                    "data": {
+                        "step": "code_generated",
+                        "success": True,
+                        "code": code,
+                        "expires_in": 300,
+                        "module_id": "evol"
+                    }
+                })
+            except Exception as e:
+                print(f"[relay] Failed to publish pairing event: {e}")
+
+        # 返回配对码给前端
+        await client_ws.send_json({
+            "type": "code_generated",
+            "code": code
+        })
+
+        # 关闭连接（前端会重新连接进行配对）
+        await client_ws.close()
+
+    async def _handle_pair(self, client_ws: WebSocket, msg: dict):
+        """处理配对请求"""
+        code = msg.get("code")
+        if not code:
+            await client_ws.send_json({
+                "type": "error",
+                "message": "Missing pairing code"
+            })
+            await client_ws.close(code=4000, reason="Missing pairing code")
+            return
+
+        # 验证配对码
+        result = self.pairing_manager.pair(code)
+        if not result:
+            # 发送配对失败事件
+            if self.evol_server and self.evol_server._ws:
+                try:
+                    await self.evol_server._publish_event({
+                        "event": "pairing.status",
+                        "data": {
+                            "step": "completed",
+                            "success": False,
+                            "reason": "Invalid pairing code"
+                        }
+                    })
+                except Exception as e:
+                    print(f"[relay] Failed to publish pairing failed event: {e}")
+
+            await client_ws.send_json({
+                "type": "error",
+                "message": "Invalid pairing code"
+            })
+            await client_ws.close(code=4001, reason="Invalid pairing code")
+            return
+
+        # 生成 session_token
+        session_token = "sess_" + secrets.token_urlsafe(6)[:6]
+
+        # 创建新连接
+        await self._create_new_connection(
+            client_ws,
+            session_token,
+            result["token"],
+            result["role"],
+            is_pairing=True
+        )
+
+    async def _handle_auth(self, client_ws: WebSocket, msg: dict):
+        """处理认证请求（已有 token）"""
+        frontend_token = msg.get("token")
+        session_token = msg.get("session_token")
+
+        if not frontend_token or not session_token:
+            await client_ws.send_json({
+                "type": "error",
+                "message": "Missing token or session_token"
+            })
+            await client_ws.close(code=4000, reason="Missing credentials")
+            return
+
+        # 验证 frontend_token
+        token_info = self.pairing_manager.verify_token(frontend_token)
+        if not token_info:
+            await client_ws.send_json({
+                "type": "error",
+                "message": "Invalid or expired token"
+            })
+            await client_ws.close(code=4001, reason="Invalid token")
+            return
+
+        # 检查是否是重连
+        if session_token in self.sessions:
+            await self._handle_reconnect(client_ws, session_token, token_info)
+        else:
+            await self._create_new_connection(
+                client_ws,
+                session_token,
+                frontend_token,
+                token_info["role"]
+            )
+
+    async def _create_new_connection(
+        self,
+        client_ws: WebSocket,
+        session_token: str,
+        frontend_token: str,
+        role: str,
+        is_pairing: bool = False
+    ):
+        """创建新的 Kernel 连接"""
+        # 生成 module_id
+        suffix = session_token.replace("sess_", "")
+        module_id = f"{self.base_module_id}-{suffix}"
+
+        try:
+            # 向 Launcher 申请 kernel_token
+            kernel_token = await self._request_kernel_token(module_id)
+
+            # 连接 Kernel
+            kernel_ws = await self._connect_kernel(module_id, kernel_token)
+
+            # 创建 session
+            session = SessionInfo(
+                session_token=session_token,
+                module_id=module_id,
+                kernel_ws=kernel_ws,
+                client_ws=client_ws,
+                frontend_token=frontend_token,
+                role=role,
+                kernel_token=kernel_token,
+            )
+            self.sessions[session_token] = session
+
+            # 返回认证成功
+            await client_ws.send_json({
+                "type": "paired" if is_pairing else "authenticated",
+                "token": frontend_token,
+                "session_token": session_token,
+                "module_id": module_id,
+                "role": role
+            })
+
+            print(f"[relay] New connection: {module_id} (role: {role})")
+
+            # 如果是配对，发送配对成功事件给 Launcher
+            if is_pairing and self.evol_server and self.evol_server._ws:
+                try:
+                    await self.evol_server._publish_event({
+                        "event": "pairing.status",
+                        "data": {
+                            "step": "completed",
+                            "success": True,
+                            "module_id": module_id,
+                            "role": role
+                        }
+                    })
+                except Exception as e:
+                    print(f"[relay] Failed to publish pairing success event: {e}")
+
+            # 开始双向转发
+            await self._relay_messages(session)
+
+        except Exception as e:
+            error_msg = str(e)
+            print(f"[relay] Failed to create connection: {error_msg}")
+
+            # 判断错误类型
+            is_fatal = self._is_fatal_error(error_msg)
+            error_code = 1011  # Internal Error (default)
+
+            if is_fatal:
+                # 永久性错误（权限、配置错误）
+                error_code = 1008  # Policy Violation
+                print(f"[relay] Fatal error detected, client should not retry")
+
+            await client_ws.send_json({
+                "type": "error",
+                "message": f"Failed to connect to Kernel: {error_msg}",
+                "fatal": is_fatal,  # 告知前端是否应该重试
+                "code": error_code
+            })
+            await client_ws.close(code=error_code, reason="Kernel connection failed")
+
+    def _is_fatal_error(self, error_msg: str) -> bool:
+        """判断是否为永久性错误（不应重试）"""
+        fatal_keywords = [
+            "Permission denied",
+            "not in relay_modules whitelist",
+            "module_id must start with",
+            "Invalid module_id",
+            "Token limit reached",
+        ]
+        return any(keyword in error_msg for keyword in fatal_keywords)
+
+    async def _handle_reconnect(
+        self,
+        client_ws: WebSocket,
+        session_token: str,
+        token_info: dict
+    ):
+        """处理重连"""
+        session = self.sessions[session_token]
+
+        # 取消超时计时器
+        if session_token in self.reconnect_timers:
+            self.reconnect_timers[session_token].cancel()
+            del self.reconnect_timers[session_token]
+
+        # 更新 client_ws
+        session.client_ws = client_ws
+        session.status = "active"
+        session.last_active = time.time()
+
+        # 返回重连成功
+        await client_ws.send_json({
+            "type": "reconnected",
+            "module_id": session.module_id,
+            "role": session.role
+        })
+
+        print(f"[relay] Reconnected: {session.module_id}")
+
+        # 继续双向转发
+        await self._relay_messages(session)
+
+    async def _request_kernel_token(self, module_id: str) -> str:
+        """向 Launcher 申请 kernel_token"""
+        if not self.evol_server or not self.evol_server._ws:
+            raise RuntimeError("evol_server not connected to Kernel")
+
+        print(f"[relay] Requesting kernel_token for {module_id} from Launcher")
+
+        try:
+            result = await self.evol_server._rpc_call(
+                "launcher.request_client_token",
+                {"module_id": module_id}
+            )
+            token = result.get("token")
+            if not token:
+                raise RuntimeError(f"Launcher returned no token: {result}")
+
+            print(f"[relay] Received kernel_token for {module_id}")
+            return token
+
+        except Exception as e:
+            print(f"[relay] Failed to request kernel_token: {e}")
+            raise
+
+    async def _connect_kernel(self, module_id: str, kernel_token: str):
+        """连接到 Kernel"""
+        url = f"ws://{self.kernel_host}:{self.kernel_port}/ws?token={kernel_token}&id={module_id}"
+        print(f"[relay] DEBUG: Connecting to Kernel with module_id={module_id}, url={url}")
+        kernel_ws = await websockets.connect(
+            url,
+            open_timeout=5,
+            ping_interval=None,
+            close_timeout=10
+        )
+        print(f"[relay] DEBUG: Connected to Kernel")
+
+        # 注册模块（不预订阅事件，由前端自己决定）
+        await self._send_to_kernel(kernel_ws, {
+            "jsonrpc": "2.0",
+            "id": str(uuid.uuid4()),
+            "method": "registry.register",
+            "params": {
+                "module_id": module_id,
+                "module_type": "web_client",
+            }
+        })
+
+        # 发送 module.ready
+        await self._send_to_kernel(kernel_ws, {
+            "jsonrpc": "2.0",
+            "id": str(uuid.uuid4()),
+            "method": "event.publish",
+            "params": {
+                "event_id": str(uuid.uuid4()),
+                "event": "module.ready",
+                "data": {
+                    "module_id": module_id,
+                    "graceful_shutdown": True,
+                }
+            }
+        })
+
+        return kernel_ws
+
+    async def _relay_messages(self, session: SessionInfo):
+        """双向转发消息"""
+        try:
+            # 创建两个任务：client → kernel, kernel → client
+            client_to_kernel = asyncio.create_task(
+                self._forward_client_to_kernel(session)
+            )
+            kernel_to_client = asyncio.create_task(
+                self._forward_kernel_to_client(session)
+            )
+
+            # 等待任一任务完成（断开）
+            done, pending = await asyncio.wait(
+                [client_to_kernel, kernel_to_client],
+                return_when=asyncio.FIRST_COMPLETED
+            )
+
+            # 取消未完成的任务
+            for task in pending:
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+
+            # 检查已完成任务的异常
+            for task in done:
+                try:
+                    task.result()
+                except Exception:
+                    pass  # 忽略异常，因为断开连接是正常的
+
+        except Exception as e:
+            print(f"[relay] Relay error: {e}")
+
+        finally:
+            # 客户端断开，启动重连等待
+            await self._on_client_disconnect(session)
+
+    async def _forward_client_to_kernel(self, session: SessionInfo):
+        """转发客户端消息到 Kernel（带权限检查）"""
+        while True:
+            raw = await session.client_ws.receive_text()
+            msg = json.loads(raw)
+
+            print(f"[relay] Client → Kernel: {msg.get('method', msg.get('type', 'unknown'))}")
+
+            # 处理心跳 ping
+            if msg.get("type") == "ping":
+                await session.client_ws.send_json({"type": "pong"})
+                continue
+
+            # 检查权限
+            if not self._check_permission(session.role, msg):
+                method = msg.get("method", "")
+                # 红色高亮显示权限错误
+                print(f"\033[91m[relay] ✗ 权限被拒绝: {method} (角色: {session.role})\033[0m")
+                # 返回权限错误
+                if "id" in msg:
+                    await session.client_ws.send_json({
+                        "jsonrpc": "2.0",
+                        "id": msg["id"],
+                        "error": {
+                            "code": -32000,
+                            "message": f"Permission denied: {method} (role: {session.role})"
+                        }
+                    })
+                continue
+
+            # 检查是否是 web.* RPC 调用
+            method = msg.get("method", "")
+            if method.startswith("web.") and "id" in msg:
+                print(f"[relay] Intercepted web RPC: {method}")
+                # 拦截并处理 web.* RPC 调用
+                await self._handle_web_rpc(session, msg)
+                continue
+
+            # 转发到 Kernel
+            print(f"[relay] Forwarding to Kernel: {method}")
+            await self._send_to_kernel(session.kernel_ws, msg)
+
+    async def _forward_kernel_to_client(self, session: SessionInfo):
+        """转发 Kernel 消息到客户端"""
+        async for raw in session.kernel_ws:
+            print(f"[relay] Kernel → Client: {len(raw)} bytes")
+            await session.client_ws.send_text(raw)
+
+    def _check_permission(self, role: str, msg: dict) -> bool:
+        """检查 RPC 权限"""
+        method = msg.get("method")
+        if not method:
+            return True  # 非 RPC 请求，放行
+
+        # 获取角色的权限列表
+        allowed = self.permissions.get(role, [])
+
+        # 检查是否匹配
+        for pattern in allowed:
+            if pattern.endswith(".*"):
+                # 通配符匹配
+                prefix = pattern[:-2]
+                if method.startswith(prefix + "."):
+                    return True
+            elif pattern == method:
+                # 精确匹配
+                return True
+
+        return False
+
+    async def _on_client_disconnect(self, session: SessionInfo):
+        """客户端断开，启动重连等待"""
+        session.status = "waiting_reconnect"
+        session.client_ws = None
+
+        print(f"[relay] Client disconnected: {session.module_id}, waiting {self.reconnect_timeout}s for reconnect")
+
+        # 启动超时计时器
+        timer = asyncio.create_task(self._reconnect_timeout(session))
+        self.reconnect_timers[session.session_token] = timer
+
+    async def _reconnect_timeout(self, session: SessionInfo):
+        """重连超时，执行优雅退出"""
+        await asyncio.sleep(self.reconnect_timeout)
+
+        # 超时后仍未重连，执行优雅退出
+        if session.status == "waiting_reconnect":
+            await self._graceful_shutdown(session)
+
+    async def _graceful_shutdown(self, session: SessionInfo):
+        """代表前端执行优雅退出"""
+        print(f"[relay] Graceful shutdown: {session.module_id}")
+
+        try:
+            # 发送 module.exiting 事件（带 token_revoked 标记）
+            await self._send_to_kernel(session.kernel_ws, {
+                "jsonrpc": "2.0",
+                "id": str(uuid.uuid4()),
+                "method": "event.publish",
+                "params": {
+                    "event_id": str(uuid.uuid4()),
+                    "event": "module.exiting",
+                    "data": {
+                        "module_id": session.module_id,
+                        "action": "none",
+                        "token_revoked": True
+                    }
+                }
+            })
+
+            # 发送 module.shutdown.ready 事件
+            await self._send_to_kernel(session.kernel_ws, {
+                "jsonrpc": "2.0",
+                "id": str(uuid.uuid4()),
+                "method": "event.publish",
+                "params": {
+                    "event_id": str(uuid.uuid4()),
+                    "event": "module.shutdown.ready",
+                    "data": {
+                        "module_id": session.module_id
+                    }
+                }
+            })
+
+            # 断开 Kernel 连接
+            await session.kernel_ws.close()
+
+        except Exception as e:
+            print(f"[relay] Graceful shutdown error: {e}")
+
+        finally:
+            # 删除 session
+            if session.session_token in self.sessions:
+                del self.sessions[session.session_token]
+            if session.session_token in self.reconnect_timers:
+                del self.reconnect_timers[session.session_token]
+
+    async def _send_to_kernel(self, kernel_ws, msg: dict):
+        """发送消息到 Kernel"""
+        await kernel_ws.send(json.dumps(msg))
+
+    async def _handle_web_rpc(self, session: SessionInfo, msg: dict):
+        """处理 web.* RPC 调用"""
+        rpc_id = msg.get("id")
+        method = msg.get("method", "")
+        params = msg.get("params", {})
+
+        print(f"[relay] Handling web RPC: {method} (id={rpc_id})")
+
+        # 去掉 web. 前缀
+        if method.startswith("web."):
+            method = method[4:]
+
+        try:
+            # 调用 evol_server 的 RPC 处理器
+            if method == "list_tokens":
+                print(f"[relay] Calling list_tokens")
+                result = await self.evol_server._rpc_list_tokens()
+            elif method == "revoke_token":
+                print(f"[relay] Calling revoke_token with params: {params}")
+                result = await self.evol_server._rpc_revoke_token(params)
+            else:
+                raise ValueError(f"Unknown method: web.{method}")
+
+            print(f"[relay] RPC success: {method}, result keys: {list(result.keys()) if isinstance(result, dict) else type(result)}")
+
+            # 返回结果
+            await session.client_ws.send_json({
+                "jsonrpc": "2.0",
+                "id": rpc_id,
+                "result": result
+            })
+        except Exception as e:
+            # 红色高亮显示 RPC 错误
+            print(f"\033[91m[relay] ✗ RPC 错误: {method}, 错误: {e}\033[0m")
+            # 返回错误
+            await session.client_ws.send_json({
+                "jsonrpc": "2.0",
+                "id": rpc_id,
+                "error": {
+                    "code": -32603,
+                    "message": str(e)
+                }
+            })
+
+    async def close_all_sessions(self):
+        """优雅关闭所有会话（用于 shutdown）"""
+        print(f"[relay] Closing {len(self.sessions)} active sessions...")
+
+        # 取消所有重连定时器
+        for timer in self.reconnect_timers.values():
+            timer.cancel()
+        self.reconnect_timers.clear()
+
+        # 关闭所有会话
+        for session in list(self.sessions.values()):
+            try:
+                # 关闭客户端连接
+                if session.client_ws:
+                    await session.client_ws.close(code=1001, reason="Server shutting down")
+                # 关闭 Kernel 连接
+                if session.kernel_ws:
+                    await session.kernel_ws.close()
+            except Exception as e:
+                print(f"[relay] Error closing session {session.session_token}: {e}")
+
+        self.sessions.clear()
+        print(f"[relay] All sessions closed")