@agentunion/kite 1.5.0 → 1.6.0

package/extensions/services/web/relay.py

@@ -14,6 +14,7 @@ Kernel WebSocket 中转服务
 
 import asyncio
 import json
+import os
 import secrets
 import time
 import uuid
@@ -23,6 +24,9 @@ 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:
     """前端会话信息"""
@@ -35,6 +39,7 @@ class SessionInfo:
         client_ws: WebSocket,
         frontend_token: str,
         role: str,
+        kernel_token: str,
     ):
         self.session_token = session_token
         self.module_id = module_id
@@ -42,6 +47,7 @@ class SessionInfo:
         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
@@ -59,7 +65,9 @@ class KernelRelay:
         reconnect_timeout: int,
         permissions: dict,
         pairing_manager,
-        web_server,  # 新增：web server 实例，用于发送事件
+        web_server,  # web server 实例，用于发送事件
+        auth_manager=None,  # 新增：AuthManager 实例，用于 Kite Token 认证
+        oauth_manager=None,  # 新增：OAuthManager 实例，用于 OAuth 认证
     ):
         """
         初始化中转服务。
@@ -73,6 +81,8 @@ class KernelRelay:
             permissions: 权限配置（角色 → RPC 白名单）
             pairing_manager: 配对管理器实例
             web_server: web server 实例
+            auth_manager: AuthManager 实例（可选，用于 Kite Token 认证）
+            oauth_manager: OAuthManager 实例（可选，用于 OAuth auth_ticket 认证）
         """
         self.kernel_host = kernel_host
         self.kernel_port = kernel_port
@@ -82,6 +92,16 @@ class KernelRelay:
         self.permissions = permissions
         self.pairing_manager = pairing_manager
         self.web_server = web_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] = {}
@@ -89,6 +109,16 @@ class KernelRelay:
         # 待重连的 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  # 秒
+
+        # MFA/TOTP 验证器
+        self._totp = TOTPVerifier()
+        # MFA 密钥存储：token → totp_secret（由外部配置注入）
+        self._mfa_secrets: dict[str, str] = {}
+
     async def handle_client(self, client_ws: WebSocket):
         """
         处理客户端连接。
@@ -100,9 +130,28 @@ class KernelRelay:
         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()
 
         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)
@@ -111,6 +160,9 @@ class KernelRelay:
             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)
@@ -120,7 +172,7 @@ class KernelRelay:
             else:
                 await client_ws.send_json({
                     "type": "error",
-                    "message": "Invalid message type, expected 'request_code', 'pair' or 'auth'"
+                    "message": "Invalid message type, expected 'challenge', 'request_code', 'pair' or 'auth'"
                 })
                 await client_ws.close(code=4000, reason="Invalid message type")
 
@@ -133,6 +185,95 @@ class KernelRelay:
             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):
         """处理请求配对码"""
         # 生成配对码
@@ -141,15 +282,17 @@ class KernelRelay:
         # 发送事件给 Launcher（通过 Kernel）
         if self.web_server and self.web_server._ws:
             try:
-                await self.web_server._publish_event({
-                    "event": "pairing.status",
-                    "data": {
+                await self.web_server._publish_event(
+                    self.web_server._ws,
+                    "pairing.status",
+                    {
                         "step": "code_generated",
                         "success": True,
                         "code": code,
-                        "expires_in": 300
+                        "expires_in": 300,
+                        "module_id": "evol"
                     }
-                })
+                )
             except Exception as e:
                 print(f"[relay] Failed to publish pairing event: {e}")
 
@@ -179,14 +322,15 @@ class KernelRelay:
             # 发送配对失败事件
             if self.web_server and self.web_server._ws:
                 try:
-                    await self.web_server._publish_event({
-                        "event": "pairing.status",
-                        "data": {
+                    await self.web_server._publish_event(
+                        self.web_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}")
 
@@ -195,6 +339,7 @@ class KernelRelay:
                 "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
@@ -213,6 +358,55 @@ class KernelRelay:
         """处理认证请求（已有 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')}")
+            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({
@@ -222,16 +416,47 @@ class KernelRelay:
             await client_ws.close(code=4000, reason="Missing credentials")
             return
 
-        # 验证 frontend_token
+        # 1. 尝试配对码认证
         token_info = self.pairing_manager.verify_token(frontend_token)
-        if not token_info:
+        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)
@@ -240,7 +465,7 @@ class KernelRelay:
                 client_ws,
                 session_token,
                 frontend_token,
-                token_info["role"]
+                role
             )
 
     async def _create_new_connection(
@@ -257,8 +482,11 @@ class KernelRelay:
         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_ws = await self._connect_kernel(module_id, kernel_token)
 
             # 创建 session
             session = SessionInfo(
@@ -268,6 +496,7 @@ class KernelRelay:
                 client_ws=client_ws,
                 frontend_token=frontend_token,
                 role=role,
+                kernel_token=kernel_token,
             )
             self.sessions[session_token] = session
 
@@ -285,15 +514,16 @@ class KernelRelay:
             # 如果是配对，发送配对成功事件给 Launcher
             if is_pairing and self.web_server and self.web_server._ws:
                 try:
-                    await self.web_server._publish_event({
-                        "event": "pairing.status",
-                        "data": {
+                    await self.web_server._publish_event(
+                        self.web_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}")
 
@@ -301,12 +531,39 @@ class KernelRelay:
             await self._relay_messages(session)
 
         except Exception as e:
-            print(f"[relay] Failed to create connection: {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: {str(e)}"
+                "message": f"Failed to connect to Kernel: {error_msg}",
+                "fatal": is_fatal,  # 告知前端是否应该重试
+                "code": error_code
             })
-            await client_ws.close(code=1011, reason="Kernel connection failed")
+            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,
@@ -339,33 +596,63 @@ class KernelRelay:
         # 继续双向转发
         await self._relay_messages(session)
 
-    async def _connect_kernel(self, module_id: str):
+    async def _request_kernel_token(self, module_id: str) -> str:
+        """向 Launcher 申请 kernel_token"""
+        if not self.web_server or not self.web_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.web_server._rpc_call(
+                self.web_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: {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={self.kernel_token}&id={module_id}"
+        url = f"ws://{self.kernel_host}:{self.kernel_port}/ws?id={module_id}"
+        print(f"[relay] DEBUG: Connecting to Kernel with module_id={module_id}")
         kernel_ws = await websockets.connect(
             url,
             open_timeout=5,
-            ping_interval=20,
-            ping_timeout=20,
+            ping_interval=None,
             close_timeout=10
         )
+        print(f"[relay] DEBUG: Connected to Kernel")
 
-        # 订阅事件
+        # 1. 先发送认证请求（Kernel 要求第一条消息必须是 auth）
+        auth_id = str(uuid.uuid4())
         await self._send_to_kernel(kernel_ws, {
             "jsonrpc": "2.0",
-            "id": str(uuid.uuid4()),
-            "method": "event.subscribe",
+            "id": auth_id,
+            "method": "auth",
             "params": {
-                "events": [
-                    "module.started",
-                    "module.stopped",
-                    "module.crashed",
-                    "module.ready",
-                ]
+                "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']}")
+        print(f"[relay] DEBUG: Kernel auth success")
+
+        # 2. 注册模块（不预订阅事件，由前端自己决定）
         await self._send_to_kernel(kernel_ws, {
             "jsonrpc": "2.0",
             "id": str(uuid.uuid4()),
@@ -373,12 +660,6 @@ class KernelRelay:
             "params": {
                 "module_id": module_id,
                 "module_type": "web_client",
-                "events_subscribe": [
-                    "module.started",
-                    "module.stopped",
-                    "module.crashed",
-                    "module.ready",
-                ]
             }
         })
 
@@ -440,10 +721,13 @@ class KernelRelay:
 
     async def _forward_client_to_kernel(self, session: SessionInfo):
         """转发客户端消息到 Kernel（带权限检查）"""
-        while True:
-            raw = await session.client_ws.receive_text()
+        try:
+            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"})
@@ -474,13 +758,35 @@ class KernelRelay:
                 await self._handle_web_rpc(session, msg)
                 continue
 
-            # 转发到 Kernel
+            # 转发到 Kernel（包括 evol.* 调用）
+            # print(f"[relay] Forwarding to Kernel: {method}")
             await self._send_to_kernel(session.kernel_ws, msg)
+        except (WebSocketDisconnect, AttributeError, asyncio.CancelledError,
+                websockets.exceptions.ConnectionClosedOK, websockets.exceptions.ConnectionClosedError):
+            pass  # 客户端断开或 Kernel 连接关闭，正常退出
 
     async def _forward_kernel_to_client(self, session: SessionInfo):
         """转发 Kernel 消息到客户端"""
-        async for raw in session.kernel_ws:
+        try:
+            async for raw in session.kernel_ws:
+            # 解析消息内容用于日志
+            # try:
+            #     msg = json.loads(raw)
+            #     if "method" in msg:
+            #         print(f"[relay] Kernel → Client: {msg['method']} ({len(raw)} bytes)")
+            #     elif "result" in msg:
+            #         print(f"[relay] Kernel → Client: response id={msg.get('id')} ({len(raw)} bytes)")
+            #     elif "error" in msg:
+            #         print(f"[relay] Kernel → Client: error id={msg.get('id')} ({len(raw)} bytes)")
+            #     else:
+            #         print(f"[relay] Kernel → Client: {len(raw)} bytes")
+            # except:
+            #     print(f"[relay] Kernel → Client: {len(raw)} bytes")
             await session.client_ws.send_text(raw)
+        except (websockets.exceptions.ConnectionClosedOK, websockets.exceptions.ConnectionClosedError):
+            pass  # Kernel 连接关闭，正常退出
+        except Exception:
+            pass  # client_ws 断开或被置 None，正常退出
 
     def _check_permission(self, role: str, msg: dict) -> bool:
         """检查 RPC 权限"""
@@ -526,9 +832,10 @@ class KernelRelay:
     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"})
 
         try:
-            # 发送 module.exiting 事件
+            # 发送 module.exiting 事件（带 token_revoked 标记）
             await self._send_to_kernel(session.kernel_ws, {
                 "jsonrpc": "2.0",
                 "id": str(uuid.uuid4()),
@@ -538,7 +845,8 @@ class KernelRelay:
                     "event": "module.exiting",
                     "data": {
                         "module_id": session.module_id,
-                        "action": "none"
+                        "action": "none",
+                        "token_revoked": True
                     }
                 }
             })
@@ -587,7 +895,7 @@ class KernelRelay:
             method = method[4:]
 
         try:
-            # 调用 web_server 的 RPC 处理器
+            # 调用 evol_server 的 RPC 处理器
             if method == "list_tokens":
                 print(f"[relay] Calling list_tokens")
                 result = await self.web_server._rpc_list_tokens()
@@ -600,23 +908,25 @@ class KernelRelay:
             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
-            })
+            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")
-            # 返回错误
-            await session.client_ws.send_json({
-                "jsonrpc": "2.0",
-                "id": rpc_id,
-                "error": {
-                    "code": -32603,
-                    "message": str(e)
-                }
-            })
+            # 返回错误（仅当连接存在时）
+            if session.client_ws:
+                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）"""
@@ -641,3 +951,72 @@ class KernelRelay:
 
         self.sessions.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.web_server or not self.web_server._ws:
+            return
+        try:
+            await self.web_server._publish_event(
+                self.web_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]