@agentunion/kite 1.3.1 → 1.4.0

package/extensions/services/web/server.py

@@ -8,13 +8,14 @@ Connects to Kernel via WebSocket JSON-RPC 2.0 for event publishing and subscript
 import asyncio
 import json
 import logging
+import os
 import time
 import uuid
 from datetime import datetime, timezone
 from pathlib import Path
 
 import websockets
-from fastapi import FastAPI
+from fastapi import FastAPI, WebSocket
 from fastapi.staticfiles import StaticFiles
 
 from vendor import config as cfg
@@ -30,10 +31,26 @@ from routes.routes_contacts import router as contacts_router
 from routes.routes_stats import router as stats_router
 from routes.routes_voicechat import router as voicechat_router
 from routes.routes_devlog import router as devlog_router
+from routes.routes_rpc import router as rpc_router, set_web_server
+from routes.routes_management_ws import router as management_ws_router, broadcast_event
+from routes.routes_test import router as test_router
+
+from config_loader import load_business_configs
+from pairing import PairingManager
+from relay import KernelRelay
 
 logger = logging.getLogger(__name__)
 
 
+# System broadcast events (received by all modules, may not need handling)
+SYSTEM_BROADCAST_EVENTS = {
+    "module.ready", "module.registered", "module.started", "module.stopped",
+    "module.crashed", "module.exiting", "module.offline",
+    "module.shutdown.ack", "module.shutdown.ready",
+    "system.ready", "registry.updated",
+}
+
+
 class WebServer:
 
     def __init__(self, token: str = "", kernel_port: int = 0,
@@ -46,10 +63,11 @@ class WebServer:
         self._ws_task: asyncio.Task | None = None
         self._test_task: asyncio.Task | None = None
         self._ws: object | None = None
-        self._ready_sent = False
         self._shutting_down = False
+        self._exit_code = 0  # Exit code for main() to use
         self._uvicorn_server = None  # set by entry.py for graceful shutdown
         self._start_time = time.time()
+        self._rpc_futures = {}  # Store pending RPC request futures
         self.bt_manager: BluetoothManager | None = None
         self.task_manager: TaskManager | None = None
         self.app = self._create_app()
@@ -60,6 +78,40 @@ class WebServer:
 
         @app.on_event("startup")
         async def _startup():
+            # Load business configurations
+            module_dir = Path(__file__).parent
+            business_configs = load_business_configs(str(module_dir))
+
+            # Get relay service config
+            relay_business = business_configs.get('relay_service')
+            if relay_business:
+                relay_config = relay_business['config']
+
+                # Initialize pairing manager
+                auth_config = relay_config['auth']
+                pairing_file = module_dir / auth_config['pairing_code_file']
+                pairing_manager = PairingManager(
+                    pairing_file=str(pairing_file),
+                    code_length=auth_config['pairing_code_length'],
+                    token_expiry=auth_config['token_expiry']
+                )
+                app.state.pairing_manager = pairing_manager
+                print(f"[web] Pairing manager initialized")
+
+                # Initialize relay service
+                relay_service = KernelRelay(
+                    kernel_host="127.0.0.1",
+                    kernel_port=server.kernel_port,
+                    kernel_token=server.token,
+                    base_module_id=relay_config['relay']['base_module_id'],
+                    reconnect_timeout=relay_config['relay']['reconnect_timeout'],
+                    permissions=relay_config['permissions'],
+                    pairing_manager=pairing_manager,
+                    web_server=server  # 传递 web server 实例
+                )
+                app.state.relay_service = relay_service
+                print(f"[web] Relay service initialized")
+
             # Load configuration
             cfg.load_config()
             load_err = cfg.get_load_error()
@@ -142,11 +194,63 @@ class WebServer:
         app.include_router(stats_router, prefix="/api")
         app.include_router(voicechat_router)  # no prefix (has own /ws/ and /api/ paths)
         app.include_router(devlog_router, prefix="/api")
+        app.include_router(rpc_router, prefix="/api")
+        app.include_router(test_router, prefix="/api")
+        app.include_router(management_ws_router)  # /ws/management
+
+        # Relay WebSocket endpoint
+        @app.websocket("/ws/relay")
+        async def relay_endpoint(ws: WebSocket):
+            """Kernel 中转服务 WebSocket 端点"""
+            relay_service = app.state.relay_service
+            if relay_service:
+                await relay_service.handle_client(ws)
+            else:
+                await ws.close(code=1011, reason="Relay service not initialized")
+
+        # Set web server reference for RPC forwarding
+        set_web_server(server)
 
         # Serve frontend static files
+        # IMPORTANT: Do NOT use app.mount("/", ...) as it will intercept WebSocket routes
+        # Instead, explicitly serve HTML files and static assets
         static_dir = Path(__file__).parent / "static"
         if static_dir.exists():
-            app.mount("/", StaticFiles(directory=str(static_dir), html=True), name="web")
+            from fastapi.responses import FileResponse, JSONResponse
+
+            @app.get("/")
+            async def serve_index():
+                """Serve index.html at root path."""
+                index_path = static_dir / "index.html"
+                if index_path.exists():
+                    return FileResponse(index_path)
+                return {"message": "Kite Web Management"}
+
+            @app.get("/pairing.html")
+            async def serve_pairing():
+                """Serve pairing.html."""
+                pairing_path = static_dir / "pairing.html"
+                if pairing_path.exists():
+                    return FileResponse(pairing_path)
+                return JSONResponse({"error": "Not found"}, status_code=404)
+
+            # Mount static files at /static prefix to avoid route conflicts
+            app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")
+
+            # Also serve common static paths directly from root for backward compatibility
+            @app.get("/js/{file_path:path}")
+            async def serve_js(file_path: str):
+                file = static_dir / "js" / file_path
+                if file.exists() and file.is_file():
+                    return FileResponse(file)
+                return JSONResponse({"error": "Not found"}, status_code=404)
+
+            @app.get("/css/{file_path:path}")
+            async def serve_css(file_path: str):
+                file = static_dir / "css" / file_path
+                if file.exists() and file.is_file():
+                    return FileResponse(file)
+                return JSONResponse({"error": "Not found"}, status_code=404)
 
         return app
 
@@ -154,107 +258,183 @@ class WebServer:
 
     async def _ws_loop(self):
         """Connect to Kernel, subscribe, register, and listen. Reconnect on failure."""
-        retry_delay = 0.5  # start with 0.5s
-        max_delay = 30  # cap at 30s
+        retry_delay = 0.3
+        max_delay = 5.0
+        max_retries = 10
+        attempt = 0
         while not self._shutting_down:
             try:
                 await self._ws_connect()
-                retry_delay = 0.5  # reset on successful connection
+                retry_delay = 0.3  # reset on successful connection
+                attempt = 0
             except asyncio.CancelledError:
+                print(f"[web] WS loop cancelled")
                 return
             except Exception as e:
-                print(f"[web] Kernel connection error: {e}, retrying in {retry_delay:.1f}s")
+                attempt += 1
+                # Auth failure — don't retry
+                if hasattr(e, 'rcvd') and e.rcvd is not None:
+                    code = e.rcvd.code if hasattr(e.rcvd, 'code') else 0
+                    if code in (4001, 4003):
+                        print(f"[web] Kernel 认证失败 (code {code})，退出")
+                        self._exit_code = 1
+                        self._shutting_down = True
+                        if self._uvicorn_server:
+                            self._uvicorn_server.should_exit = True
+                        return
+                if attempt >= max_retries:
+                    print(f"[web] Kernel 重连失败 {max_retries} 次，退出")
+                    self._exit_code = 1
+                    self._shutting_down = True
+                    if self._uvicorn_server:
+                        self._uvicorn_server.should_exit = True
+                    return
+                if self._shutting_down:
+                    print(f"[web] Shutting down, not retrying connection")
+                    return
+                print(f"[web] Kernel connection error: {e}, retrying in {retry_delay:.1f}s ({attempt}/{max_retries})")
             self._ws = None
             if self._shutting_down:
+                print(f"[web] Shutting down, exiting WS loop")
                 return
             await asyncio.sleep(retry_delay)
-            retry_delay = min(retry_delay * 2, max_delay)  # exponential backoff
+            retry_delay = min(retry_delay * 2, max_delay)
 
     async def _ws_connect(self):
         """Single WebSocket session: connect, register, subscribe, receive loop."""
         url = f"ws://127.0.0.1:{self.kernel_port}/ws?token={self.token}&id=web"
         print(f"[web] WS connecting to Kernel")
-        async with websockets.connect(url, open_timeout=5, ping_interval=None, ping_timeout=None, close_timeout=10) as ws:
-            self._ws = ws
-            elapsed = time.monotonic() - self.boot_t0 if self.boot_t0 else 0
-            elapsed_str = f" ({elapsed:.1f}s)" if elapsed else ""
-            print(f"[web] Connected to Kernel{elapsed_str}")
-
-            # Subscribe to events
-            await self._rpc_call(ws, "event.subscribe", {
-                "events": [
-                    "module.started",
-                    "module.stopped",
-                    "module.shutdown",
-                ],
-            })
-
-            # Register to Kernel Registry via RPC
-            await self._rpc_call(ws, "registry.register", {
-                "module_id": "web",
-                "module_type": "service",
-                "api_endpoint": f"http://127.0.0.1:{self.port}",
-                "health_endpoint": "/health",
-                "events_publish": {
-                    "web.test": {"description": "Test event from web module"},
-                    "web.started": {"description": "Web UI started with access URL"},
-                },
-                "events_subscribe": [
-                    "module.started",
-                    "module.stopped",
-                    "module.shutdown",
-                ],
-            })
-            print(f"[web] Registered to Kernel{elapsed_str}")
-
-            # Send module.ready (once) so Launcher knows we're up
-            if not self._ready_sent:
-                await self._rpc_call(ws, "event.publish", {
-                    "event_id": str(uuid.uuid4()),
-                    "event": "module.ready",
-                    "data": {
-                        "module_id": "web",
-                        "graceful_shutdown": True,
-                    },
-                })
-                self._ready_sent = True
+        try:
+            async with websockets.connect(url, open_timeout=5, ping_interval=20, ping_timeout=20, close_timeout=10) as ws:
+                self._ws = ws
                 elapsed = time.monotonic() - self.boot_t0 if self.boot_t0 else 0
                 elapsed_str = f" ({elapsed:.1f}s)" if elapsed else ""
-                print(f"[web] module.ready sent{elapsed_str}")
-
-                # Publish web.started event with access URL
-                display_host = "localhost" if self.host == "0.0.0.0" else self.host
-                access_url = f"http://{display_host}:{self.port}"
-                await self._publish_event({
-                    "event": "web.started",
-                    "data": {
-                        "module_id": "web",
-                        "url": access_url,
-                        "host": self.host,
-                        "port": self.port,
-                    },
+                print(f"[web] Connected to Kernel{elapsed_str}")
+
+                # Subscribe to events
+                await self._rpc_call(ws, "event.subscribe", {
+                    "events": [
+                        "module.started",
+                        "module.stopped",
+                        "module.crashed",
+                        "module.ready",
+                        "module.exiting",
+                        "module.shutdown",
+                        "module.shutdown.ack",
+                        "module.shutdown.ready",
+                    ],
                 })
 
-            # Receive loop
-            async for raw in ws:
-                try:
-                    msg = json.loads(raw)
-                except (json.JSONDecodeError, TypeError):
-                    continue
+                # Register to Kernel Registry via RPC
+                await self._rpc_call(ws, "registry.register", {
+                    "module_id": "web",
+                    "module_type": "service",
+                    "api_endpoint": f"http://127.0.0.1:{self.port}",
+                    "health_endpoint": "/health",
+                    "tools": {
+                        "rpc": {
+                            "module": {
+                                "health": {"method": "health", "description": "健康检查"},
+                                "status": {"method": "status", "description": "状态查询"}
+                            },
+                            "web": {
+                                "list_tokens": {"method": "list_tokens", "description": "列出所有令牌"},
+                                "revoke_token": {"method": "revoke_token", "description": "撤销令牌"}
+                            }
+                        }
+                    },
+                    "events_publish": {
+                        "web": {
+                            "test": {"description": "Test event from web module"},
+                            "started": {"description": "Web UI started with access URL"},
+                        }
+                    },
+                    "events_subscribe": [
+                        "module.started",
+                        "module.stopped",
+                        "module.crashed",
+                        "module.ready",
+                        "module.exiting",
+                        "module.shutdown",
+                        "module.shutdown.ack",
+                        "module.shutdown.ready",
+                    ],
+                })
+                print(f"[web] Registered to Kernel{elapsed_str}")
+
+                # Send module.ready (every reconnect, not just first time)
+                if not self._shutting_down:
+                    await self._rpc_call(ws, "event.publish", {
+                        "event_id": str(uuid.uuid4()),
+                        "event": "module.ready",
+                        "data": {
+                            "module_id": "web",
+                            "graceful_shutdown": True,
+                        },
+                    })
+                    elapsed = time.monotonic() - self.boot_t0 if self.boot_t0 else 0
+                    elapsed_str = f" ({elapsed:.1f}s)" if elapsed else ""
+                    print(f"[web] module.ready sent{elapsed_str}")
+
+                    # Publish web.started event with access URL
+                    display_host = "localhost" if self.host == "0.0.0.0" else self.host
+                    access_url = f"http://{display_host}:{self.port}"
+                    await self._publish_event({
+                        "event": "web.started",
+                        "data": {
+                            "module_id": "web",
+                            "url": access_url,
+                            "host": self.host,
+                            "port": self.port,
+                        },
+                    })
+
+                # Receive loop
+                # CRITICAL: RPC 死锁防范
+                # - 入站 RPC 请求必须用 create_task() 异步执行，不可 await
+                # - 原因：如果 handler 内部调用 rpc_call() 发出站请求，出站响应需要本接收循环来分发
+                # - 如果接收循环被 await handler 阻塞，出站响应永远收不到 → 超时死锁
+                # - 事件通知和 RPC 响应可以同步处理（它们不会反向调用 rpc_call）
+                print(f"[web] Entering receive loop")
+
+                # Start heartbeat loop
+                heartbeat_task = asyncio.create_task(self._heartbeat_loop(ws))
 
                 try:
-                    has_method = "method" in msg
-                    has_id = "id" in msg
-
-                    if has_method and not has_id:
-                        # Event Notification
-                        await self._handle_event_notification(msg)
-                    elif has_method and has_id:
-                        # Incoming RPC request
-                        await self._handle_rpc_request(ws, msg)
-                    # Ignore RPC responses (we don't await them in this simple impl)
+                    async for raw in ws:
+                        try:
+                            msg = json.loads(raw)
+                        except (json.JSONDecodeError, TypeError):
+                            continue
+
+                        try:
+                            has_method = "method" in msg
+                            has_id = "id" in msg
+                            has_result_or_error = "result" in msg or "error" in msg
+
+                            if has_method and not has_id:
+                                # Event Notification
+                                await self._handle_event_notification(msg)
+                            elif has_method and has_id:
+                                # Incoming RPC request — run in background to prevent deadlock
+                                asyncio.create_task(self._handle_rpc_request(ws, msg))
+                            elif has_id and has_result_or_error:
+                                # RPC response — resolve pending future
+                                rpc_id = msg.get("id")
+                                if rpc_id in self._rpc_futures:
+                                    self._rpc_futures[rpc_id].set_result(msg)
+                        except Exception as e:
+                            print(f"[web] 消息处理异常（已忽略）: {e}")
                 except Exception as e:
-                    print(f"[web] 消息处理异常（已忽略）: {e}")
+                    print(f"[web] Receive loop exited with exception: {e}")
+                finally:
+                    print(f"[web] Receive loop ended")
+        except Exception as e:
+            print(f"[web] WebSocket connection error: {e}")
+            raise
+        finally:
+            print(f"[web] WebSocket connection closed")
+            self._ws = None
 
     async def _rpc_call(self, ws, method: str, params: dict = None):
         """Send a JSON-RPC 2.0 request (fire-and-forget, no response awaited)."""
@@ -263,19 +443,60 @@ class WebServer:
             msg["params"] = params
         await ws.send(json.dumps(msg))
 
+    async def _heartbeat_loop(self, ws):
+        """Send registry.heartbeat every 30 seconds to prevent TTL expiration."""
+        while True:
+            try:
+                await asyncio.sleep(30)
+                if not self._shutting_down:
+                    await self._rpc_call(ws, "registry.heartbeat", {"module_id": "web"})
+            except Exception as e:
+                print(f"[web] Heartbeat error: {e}")
+                break
+
     async def _handle_event_notification(self, msg: dict):
         """Handle an event notification (JSON-RPC 2.0 Notification with method='event')."""
         params = msg.get("params", {})
         event_type = params.get("event", "")
         data = params.get("data", {})
 
-        # Special handling for module.shutdown targeting web
-        if event_type == "module.shutdown" and data.get("module_id") == "web":
-            await self._handle_shutdown()
+        # Log all events for debugging
+        print(f"[web] Event received: {event_type}, data: {data}")
+
+        # Special handling for module.shutdown
+        if event_type == "module.shutdown":
+            target = data.get("module_id", "")
+            reason = data.get("reason", "")
+            print(f"[web] Shutdown event: target={target}, reason={reason}")
+            # Handle both targeted shutdown (module_id == "web") and broadcast shutdown (no module_id or launcher_lost)
+            if target == "web" or not target or reason == "launcher_lost":
+                print(f"[web] Handling shutdown...")
+                await self._handle_shutdown()
+                return
+            else:
+                print(f"[web] Ignoring shutdown (not for us)")
+                return
+
+        # Forward module status events to management WebSocket clients
+        if event_type in (
+            "module.started",
+            "module.stopped",
+            "module.crashed",
+            "module.ready",
+            "module.exiting",
+            "module.shutdown.ack",
+            "module.shutdown.ready",
+        ):
+            await broadcast_event(event_type, data)
             return
 
-        # Log other events
-        print(f"[web] Event received: {event_type}")
+        # Layer 2: 忽略其他系统广播事件
+        if event_type in SYSTEM_BROADCAST_EVENTS:
+            return
+
+        # Layer 3: 警告未知事件（仅开发环境）
+        if os.environ.get("KITE_ENV") == "development":
+            print(f"[web] Debug: Unhandled event: {event_type}")
 
     async def _handle_rpc_request(self, ws, msg: dict):
         """Handle an incoming RPC request (web.* methods)."""
@@ -283,9 +504,15 @@ class WebServer:
         method = msg.get("method", "")
         params = msg.get("params", {})
 
+        # Strip "web." prefix if present
+        if method.startswith("web."):
+            method = method[4:]
+
         handlers = {
             "health": lambda p: self._rpc_health(),
             "status": lambda p: self._rpc_status(),
+            "list_tokens": lambda p: self._rpc_list_tokens(),
+            "revoke_token": lambda p: self._rpc_revoke_token(p),
         }
         handler = handlers.get(method)
         if handler:
@@ -320,44 +547,163 @@ class WebServer:
             "uptime_seconds": round(time.time() - self._start_time),
         }
 
+    async def _rpc_list_tokens(self) -> dict:
+        """RPC handler for web.list_tokens."""
+        from extensions.services.web.pairing import PairingManager
+        from pathlib import Path
+
+        # Get pairing manager from app state
+        pairing_manager = self.app.state.pairing_manager
+        if not pairing_manager:
+            # Fallback: create temporary instance
+            pairing_file = Path(__file__).parent / "pairing_codes.jsonl"
+            pairing_manager = PairingManager(str(pairing_file))
+
+        # Read all token records
+        codes = pairing_manager._read_codes()
+
+        # Build a set of revoked tokens
+        revoked_tokens = {
+            record.get("token")
+            for record in codes
+            if record.get("status") == "revoked" and record.get("token")
+        }
+
+        # Filter only used tokens (status="used") that are not revoked
+        tokens = [
+            {
+                "token": record.get("token"),
+                "role": record.get("role", "viewer"),
+                "paired_at": record.get("paired_at"),
+                "expires_at": record.get("expires_at"),
+                "code": record.get("code", "")
+            }
+            for record in codes
+            if record.get("status") == "used"
+               and record.get("token")
+               and record.get("token") not in revoked_tokens
+        ]
+
+        return {"tokens": tokens}
+
+    async def _rpc_revoke_token(self, params: dict) -> dict:
+        """RPC handler for web.revoke_token."""
+        from extensions.services.web.pairing import PairingManager
+        from pathlib import Path
+
+        token = params.get("token")
+        if not token:
+            raise ValueError("Missing token parameter")
+
+        # Get pairing manager from app state
+        pairing_manager = self.app.state.pairing_manager
+        if not pairing_manager:
+            # Fallback: create temporary instance
+            pairing_file = Path(__file__).parent / "pairing_codes.jsonl"
+            pairing_manager = PairingManager(str(pairing_file))
+
+        # Verify token exists
+        token_info = pairing_manager.verify_token(token)
+        if not token_info:
+            raise ValueError("Token not found or already expired")
+
+        # Write a revoked record
+        revoked_record = {
+            "token": token,
+            "status": "revoked",
+            "revoked_at": datetime.now(timezone.utc).isoformat()
+        }
+        pairing_manager._write_code(revoked_record)
+
+        return {"success": True, "message": "Token revoked successfully"}
+
     async def _handle_shutdown(self):
-        """Handle module.shutdown: ack → cleanup → ready → exit."""
+        """Handle module.shutdown: ack → exiting → cleanup → ready → close connections → exit."""
+        print("[web] ========== SHUTDOWN STARTED ==========")
         print("[web] Received module.shutdown")
         self._shutting_down = True
 
-        # Step 1: Send ack
+        # Step 1: Send ack (立即确认收到)
+        print("[web] Sending shutdown ack...")
         await self._publish_event({
             "event": "module.shutdown.ack",
-            "data": {"module_id": "web", "estimated_cleanup": 2},
+            "data": {"module_id": "web"},
         })
         print("[web] shutdown ack sent")
 
-        # Step 2: Cleanup (cancel background tasks)
+        # Step 2: Send module.exiting (开始清理)
+        print("[web] Sending module.exiting...")
+        await self._publish_event({
+            "event": "module.exiting",
+            "data": {
+                "module_id": "web",
+                "type": "passive",
+                "reason": "shutdown_requested",
+                "restart": "auto",
+                "action": "none",
+                "timeout": 3.0,
+                "restart_delay": 0.0,
+            },
+        })
+        print("[web] module.exiting sent")
+
+        # Step 3: Cleanup (cancel background tasks)
+        print("[web] Cleaning up background tasks...")
         if self._test_task:
             self._test_task.cancel()
+            print("[web] Test task cancelled")
         if self.bt_manager:
             await self.bt_manager.stop()
+            print("[web] Bluetooth manager stopped")
+
+        # Step 3: Close all WebSocket connections gracefully
+        print("[web] Closing WebSocket connections...")
+
+        # Close relay sessions
+        if hasattr(self.app.state, 'relay_service'):
+            await self.app.state.relay_service.close_all_sessions()
+
+        # Close management clients
+        from .routes.routes_management_ws import close_all_clients
+        await close_all_clients()
 
-        # Step 3: Send ready (before closing WS!)
+        print("[web] All WebSocket connections closed")
+
+        # Step 4: Send ready (after closing connections)
+        print("[web] Sending shutdown ready...")
         await self._publish_event({
             "event": "module.shutdown.ready",
             "data": {"module_id": "web"},
         })
-        print("[web] Shutdown complete")
+        print("[web] shutdown ready sent")
 
-        # Step 4: Trigger uvicorn exit (WS will close when uvicorn shuts down)
+        # Step 5: Trigger uvicorn graceful shutdown
+        print("[web] Triggering uvicorn graceful shutdown...")
         if self._uvicorn_server:
             self._uvicorn_server.should_exit = True
+            print("[web] uvicorn.should_exit = True")
+
+        print("[web] ========== SHUTDOWN COMPLETE ==========")
+
+        # Give uvicorn a moment to start shutdown, then force exit
+        # This prevents hanging on lingering connections
+        await asyncio.sleep(0.5)
+        import sys
+        sys.exit(0)
 
     async def _publish_event(self, event: dict):
         """Publish an event via RPC event.publish."""
         if not self._ws:
+            print(f"[web] WARNING: Cannot publish event {event.get('event')}, WebSocket not connected")
             return
-        await self._rpc_call(self._ws, "event.publish", {
-            "event_id": str(uuid.uuid4()),
-            "event": event.get("event", ""),
-            "data": event.get("data", {}),
-        })
+        try:
+            await self._rpc_call(self._ws, "event.publish", {
+                "event_id": str(uuid.uuid4()),
+                "event": event.get("event", ""),
+                "data": event.get("data", {}),
+            })
+        except Exception as e:
+            print(f"[web] ERROR: Failed to publish event {event.get('event')}: {e}")
 
     # ── Test event loop ──