@agentunion/kite 1.4.0 → 1.6.0

package/extensions/agents/assistant/server.py

@@ -19,21 +19,36 @@ SYSTEM_BROADCAST_EVENTS = {
     "module.crashed", "module.exiting", "module.offline",
     "module.shutdown.ack", "module.shutdown.ready",
     "system.ready", "registry.updated",
+    "system.instance.started", "system.instance.stopped",
 }
 
 
 class AssistantServer:
 
-    def __init__(self, token: str = "", kernel_port: int = 0, boot_t0: float = 0):
+    def __init__(self, token: str = "", kernel_port: int = 0, boot_t0: float = 0, max_connections: int = 1,
+                 max_instances: int = 1, instance_elastic: bool = False,
+                 gateway_url: str = "", kite_token: str = ""):
         self.token = token
         self.kernel_port = kernel_port
         self.boot_t0 = boot_t0
+        self.max_connections = max_connections
+        self.max_instances = max_instances
+        self.instance_elastic = instance_elastic
+        self.gateway_url = gateway_url
+        self.kite_token = kite_token
         self._ws_task: asyncio.Task | None = None
         self._test_task: asyncio.Task | None = None
         self._ws: object | None = None
         self._shutting_down = False
         self._exit_code = 0  # Exit code for main() to use
+        self._auth_failed = False
         self._start_time = time.time()
+        self._extra_ws: dict = {}           # slot → WebSocket
+        self._extra_ws_tasks: dict = {}     # slot → recv loop Task
+        self._has_registered = False
+        self._pending_rpc: dict[str, asyncio.Future] = {}
+        self._active_handlers = 0   # 当前正在处理的入站请求数
+        self._handler_capacity = 50 # 并发处理能力上限（用于压力计算）
 
     async def run(self):
         """Main entry point: start WebSocket loop and test event loop."""
@@ -71,6 +86,7 @@ class AssistantServer:
                     if code in (4001, 4003):
                         print(f"[assistant] Kernel 认证失败 (code {code})，退出")
                         self._exit_code = 1
+                        self._auth_failed = True
                         self._shutting_down = True
                         return
                 if attempt >= max_retries:
@@ -87,20 +103,98 @@ class AssistantServer:
 
     async def _ws_connect(self):
         """Single WebSocket session: connect, subscribe, register, ready, receive loop."""
-        url = f"ws://127.0.0.1:{self.kernel_port}/ws?token={self.token}&id=assistant"
+        url = f"ws://127.0.0.1:{self.kernel_port}/ws?id=assistant"
         print(f"[assistant] Connecting to Kernel (port {self.kernel_port})")
-        async with websockets.connect(url, open_timeout=3, ping_interval=20, ping_timeout=20, close_timeout=10) as ws:
+        async with websockets.connect(url, open_timeout=3, ping_interval=None, close_timeout=10) as ws:
+            # Send auth message first
+            auth_req = {
+                "jsonrpc": "2.0",
+                "id": "auth",
+                "method": "auth",
+                "params": {"token": self.token}
+            }
+            await ws.send(json.dumps(auth_req))
+
+            # Wait for auth response
+            auth_resp_raw = await asyncio.wait_for(ws.recv(), timeout=5)
+            auth_resp = json.loads(auth_resp_raw)
+            if "error" in auth_resp:
+                raise Exception(f"Auth failed: {auth_resp['error']}")
+
             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"[assistant] Connected to Kernel{elapsed_str}")
 
-            # Step 1: Subscribe to events (先订阅)
+            # Start receiver task before subscribe/register so responses can be routed
+            receiver_task = asyncio.create_task(self._ws_receiver(ws))
+
+            try:
+                # subscribe/register/module.ready 由 _do_init 在收到 system.require_init 后执行
+
+                # Wait for receiver to finish (connection closed or error)
+                await receiver_task
+            finally:
+                receiver_task.cancel()
+                for fut in self._pending_rpc.values():
+                    if not fut.done():
+                        fut.set_result({"error": {"code": -32001, "message": "Connection lost"}})
+                self._pending_rpc.clear()
+
+    async def _ws_receiver(self, ws):
+        """Receive loop for main WebSocket connection.
+
+        CRITICAL: RPC 死锁防范
+        - 入站 RPC 请求必须用 create_task() 异步执行，不可 await
+        - 原因：如果 handler 内部调用 rpc_call() 发出站请求，出站响应需要本接收循环来分发
+        - 如果接收循环被 await handler 阻塞，出站响应永远收不到 → 超时死锁
+        - 事件通知和 RPC 响应可以同步处理（它们不会反向调用 rpc_call）
+        """
+        try:
+            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
+
+                    # 检测 system.require_init 事件，触发注册流程
+                    if has_method and not has_id:
+                        params = msg.get("params", {})
+                        event = params.get("event", "")
+                        if event == "system.require_init":
+                            asyncio.create_task(self._do_init(ws))
+                            continue
+
+                    if has_method and not has_id:
+                        # JSON-RPC Notification (event delivery)
+                        # Run in background to prevent deadlock if handler calls RPC
+                        asyncio.create_task(self._handle_event_notification(msg))
+
+                    elif not has_method and has_id:
+                        # JSON-RPC Response (to our RPC calls)
+                        self._handle_rpc_response(msg)
+                except Exception as e:
+                    print(f"[assistant] 事件处理异常: {e}")
+        except websockets.exceptions.ConnectionClosed:
+            pass  # Kernel 连接关闭，正常退出接收循环
+
+    async def _do_init(self, ws):
+        """收到 system.require_init 后执行：订阅 + 注册 + module.ready"""
+        try:
+            reason = "startup" if not self._has_registered else "recovery"
+
+            # Step 1: Subscribe to events
             await self._rpc_call(ws, "event.subscribe", {
-                "events": ["module.started", "module.stopped", "module.shutdown"],
+                "events": [
+                    "module.started", "module.stopped", "module.shutdown",
+                ],
             })
 
-            # Step 2: Register to Kernel (再注册)
+            # Step 2: Register to Kernel
             await self._rpc_call(ws, "registry.register", {
                 "module_id": "assistant",
                 "module_type": "agent",
@@ -117,30 +211,88 @@ class AssistantServer:
                 ],
             })
 
-            # Step 3: Publish module.ready (every reconnect)
+            # Step 3: Publish module.ready
             if not self._shutting_down:
-                await self._rpc_call(ws, "event.publish", {
-                    "event_id": str(uuid.uuid4()),
-                    "event": "module.ready",
-                    "data": {
-                        "module_id": "assistant",
-                        "graceful_shutdown": True,
-                    },
+                startup_time = time.monotonic() - self.boot_t0 if self.boot_t0 else 0
+                await self._publish_event(ws, "module.ready", {
+                    "module_id": "assistant",
+                    "graceful_shutdown": True,
+                    "startup_time": startup_time,
+                    "reason": reason,
                 })
-                elapsed = time.monotonic() - self.boot_t0 if self.boot_t0 else 0
-                elapsed_str = f" ({elapsed:.1f}s)" if elapsed else ""
-                print(f"[assistant] module.ready published{elapsed_str}")
+                elapsed_str = self._fmt_elapsed(self.boot_t0)
+                print(f"[assistant] module.ready published (reason={reason}, {elapsed_str})")
+                self._has_registered = True
+        except Exception as e:
+            print(f"[assistant] _do_init failed: {e}")
+            self._has_registered = False
+
+    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", {})
+
+        # Handle system.ping event
+        if event_type == "system.ping":
+            await self._handle_ping_event(data)
+            return
 
-            # Receive loop
-            # CRITICAL: RPC 死锁防范
-            # - 入站 RPC 请求必须用 create_task() 异步执行，不可 await
-            # - 原因：如果 handler 内部调用 rpc_call() 发出站请求，出站响应需要本接收循环来分发
-            # - 如果接收循环被 await handler 阻塞，出站响应永远收不到 → 超时死锁
-            # - 事件通知和 RPC 响应可以同步处理（它们不会反向调用 rpc_call）
+        # 弹性连接 offer/release
+        if event_type == "system.connection.offer":
+            asyncio.create_task(self._handle_connection_offer(data))
+            return
+        if event_type == "system.connection.release":
+            asyncio.create_task(self._handle_connection_release(data))
+            return
+
+        # Layer 1: 处理订阅的事件
+        if event_type == "module.shutdown":
+            target = data.get("module_id", "")
+            reason = data.get("reason", "")
+            # Handle both targeted shutdown (module_id == "assistant") and broadcast shutdown (no module_id or launcher_lost)
+            if target.startswith("assistant") or not target or reason == "launcher_lost":
+                await self._handle_shutdown(self._ws)
+                return
+
+        # Layer 2: 忽略系统广播事件
+        if event_type in SYSTEM_BROADCAST_EVENTS:
+            return
 
-            # Start heartbeat loop
-            heartbeat_task = asyncio.create_task(self._heartbeat_loop(ws))
+        # Layer 3: 警告未知事件（仅开发环境）
+        if os.environ.get("KITE_ENV") == "development":
+            print(f"[assistant] Debug: Unhandled event: {event_type}")
+
+    async def _handle_connection_offer(self, data):
+        """处理 Kernel 下发的 slot token，建立附加连接。"""
+        slots = data.get("slots", {})
+        for slot_str, info in slots.items():
+            slot = int(slot_str)
+            token = info.get("token", "")
+            if not token or slot in self._extra_ws:
+                continue
+            asyncio.create_task(self._connect_slot(slot, token))
+
+    async def _connect_slot(self, slot, token):
+        """建立单个 slot 附加连接。"""
+        ws_url = f"ws://127.0.0.1:{self.kernel_port}/ws"
+        try:
+            ws = await websockets.connect(ws_url, open_timeout=5, ping_interval=None, close_timeout=5)
+            auth_req = {"jsonrpc": "2.0", "id": f"auth-slot-{slot}", "method": "auth", "params": {"token": token}}
+            await ws.send(json.dumps(auth_req))
+            resp = json.loads(await asyncio.wait_for(ws.recv(), timeout=5))
+            if "error" in resp:
+                await ws.close()
+                return
+            self._extra_ws[slot] = ws
+            self._extra_ws_tasks[slot] = asyncio.create_task(self._slot_recv_loop(slot, ws))
+            print(f"[assistant] Slot {slot} connected")
+        except Exception as e:
+            print(f"[assistant] Slot {slot} connect failed: {e}")
 
+    async def _slot_recv_loop(self, slot, ws):
+        """附加 slot 接收循环：与主连接平等处理所有消息。"""
+        try:
             async for raw in ws:
                 try:
                     msg = json.loads(raw)
@@ -152,112 +304,162 @@ class AssistantServer:
                     has_id = "id" in msg
 
                     if has_method and not has_id:
-                        # JSON-RPC Notification (event delivery)
-                        params = msg.get("params", {})
-                        event_type = params.get("event", "")
-                        data = params.get("data", {})
-
-                        # Layer 1: 处理订阅的事件
-                        if event_type == "module.shutdown":
-                            target = data.get("module_id", "")
-                            reason = data.get("reason", "")
-                            # Handle both targeted shutdown (module_id == "assistant") and broadcast shutdown (no module_id or launcher_lost)
-                            if target == "assistant" or not target or reason == "launcher_lost":
-                                await self._handle_shutdown(ws)
-                                return
-
-                        # Layer 2: 忽略系统广播事件
-                        if event_type in SYSTEM_BROADCAST_EVENTS:
-                            continue
-
-                        # Layer 3: 警告未知事件（仅开发环境）
-                        if os.environ.get("KITE_ENV") == "development":
-                            print(f"[assistant] Debug: Unhandled event: {event_type}")
-
+                        asyncio.create_task(self._handle_event_notification(msg))
                     elif not has_method and has_id:
-                        # JSON-RPC Response (to our RPC calls)
-                        pass
+                        self._handle_rpc_response(msg)
                 except Exception as e:
-                    print(f"[assistant] 事件处理异常: {e}")
+                    print(f"[assistant] Slot {slot} 消息处理异常（已忽略）: {e}")
+        except Exception as e:
+            print(f"[assistant] Slot {slot} recv loop exited: {e}")
+        finally:
+            self._extra_ws.pop(slot, None)
+            self._extra_ws_tasks.pop(slot, None)
+
+    async def _handle_connection_release(self, data):
+        """Kernel 请求释放 slot，优雅关闭。"""
+        for slot in data.get("slots", []):
+            ws = self._extra_ws.pop(slot, None)
+            task = self._extra_ws_tasks.pop(slot, None)
+            if ws:
+                try:
+                    await ws.close(code=1000, reason="release")
+                except Exception:
+                    pass
+            if task:
+                task.cancel()
 
     async def _handle_shutdown(self, ws):
-        """Handle module.shutdown: ack → exiting → cleanup → ready → exit."""
+        """Handle module.shutdown: ack → exiting → ready → cleanup → exit."""
         print("[assistant] Received module.shutdown")
         self._shutting_down = True
 
         # Step 1: Send ack (立即确认收到)
-        await self._rpc_call(ws, "event.publish", {
-            "event_id": str(uuid.uuid4()),
-            "event": "module.shutdown.ack",
-            "data": {"module_id": "assistant", "estimated_cleanup": 2},
+        await self._publish_event(ws, "module.shutdown.ack", {
+            "module_id": "assistant", "estimated_cleanup": 2,
         })
         print("[assistant] shutdown ack sent")
 
         # Step 2: Send module.exiting (开始清理)
-        await self._rpc_call(ws, "event.publish", {
-            "event_id": str(uuid.uuid4()),
-            "event": "module.exiting",
-            "data": {"module_id": "assistant", "action": "none"},
+        await self._publish_event(ws, "module.exiting", {
+            "module_id": "assistant", "action": "none",
         })
 
-        # Step 3: Cleanup (cancel background tasks)
-        if self._test_task:
-            self._test_task.cancel()
-
-        # Step 4: Send ready (清理完成)
-        await self._rpc_call(ws, "event.publish", {
-            "event_id": str(uuid.uuid4()),
-            "event": "module.shutdown.ready",
-            "data": {"module_id": "assistant"},
+        # Step 3: Send ready（必须在关闭连接之前发送）
+        await self._publish_event(ws, "module.shutdown.ready", {
+            "module_id": "assistant",
         })
         print("[assistant] Shutdown ready sent")
 
-        # Step 4: Exit process
-        import sys
-        sys.exit(0)
+        # Step 4: 等待 Kernel 处理 shutdown.ready
+        await asyncio.sleep(0.01)
 
-    async def _rpc_call(self, ws, method: str, params: dict = None):
-        """Send a JSON-RPC 2.0 request."""
-        msg = {"jsonrpc": "2.0", "id": str(uuid.uuid4()), "method": method}
+        # Step 5: Cleanup - 关闭所有附加连接
+        for _s, _w in list(self._extra_ws.items()):
+            try:
+                await _w.close(code=1000, reason="shutdown")
+            except Exception:
+                pass
+        for _t in self._extra_ws_tasks.values():
+            _t.cancel()
+        self._extra_ws.clear()
+        self._extra_ws_tasks.clear()
+
+        if self._test_task:
+            self._test_task.cancel()
+
+        # Step 6: Close WebSocket connection gracefully
+        try:
+            await ws.close(code=1000, reason="Graceful shutdown")
+            print("[assistant] WebSocket closed")
+        except Exception as e:
+            print(f"[assistant] Failed to close WebSocket: {e}")
+
+    async def _rpc_call(self, ws, method: str, params: dict = None,
+                        wait_response: bool = True, timeout: float = 3.0) -> dict:
+        """JSON-RPC 2.0 request。默认等待响应。"""
+        rpc_id = str(uuid.uuid4())
+        msg = {"jsonrpc": "2.0", "id": rpc_id, "method": method}
         if params:
             msg["params"] = params
+        if not wait_response:
+            await ws.send(json.dumps(msg))
+            return
+        future = asyncio.get_event_loop().create_future()
+        self._pending_rpc[rpc_id] = future
         await ws.send(json.dumps(msg))
+        try:
+            return await asyncio.wait_for(future, timeout=timeout)
+        except asyncio.TimeoutError:
+            self._pending_rpc.pop(rpc_id, None)
+            return {"error": {"code": -32000, "message": f"RPC timeout: {method} ({timeout}s)"}}
+
+    def _compute_pressure_level(self) -> str:
+        """根据并发处理量计算压力等级。"""
+        depth = self._active_handlers + len(self._pending_rpc)
+        ratio = depth / self._handler_capacity if self._handler_capacity > 0 else 0
+        if ratio >= 0.9:
+            return "critical"
+        elif ratio >= 0.5:
+            return "overload"
+        elif depth > 0:
+            return "normal"
+        return "idle"
+
+    def _handle_rpc_response(self, msg: dict):
+        """Resolve pending RPC future with the response."""
+        rpc_id = msg.get("id")
+        future = self._pending_rpc.pop(rpc_id, None)
+        if future and not future.done():
+            future.set_result(msg)
+
+    async def _handle_ping_event(self, data: dict):
+        """Handle system.ping event and reply with system.pong."""
+        if self._ws is None:
+            return
+        t1 = data.get("ping_time")
+        t2 = time.time()
 
-    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": "assistant"})
-            except Exception as e:
-                print(f"[assistant] Heartbeat error: {e}")
-                break
+        await self._publish_event(self._ws, "system.pong", {
+            "module_id": "assistant",
+            "ping_time": t1,
+            "pong_time": t2,
+        })
 
-    async def _publish_event(self, event: dict):
-        """Publish an event via JSON-RPC event.publish."""
-        if not self._ws:
-            return
-        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"[assistant] Failed to publish event: {e}")
+    def _fmt_elapsed(self, t0: float) -> str:
+        """Format elapsed time since t0."""
+        d = time.monotonic() - t0 if t0 else 0
+        if d < 1:
+            return f"{d * 1000:.0f}ms"
+        if d < 10:
+            return f"{d:.1f}s"
+        return f"{d:.0f}s"
+
+    async def _publish_event(self, ws, event: str, data: dict = None):
+        """Publish event via event.publish RPC (fire-and-forget)."""
+        await self._rpc_call(ws, "event.publish", {
+            "event_id": str(uuid.uuid4()),
+            "event": event,
+            "data": data or {},
+        }, wait_response=False)
 
     # ── Test event loop ──
 
     async def _test_event_loop(self):
         """Publish a test event every 10 seconds."""
-        while True:
-            await asyncio.sleep(10)
-            await self._publish_event({
-                "event": "assistant.test",
-                "data": {
-                    "message": "test event from assistant",
-                    "timestamp": datetime.now(timezone.utc).isoformat(),
-                },
-            })
-            print("[assistant] test event published")
+        try:
+            while True:
+                await asyncio.sleep(10)
+                if self._ws:
+                    await self._publish_event(self._ws, "assistant.test", {
+                            "message": "test event from assistant",
+                            "timestamp": datetime.now(timezone.utc).isoformat(),
+                        })
+                    # 上报队列压力（供 Kernel 自动缩容决策）
+                    depth = self._active_handlers + len(self._pending_rpc)
+                    await self._publish_event(self._ws, "module.queue_pressure", {
+                        "level": self._compute_pressure_level(),
+                        "depth": depth,
+                        "capacity": self._handler_capacity,
+                    })
+        except Exception:
+            pass