@agentunion/kite 1.4.0 → 1.6.0

package/extensions/services/kite_console/relay.py

@@ -0,0 +1,1350 @@
+"""
+Kernel WebSocket 中转服务
+
+功能：
+- 接受前端 WebSocket 连接
+- 处理配对流程（通过 WebSocket）
+- 为每个前端创建到 Kernel 的连接
+- 双向转发 JSON-RPC 消息
+- 管理前端模块的生命周期（注册、重连、优雅退出）
+- RPC 权限控制（白名单）
+
+零共享代码依赖 - 此文件可以独立拷贝到其他模块使用。
+"""
+
+import asyncio
+import json
+import os
+import secrets
+import time
+import uuid
+from collections import deque
+from datetime import datetime, timezone
+from typing import Optional
+
+import websockets
+from fastapi import WebSocket, WebSocketDisconnect
+
+from extensions.services.kite_console.nonce_pool import NoncePool
+from extensions.services.kite_console.mfa_totp import TOTPVerifier
+
+
+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
+        self.relay_task: Optional[asyncio.Task] = None  # 当前 relay 任务，用于重连时取消旧 relay
+        self._msg_timestamps: deque[float] = deque()  # Layer 3: 消息速率跟踪
+        self._reconnecting = False  # 重连期间标志，防止 cancel 触发 offline
+
+
+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 实例，用于发送事件
+        auth_manager=None,  # 新增：AuthManager 实例，用于 Kite Token 认证
+        oauth_manager=None,  # 新增：OAuthManager 实例，用于 OAuth 认证
+    ):
+        """
+        初始化中转服务。
+
+        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 实例
+            auth_manager: AuthManager 实例（可选，用于 Kite Token 认证）
+            oauth_manager: OAuthManager 实例（可选，用于 OAuth auth_ticket 认证）
+        """
+        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
+        self.auth_manager = auth_manager
+        self.oauth_manager = oauth_manager
+
+        # Nonce 池（challenge-response 防重放）
+        self.nonce_pool = NoncePool(max_size=10000, ttl=600)
+
+        # TLS configuration
+        # 默认值根据环境变量：开发环境允许 WS，生产环境强制 WSS
+        env = os.environ.get("KITE_ENV", "development").lower()
+        self.require_tls = os.environ.get("KITE_REQUIRE_TLS", "true" if env == "production" else "false").lower() == "true"
+
+        # session_token → SessionInfo
+        self.sessions: dict[str, SessionInfo] = {}
+
+        # 待重连的 session（超时清理）
+        self.reconnect_timers: dict[str, asyncio.Task] = {}
+
+        # 速率限制：(ip, device_id) → [timestamp, timestamp, ...]
+        self._rate_limits: dict[tuple[str, str], list[float]] = {}
+        self._rate_limit_max = 20  # 10 秒内最多 20 次
+        self._rate_limit_window = 10.0  # 秒
+
+        # Layer 1: Session 级重连节流
+        self._reconnect_timestamps: dict[str, list[float]] = {}
+        self._reconnect_limit_max = 5       # 最多 N 次
+        self._reconnect_limit_window = 30.0  # 每 M 秒
+
+        # Layer 2: 全局并发 Session 上限
+        self._max_sessions = 50
+
+        # Layer 3: 单连接消息速率限制
+        self._msg_rate_limit_max = 200      # 最多 N 条
+        self._msg_rate_limit_window = 10.0  # 每 M 秒
+
+        # Layer 4: 消息大小限制
+        self._max_message_size = 1 * 1024 * 1024  # 1MB
+
+        # MFA/TOTP 验证器
+        self._totp = TOTPVerifier()
+        # MFA 密钥存储：token → totp_secret（由外部配置注入）
+        self._mfa_secrets: dict[str, str] = {}
+
+        # Kernel 离线状态（全局，影响所有客户端）
+        self._kernel_offline = False
+        self._kernel_offline_lock = asyncio.Lock()
+        self.kernel_reconnect_info = {"state": "connected", "attempt": 0, "max_attempts": 10, "next_retry_ms": 0}
+
+    async def handle_client(self, client_ws: WebSocket):
+        """
+        处理客户端连接。
+
+        流程：
+        1. 接受 WebSocket 连接
+        2. 接收认证/配对/请求配对码消息
+        3. 验证身份
+        4. 创建或恢复 Kernel 连接
+        5. 双向转发消息
+        """
+        # Check TLS requirement before accepting
+        if self.require_tls:
+            # Check if connection is secure (WSS)
+            if client_ws.url.scheme != "wss":
+                await client_ws.close(code=1008, reason="TLS required")
+                print(f"[relay] Rejected non-TLS connection from {client_ws.client.host if client_ws.client else 'unknown'}")
+                return
+
+        await client_ws.accept()
+
+        client_ip = client_ws.client.host if client_ws.client else "unknown"
+
+        # Kernel 离线检查：发送状态信息后关闭
+        if self._kernel_offline:
+            try:
+                info = self.kernel_reconnect_info
+                # 根据服务端重连状态计算建议客户端延迟
+                if info["state"] == "fatal" or info["state"] == "exiting":
+                    suggest_ms = 0  # 服务端即将退出，不建议重连
+                elif info["next_retry_ms"] > 0:
+                    suggest_ms = info["next_retry_ms"] + 1000  # 服务端重连后再等 1s
+                else:
+                    suggest_ms = 2000
+                await client_ws.send_json({
+                    "type": "error",
+                    "message": "Kernel is currently unavailable",
+                    "code": 4050,
+                    "retry_after_ms": suggest_ms,
+                    "kernel_status": info,
+                })
+                await client_ws.close(code=4050, reason="Kernel unavailable")
+            except Exception:
+                pass
+            return
+
+        try:
+            # 速率限制检查
+            client_ip = client_ws.client.host if client_ws.client else "unknown"
+            if not self._check_rate_limit(client_ip, "unknown"):
+                await client_ws.send_json({
+                    "type": "error",
+                    "message": "Rate limit exceeded"
+                })
+                await client_ws.close(code=4020, reason="Rate limit exceeded")
+                await self._audit_log("auth.rate_limited", {"ip": client_ip})
+                return
+
+            # 接收认证/配对/请求配对码消息
+            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 == "challenge":
+                # Challenge-response 握手（第一步：客户端请求 challenge）
+                await self._handle_challenge(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 'challenge', 'request_code', 'pair' or 'auth'"
+                })
+                await client_ws.close(code=4000, reason="Invalid message type")
+
+        except WebSocketDisconnect:
+            pass
+        except asyncio.CancelledError:
+            pass  # relay 被 cancel（重连或关闭），正常退出
+        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_challenge(self, client_ws: WebSocket, msg: dict):
+        """处理 challenge 请求（握手第一步）"""
+        client_ip = client_ws.client.host if client_ws.client else "unknown"
+        nonce = self.nonce_pool.generate(metadata={"client_ip": client_ip})
+
+        await client_ws.send_json({
+            "type": "challenge",
+            "nonce": nonce,
+            "protocol_version": "1.0",
+            "auth_methods": ["pairing_code", "kite_token", "oauth"],
+            "expires_in": self.nonce_pool.ttl,
+            "server_time": time.time(),  # 客户端可据此校准时钟偏移
+        })
+
+        # 等待客户端的 connect 消息（带 nonce）
+        try:
+            raw = await asyncio.wait_for(client_ws.receive_text(), timeout=30)
+            connect_msg = json.loads(raw)
+
+            if connect_msg.get("type") != "connect":
+                await client_ws.send_json({
+                    "type": "error",
+                    "message": "Expected 'connect' message after challenge"
+                })
+                await client_ws.close(code=4000, reason="Protocol error")
+                return
+
+            # 验证 nonce
+            returned_nonce = connect_msg.get("nonce")
+            if not returned_nonce:
+                await client_ws.send_json({
+                    "type": "error",
+                    "message": "Missing nonce in connect message"
+                })
+                await client_ws.close(code=4000, reason="Missing nonce")
+                return
+
+            nonce_meta = self.nonce_pool.verify(returned_nonce)
+            if nonce_meta is None:
+                await client_ws.send_json({
+                    "type": "error",
+                    "message": "Invalid or expired nonce"
+                })
+                await client_ws.close(code=4001, reason="Invalid nonce")
+                return
+
+            # Nonce 验证通过，根据 auth_method 分发
+            auth_method = connect_msg.get("auth_method", "")
+
+            if auth_method == "pairing_code":
+                code = connect_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 code")
+                    return
+                await self._handle_pair(client_ws, {"code": code})
+
+            elif auth_method == "kite_token":
+                token = connect_msg.get("token")
+                session_token = connect_msg.get("session_token")
+                if not 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
+                await self._handle_auth(client_ws, {"token": token, "session_token": session_token})
+
+            elif auth_method == "oauth":
+                auth_ticket = connect_msg.get("auth_ticket")
+                if not auth_ticket:
+                    await client_ws.send_json({"type": "error", "message": "Missing auth_ticket"})
+                    await client_ws.close(code=4000, reason="Missing auth_ticket")
+                    return
+                await self._handle_auth(client_ws, {
+                    "method": "oauth",
+                    "auth_ticket": auth_ticket,
+                    "session_token": connect_msg.get("session_token", "")
+                })
+
+            else:
+                await client_ws.send_json({
+                    "type": "error",
+                    "message": f"Unsupported auth_method: {auth_method}"
+                })
+                await client_ws.close(code=4000, reason="Unsupported auth method")
+
+        except asyncio.TimeoutError:
+            await client_ws.send_json({"type": "error", "message": "Connect timeout"})
+            await client_ws.close(code=4000, reason="Timeout")
+
+    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(
+                    self.evol_server._ws,
+                    "pairing.status",
+                    {
+                        "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(
+                        self.evol_server._ws,
+                        "pairing.status",
+                        {
+                            "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")
+            await self._audit_log("auth.pair_failed", {"reason": "invalid_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")
+        auth_method = msg.get("method", "")  # 可选：oauth
+
+        # OAuth auth_ticket 认证
+        if auth_method == "oauth":
+            auth_ticket = msg.get("auth_ticket")
+            if not auth_ticket:
+                await client_ws.send_json({
+                    "type": "error",
+                    "message": "Missing auth_ticket for OAuth auth"
+                })
+                await client_ws.close(code=4000, reason="Missing auth_ticket")
+                return
+
+            if not self.oauth_manager:
+                await client_ws.send_json({
+                    "type": "error",
+                    "message": "OAuth not configured"
+                })
+                await client_ws.close(code=4000, reason="OAuth not configured")
+                return
+
+            ticket_info = self.oauth_manager.verify_auth_ticket(auth_ticket)
+            if not ticket_info:
+                await client_ws.send_json({
+                    "type": "error",
+                    "message": "Invalid or expired auth_ticket"
+                })
+                await client_ws.close(code=4001, reason="Invalid auth_ticket")
+                await self._audit_log("auth.failed", {"reason": "invalid_auth_ticket", "auth_type": "oauth"})
+                return
+            print(f"[Relay] Authenticated via OAuth ({ticket_info['provider']}), user={ticket_info['user_info'].get('name', 'unknown')}")
+            role = "admin"  # OAuth 用户默认 admin
+            await self._audit_log("auth.login", {
+                "auth_type": "oauth",
+                "provider": ticket_info["provider"],
+                "user": ticket_info["user_info"].get("name", "unknown"),
+                "role": role,
+            })
+
+            # 生成 session_token
+            if not session_token:
+                session_token = "sess_" + secrets.token_urlsafe(6)[:6]
+
+            await self._create_new_connection(
+                client_ws,
+                session_token,
+                f"oauth:{ticket_info['provider']}:{ticket_info['user_info'].get('id', '')}",
+                role
+            )
+            return
+
+        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
+
+        # 1. 尝试配对码认证
+        token_info = self.pairing_manager.verify_token(frontend_token)
+        if token_info:
+            role = token_info["role"]
+            auth_type = "pairing"
+            print(f"[Relay] Authenticated via pairing code, role={role}")
+
+        # 2. 尝试 Kite Token 认证
+        elif self.auth_manager and self.auth_manager.verify_kite_token(frontend_token):
+            role = "admin"  # Kite Token 用户默认为 admin
+            auth_type = "kite_token"
+            print(f"[Relay] Authenticated via Kite Token, role={role}")
+            # 构造 token_info 格式以保持兼容
+            token_info = {"role": role}
+
+        # 3. 认证失败
+        else:
+            await client_ws.send_json({
+                "type": "error",
+                "message": "Invalid or expired token"
+            })
+            await client_ws.close(code=4001, reason="Invalid token")
+            await self._audit_log("auth.failed", {"reason": "invalid_token", "session_token": session_token})
+            return
+
+        # 审计：认证成功
+        await self._audit_log("auth.login", {"auth_type": auth_type, "role": role, "session_token": session_token})
+
+        # MFA 检查（可选：如果该 token 配置了 MFA 密钥）
+        mfa_secret = self._mfa_secrets.get(frontend_token)
+        if mfa_secret:
+            mfa_code = msg.get("mfa_code", "")
+            if not mfa_code or not self._totp.verify(mfa_secret, mfa_code):
+                await client_ws.send_json({
+                    "type": "error",
+                    "message": "MFA verification failed"
+                })
+                await client_ws.close(code=4004, reason="MFA failed")
+                await self._audit_log("auth.mfa_failed", {"session_token": session_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,
+                role
+            )
+
+    async def _create_new_connection(
+        self,
+        client_ws: WebSocket,
+        session_token: str,
+        frontend_token: str,
+        role: str,
+        is_pairing: bool = False
+    ):
+        """创建新的 Kernel 连接"""
+        # Layer 2: 全局并发 Session 上限
+        if len(self.sessions) >= self._max_sessions:
+            await client_ws.send_json({
+                "type": "error",
+                "message": "Too many sessions",
+                "max_sessions": self._max_sessions
+            })
+            await client_ws.close(code=4029, reason="Too many sessions")
+            await self._audit_log("auth.session_cap", {
+                "current": len(self.sessions),
+                "max": self._max_sessions
+            })
+            return
+
+        # 生成 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"\033[32m[relay] ← 远程模块连入: {module_id} (role: {role}, ip: {client_ws.client.host if client_ws.client else 'unknown'})\033[0m")
+
+            # 如果是配对，发送配对成功事件给 Launcher
+            if is_pairing and self.evol_server and self.evol_server._ws:
+                try:
+                    await self.evol_server._publish_event(
+                        self.evol_server._ws,
+                        "pairing.status",
+                        {
+                            "step": "completed",
+                            "success": True,
+                            "module_id": module_id,
+                            "role": role
+                        }
+                    )
+                except Exception as e:
+                    print(f"[relay] Failed to publish pairing success event: {e}")
+
+            # 开始双向转发
+            session.relay_task = asyncio.current_task()
+            await self._relay_messages(session)
+
+        except asyncio.CancelledError:
+            # 被 _handle_reconnect cancel，正常退出（不逃逸到 ASGI 层）
+            pass
+
+        except Exception as e:
+            error_msg = str(e)
+            is_fatal = self._is_fatal_error(error_msg)
+            error_code = 1011  # Internal Error (default)
+
+            if is_fatal:
+                # 永久性错误（权限、配置错误）
+                error_code = 1008  # Policy Violation
+                print(f"\033[31m[relay] 致命错误: {error_msg}\033[0m")
+                if "relay_modules whitelist" in error_msg:
+                    print(f"\033[31m[relay] 请在 launcher/module.md 的 relay.modules 中添加本模块名\033[0m")
+            else:
+                print(f"[relay] Failed to create connection: {error_msg}")
+
+            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]
+
+        # Layer 1: Session 级重连节流
+        allowed, cooldown = self._check_reconnect_throttle(session_token)
+        if not allowed:
+            await client_ws.send_json({
+                "type": "error",
+                "message": "Reconnect rate limit exceeded",
+                "cooldown_seconds": cooldown
+            })
+            await client_ws.close(code=4020, reason="Reconnect rate limit exceeded")
+            await self._audit_log("auth.reconnect_throttled", {
+                "session_token": session_token,
+                "cooldown_seconds": cooldown
+            })
+            return
+
+        # 取消超时计时器
+        if session_token in self.reconnect_timers:
+            self.reconnect_timers[session_token].cancel()
+            del self.reconnect_timers[session_token]
+
+        # 标记 session 正在重连，防止旧 relay cancel 触发 offline
+        session._reconnecting = True
+
+        # 先更新 client_ws，使旧 relay finally 的 "session.client_ws is my_client_ws"
+        # 检查返回 False，从而跳过 _on_client_disconnect（避免竞态）
+        session.client_ws = client_ws
+        session.status = "active"
+        session.last_active = time.time()
+
+        # 取消旧的 relay 任务并等待它退出
+        if session.relay_task and not session.relay_task.done():
+            session.relay_task.cancel()
+            try:
+                await session.relay_task
+            except asyncio.CancelledError:
+                pass
+
+        # 复用现有 kernel_ws（前端重连不需要重建 kernel 连接）
+        # 仅当 kernel_ws 已断开时才重新创建
+        kernel_ws = session.kernel_ws
+        kernel_dead = not kernel_ws or kernel_ws.close_code is not None
+        if kernel_dead:
+            try:
+                if kernel_ws:
+                    await kernel_ws.close()
+            except Exception:
+                pass
+            kernel_token = await self._request_kernel_token(session.module_id)
+            kernel_ws = await self._connect_kernel(session.module_id, kernel_token)
+            session.kernel_ws = kernel_ws
+            print(f"[relay] Kernel connection rebuilt for {session.module_id}")
+
+        session._reconnecting = False
+
+        # 返回重连成功
+        await client_ws.send_json({
+            "type": "reconnected",
+            "module_id": session.module_id,
+            "role": session.role
+        })
+
+        print(f"\033[32m[relay] ← 远程模块重连: {session.module_id} (role: {session.role}, ip: {client_ws.client.host if client_ws.client else 'unknown'})\033[0m")
+
+        # 继续双向转发
+        session.relay_task = asyncio.current_task()
+        try:
+            await self._relay_messages(session)
+        except asyncio.CancelledError:
+            # 被下一次 _handle_reconnect cancel，正常退出
+            pass
+
+    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:
+            resp = await self.evol_server._rpc_call(
+                self.evol_server._ws,
+                "launcher.request_client_token",
+                {"module_id": module_id}
+            )
+            # _rpc_call 返回完整 JSON-RPC 响应，需要解包 result 层
+            inner = resp.get("result", resp)
+            token = inner.get("token")
+            if not token:
+                raise RuntimeError(f"Launcher returned no token: {resp}")
+
+            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?id={module_id}"
+        kernel_ws = await websockets.connect(
+            url,
+            open_timeout=5,
+            ping_interval=None,
+            close_timeout=10
+        )
+
+        # 1. 先发送认证请求（Kernel 要求第一条消息必须是 auth）
+        auth_id = str(uuid.uuid4())
+        await self._send_to_kernel(kernel_ws, {
+            "jsonrpc": "2.0",
+            "id": auth_id,
+            "method": "auth",
+            "params": {
+                "token": kernel_token
+            }
+        })
+
+        # 等待认证响应
+        auth_response = await asyncio.wait_for(kernel_ws.recv(), timeout=5.0)
+        auth_msg = json.loads(auth_response)
+        if "error" in auth_msg:
+            raise RuntimeError(f"Kernel auth failed: {auth_msg['error']}")
+
+        # 2. 注册模块（不预订阅事件，由前端自己决定）
+        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",
+                "relay_source": self.evol_server.module_name,
+            }
+        })
+
+        # 发送 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):
+        """双向转发消息"""
+        # 记住本次 relay 持有的 client_ws，用于 finally 中判断是否已被新连接替换
+        my_client_ws = session.client_ws
+        client_to_kernel = None
+        kernel_to_client = None
+        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
+
+            # 用标志位判断是否真的是 Kernel 断开（而非 client send 失败）
+            kernel_disconnected = getattr(session, '_kernel_disconnected', False)
+            graceful_shutdown = getattr(session, '_graceful_shutdown', False)
+
+            # 仅当 Kernel 真正断开时才触发全局 offline（排除主动断开和重连期间）
+            if kernel_disconnected and not graceful_shutdown and not getattr(session, '_reconnecting', False):
+                print(f"[relay] Kernel connection lost for {session.module_id}")
+                asyncio.create_task(self.set_kernel_offline(True))
+
+            # 检查已完成任务的异常
+            for task in done:
+                try:
+                    task.result()
+                except Exception:
+                    pass  # 忽略异常，因为断开连接是正常的
+
+        except asyncio.CancelledError:
+            # 被 _handle_reconnect cancel，必须清理子任务防止竞态
+            await self._cancel_relay_tasks([client_to_kernel, kernel_to_client])
+
+        except Exception as e:
+            print(f"[relay] Relay error: {e}")
+
+        finally:
+            # 兜底：确保子任务不会泄漏（正常路径已在 try/except 中清理，这里处理遗漏）
+            await self._cancel_relay_tasks([client_to_kernel, kernel_to_client])
+
+            # 只有当 session.client_ws 仍是本次 relay 持有的 ws 时才触发断开流程
+            # 如果已被 _handle_reconnect 替换为新连接，跳过（避免清掉新连接）
+            if session.client_ws is my_client_ws:
+                await self._on_client_disconnect(session)
+
+    async def _cancel_relay_tasks(self, tasks, timeout=3):
+        """取消转发子任务，带超时保护防止卡死"""
+        pending = []
+        for task in tasks:
+            if task and not task.done():
+                task.cancel()
+                pending.append(task)
+        if not pending:
+            return
+        # 批量等待所有子任务退出，最多等 timeout 秒
+        done, stuck = await asyncio.wait(pending, timeout=timeout)
+        if stuck:
+            print(f"[relay] Warning: {len(stuck)} relay subtask(s) did not exit within {timeout}s, abandoning")
+        # 收集已完成任务的异常，防止 "Task exception was never retrieved"
+        # 注意：Python 3.9+ CancelledError 继承自 BaseException，不是 Exception
+        for task in done:
+            try:
+                task.result()
+            except BaseException:
+                pass
+
+    async def _forward_client_to_kernel(self, session: SessionInfo):
+        """转发客户端消息到 Kernel（带权限检查）"""
+        try:
+            while True:
+                raw = await session.client_ws.receive_text()
+
+                # Kernel 已离线，不再转发
+                if self._kernel_offline:
+                    break
+
+                # Layer 4: 消息大小限制
+                if len(raw) > self._max_message_size:
+                    await session.client_ws.send_json({
+                        "type": "error",
+                        "message": "Message too large",
+                        "max_bytes": self._max_message_size,
+                        "actual_bytes": len(raw)
+                    })
+                    await session.client_ws.close(code=4008, reason="Message too large")
+                    await self._audit_log("relay.msg_too_large", {
+                        "session_token": session.session_token,
+                        "module_id": session.module_id,
+                        "size": len(raw)
+                    })
+                    return
+
+                # Layer 3: 单连接消息速率限制（超限时丢弃消息 + 警告，不断连）
+                if not self._check_msg_rate_limit(session):
+                    # 每次超限只发一次警告（避免警告本身也占带宽）
+                    if not getattr(session, '_rate_warn_sent', False):
+                        try:
+                            await session.client_ws.send_json({
+                                "type": "warning",
+                                "message": "Message rate limit exceeded, messages are being dropped"
+                            })
+                        except Exception:
+                            pass
+                        session._rate_warn_sent = True
+                        print(f"[relay] Message rate limit hit for {session.module_id}, dropping messages")
+                    continue
+                else:
+                    # 速率恢复，重置警告标志
+                    if getattr(session, '_rate_warn_sent', False):
+                        session._rate_warn_sent = False
+
+                msg = json.loads(raw)
+
+                # 处理心跳 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}")
+                    await self._handle_web_rpc(session, msg)
+                    continue
+
+                # 转发到 Kernel
+                await self._send_to_kernel(session.kernel_ws, msg)
+        except (WebSocketDisconnect, AttributeError, asyncio.CancelledError,
+                websockets.exceptions.ConnectionClosedOK, websockets.exceptions.ConnectionClosedError):
+            pass  # 客户端断开、Kernel 连接关闭或 client_ws 已被置 None，正常退出
+
+    async def _forward_kernel_to_client(self, session: SessionInfo):
+        """转发 Kernel 消息到客户端
+
+        通过 session._kernel_disconnected 标志区分断开原因：
+        - True: kernel_ws 迭代结束或异常（Kernel 真正断开）
+        - False: client_ws.send_text 失败（客户端断开，非 Kernel 问题）
+        """
+        session._kernel_disconnected = False
+        try:
+            async for raw in session.kernel_ws:
+                if session.client_ws is None:
+                    break
+                try:
+                    await session.client_ws.send_text(raw)
+                except Exception:
+                    return  # client send 失败，_kernel_disconnected 保持 False
+            # kernel_ws 迭代正常结束 → Kernel 断开
+            session._kernel_disconnected = True
+        except Exception:
+            # kernel_ws 异常（连接关闭等）→ Kernel 断开
+            session._kernel_disconnected = True
+
+    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}")
+        await self._audit_log("auth.logout", {"module_id": session.module_id, "reason": "reconnect_timeout"})
+
+        # 标记这是主动断开，防止触发全局 Kernel 离线
+        session._graceful_shutdown = True
+
+        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]
+            if session.session_token in self._reconnect_timestamps:
+                del self._reconnect_timestamps[session.session_token]
+
+    async def set_kernel_offline(self, offline: bool):
+        """设置 kernel 离线状态（由 server.py 调用）"""
+        async with self._kernel_offline_lock:
+            if self._kernel_offline == offline:
+                return
+            self._kernel_offline = offline
+
+            if offline:
+                print(f"[relay] Kernel offline detected, closing all {len(self.sessions)} client connections")
+                await self._close_all_clients_with_reason(
+                    code=4050,
+                    reason="Kernel unavailable",
+                    retry_after_ms=2000
+                )
+            else:
+                print(f"[relay] Kernel back online, accepting new connections")
+
+    async def _close_all_clients_with_reason(self, code: int, reason: str, retry_after_ms: int):
+        """关闭所有客户端连接并发送错误信息"""
+        info = self.kernel_reconnect_info
+        for session in list(self.sessions.values()):
+            try:
+                if session.client_ws:
+                    await session.client_ws.send_json({
+                        "type": "error",
+                        "message": reason,
+                        "code": code,
+                        "retry_after_ms": retry_after_ms,
+                        "kernel_status": info,
+                    })
+                    await session.client_ws.close(code=code, reason=reason)
+            except Exception as e:
+                print(f"[relay] Failed to close client {session.module_id}: {e}")
+
+            try:
+                if session.kernel_ws:
+                    await session.kernel_ws.close()
+            except Exception:
+                pass
+
+        # 清理所有 session
+        self.sessions.clear()
+        self.reconnect_timers.clear()
+        self._reconnect_timestamps.clear()
+
+    async def _send_to_kernel(self, kernel_ws, msg: dict):
+        """发送消息到 Kernel"""
+        try:
+            await kernel_ws.send(json.dumps(msg))
+        except (websockets.exceptions.ConnectionClosedOK, websockets.exceptions.ConnectionClosedError):
+            raise  # 向上抛出，由调用方处理
+
+    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)
+            elif method == "disconnect_client":
+                result = await self._rpc_disconnect_client(session, 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)}")
+
+            # 返回结果
+            if session.client_ws:
+                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")
+            # 返回错误（仅当连接存在时）
+            if session.client_ws:
+                await session.client_ws.send_json({
+                    "jsonrpc": "2.0",
+                    "id": rpc_id,
+                    "error": {
+                        "code": -32603,
+                        "message": str(e)
+                    }
+                })
+
+    async def _rpc_disconnect_client(self, caller_session: SessionInfo, params: dict) -> dict:
+        """断开指定远程客户端的连接并吊销其 token"""
+        target_module_id = params.get("module_id")
+        if not target_module_id:
+            raise ValueError("module_id required")
+
+        # 不允许断开自己
+        if target_module_id == caller_session.module_id:
+            raise ValueError("Cannot disconnect self")
+
+        # 查找目标 session
+        target_session = None
+        for s in self.sessions.values():
+            if s.module_id == target_module_id:
+                target_session = s
+                break
+
+        if not target_session:
+            raise ValueError(f"Session not found: {target_module_id}")
+
+        print(f"[relay] Disconnecting client {target_module_id} (requested by {caller_session.module_id})")
+
+        # 吊销 token，使其需要重新配对
+        revoked = self.pairing_manager.revoke_token(target_session.frontend_token)
+        if revoked:
+            print(f"[relay] Token revoked for {target_module_id}")
+        else:
+            print(f"[relay] Token revoke skipped for {target_module_id} (not a pairing token or already revoked)")
+
+        # 先通知被断开的客户端（用 4001 关闭码，前端会停止重连）
+        try:
+            if target_session.client_ws:
+                await target_session.client_ws.close(code=4001, reason="Token revoked")
+        except Exception:
+            pass
+
+        # 执行优雅退出（清理 Kernel 连接和 session）
+        await self._graceful_shutdown(target_session)
+
+        return {"disconnected": target_module_id}
+
+    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()
+
+        # 关闭所有会话（不等待对端确认，shutdown 场景直接断开）
+        for session in list(self.sessions.values()):
+            try:
+                if session.kernel_ws:
+                    await session.kernel_ws.close()
+            except Exception:
+                pass
+            try:
+                if session.client_ws:
+                    # 客户端 close 不 await，避免等浏览器回 close frame 卡住
+                    asyncio.create_task(session.client_ws.close(code=1001, reason="Server shutting down"))
+            except Exception:
+                pass
+
+        self.sessions.clear()
+        self._reconnect_timestamps.clear()
+        print(f"[relay] All sessions closed")
+
+    # ── 审计日志 ──
+
+    async def _audit_log(self, event: str, data: dict):
+        """
+        发布 auth.* 审计事件到 Kernel。
+
+        Args:
+            event: 事件名（如 auth.login, auth.logout, auth.failed）
+            data: 事件数据
+        """
+        if not self.evol_server or not self.evol_server._ws:
+            return
+        try:
+            await self.evol_server._publish_event(
+                self.evol_server._ws,
+                event,
+                {
+                    **data,
+                    "timestamp": time.time(),
+                    "source": "relay",
+                }
+            )
+        except Exception as e:
+            print(f"[relay] Audit log failed: {e}")
+
+    # ── 速率限制 ──
+
+    def _check_rate_limit(self, ip: str, device_id: str) -> bool:
+        """
+        检查 (IP, device_id) 是否超出速率限制。
+
+        Returns:
+            True 表示允许，False 表示被限制
+        """
+        key = (ip, device_id)
+        now = time.time()
+        cutoff = now - self._rate_limit_window
+
+        # 获取或创建时间戳列表
+        timestamps = self._rate_limits.get(key, [])
+
+        # 清理过期的时间戳
+        timestamps = [t for t in timestamps if t > cutoff]
+
+        # 检查是否超限
+        if len(timestamps) >= self._rate_limit_max:
+            self._rate_limits[key] = timestamps
+            return False
+
+        # 记录本次请求
+        timestamps.append(now)
+        self._rate_limits[key] = timestamps
+
+        # 定期清理长期不活跃的 key（避免内存泄漏）
+        if len(self._rate_limits) > 1000:
+            self._cleanup_rate_limits(now)
+
+        return True
+
+    def _cleanup_rate_limits(self, now: float):
+        """清理不活跃的速率限制条目"""
+        cutoff = now - self._rate_limit_window * 10  # 10 个窗口期未活跃则清理
+        expired = [
+            k for k, ts_list in self._rate_limits.items()
+            if not ts_list or ts_list[-1] < cutoff
+        ]
+        for k in expired:
+            del self._rate_limits[k]
+
+    @staticmethod
+    def _sliding_window_check(ts: deque, window: float, max_count: int) -> bool:
+        """滑动窗口速率检查（公共逻辑）。返回 True=允许, False=超限"""
+        now = time.time()
+        cutoff = now - window
+        while ts and ts[0] < cutoff:
+            ts.popleft()
+        if len(ts) >= max_count:
+            return False
+        ts.append(now)
+        return True
+
+    def _check_reconnect_throttle(self, session_token: str) -> tuple[bool, float]:
+        """Layer 1: Session 级重连节流。返回 (allowed, cooldown_seconds)"""
+        if session_token not in self._reconnect_timestamps:
+            self._reconnect_timestamps[session_token] = deque()
+
+        ts = self._reconnect_timestamps[session_token]
+        if not self._sliding_window_check(ts, self._reconnect_limit_window, self._reconnect_limit_max):
+            cooldown = ts[0] + self._reconnect_limit_window - time.time()
+            return False, max(cooldown, 1.0)
+
+        # 防内存泄漏：超过 1000 个 session 时清理不活跃的
+        if len(self._reconnect_timestamps) > 1000:
+            stale_cutoff = time.time() - self._reconnect_limit_window * 10
+            stale = [k for k, v in self._reconnect_timestamps.items() if not v or v[-1] < stale_cutoff]
+            for k in stale:
+                del self._reconnect_timestamps[k]
+
+        return True, 0
+
+    def _check_msg_rate_limit(self, session: 'SessionInfo') -> bool:
+        """Layer 3: 单连接消息速率限制"""
+        return self._sliding_window_check(
+            session._msg_timestamps, self._msg_rate_limit_window, self._msg_rate_limit_max
+        )