@agentunion/kite 1.4.0 → 1.6.0

package/extensions/channels/acp_channel/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 AcpChannelServer:
 
-    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._pending_rpc = {}  # Track pending RPC calls: {id: future}
+        self._extra_ws: dict = {}           # slot → WebSocket
+        self._extra_ws_tasks: dict = {}     # slot → recv loop Task
+        self._has_registered = False
+        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 AcpChannelServer:
                     if code in (4001, 4003):
                         print(f"[acp_channel] Kernel 认证失败 (code {code})，退出")
                         self._exit_code = 1
+                        self._auth_failed = True
                         self._shutting_down = True
                         return
                 if attempt >= max_retries:
@@ -80,31 +96,111 @@ class AcpChannelServer:
                     return
                 print(f"[acp_channel] Kernel connection error: {e}, retrying in {retry_delay:.1f}s ({attempt}/{max_retries})")
             self._ws = None
+            self._cancel_pending_rpc()
             if self._shutting_down:
                 return
             await asyncio.sleep(retry_delay)
             retry_delay = min(retry_delay * 2, max_delay)
 
     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=acp_channel"
+        """Single WebSocket session: connect, auth, start receiver, subscribe, register, ready."""
+        url = f"ws://127.0.0.1:{self.kernel_port}/ws?id=acp_channel"
         print(f"[acp_channel] 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"[acp_channel] Connected to Kernel{elapsed_str}")
 
-            # Step 1: Subscribe to events (先订阅)
+            # 启动接收循环（先于 RPC 调用，防止 wait_response 死锁）
+            recv_task = asyncio.create_task(self._ws_receiver(ws))
+
+            try:
+                # subscribe/register/module.ready 由 _do_init 在收到 system.require_init 后执行
+
+                # 等待接收循环结束（连接断开时自然退出）
+                await recv_task
+            finally:
+                recv_task.cancel()
+                self._cancel_pending_rpc()
+
+    async def _ws_receiver(self, ws):
+        """Receive loop: dispatch RPC requests, event notifications, and RPC responses.
+
+        CRITICAL: RPC 死锁防范
+        - 入站 RPC 请求必须用 create_task() 异步执行，不可 await
+        - 原因：如果 handler 内部调用 rpc_call() 发出站请求，出站响应需要本接收循环来分发
+        - 如果接收循环被 await handler 阻塞，出站响应永远收不到 → 超时死锁
+        """
+        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 has_id:
+                        asyncio.create_task(self._handle_rpc_request(ws, msg))
+                    elif has_method and not has_id:
+                        asyncio.create_task(self._handle_event_notification(msg))
+                    elif not has_method and has_id:
+                        self._handle_rpc_response(msg)
+                except Exception as e:
+                    print(f"[acp_channel] 消息处理异常（已忽略）: {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", "phone_channel.test"],
             })
 
-            # Step 2: Register to Kernel (再注册)
+            # Step 2: Register to Kernel
             await self._rpc_call(ws, "registry.register", {
                 "module_id": "acp_channel",
                 "module_type": "channel",
                 "name": "ACP Channel",
+                "tools": {
+                    "rpc": {
+                        "module": {
+                            "echo": {"method": "echo", "description": "Echo test"},
+                            "get_status": {"method": "get_status", "description": "Get channel status"},
+                            "nested_call": {"method": "nested_call", "description": "Nested RPC call test"},
+                        }
+                    }
+                },
                 "events_publish": {
                     "acp_channel": {
                         "test": {},
@@ -114,150 +210,361 @@ class AcpChannelServer:
                     "module.started",
                     "module.stopped",
                     "module.shutdown",
+                    "phone_channel.test",
                 ],
             })
 
-            # 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": "acp_channel",
-                        "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": "acp_channel",
+                    "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"[acp_channel] module.ready published{elapsed_str}")
+                elapsed_str = self._fmt_elapsed(self.boot_t0)
+                print(f"[acp_channel] module.ready published (reason={reason}, {elapsed_str})")
+                self._has_registered = True
+        except Exception as e:
+            print(f"[acp_channel] _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
+
+        # 弹性连接 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
 
-            # Receive loop
-            # CRITICAL: RPC 死锁防范
-            # - 入站 RPC 请求必须用 create_task() 异步执行，不可 await
-            # - 原因：如果 handler 内部调用 rpc_call() 发出站请求，出站响应需要本接收循环来分发
-            # - 如果接收循环被 await handler 阻塞，出站响应永远收不到 → 超时死锁
-            # - 事件通知和 RPC 响应可以同步处理（它们不会反向调用 rpc_call）
+        # Handle phone_channel.test event (跨模块事件测试)
+        if event_type == "phone_channel.test":
+            print(f"[acp_channel] 收到 phone_channel.test 事件: {data.get('message', '')}")
+            return
+
+        # Layer 1: 处理订阅的事件
+        if event_type == "module.shutdown":
+            target = data.get("module_id", "")
+            reason = data.get("reason", "")
+            # Handle both targeted shutdown (module_id == "acp_channel") and broadcast shutdown (no module_id or launcher_lost)
+            if target.startswith("acp_channel") 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"[acp_channel] 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"[acp_channel] Slot {slot} connected")
+        except Exception as e:
+            print(f"[acp_channel] Slot {slot} connect failed: {e}")
 
+    async def _slot_recv_loop(self, slot, ws):
+        """附加连接的接收循环：与主连接平等处理所有消息。"""
+        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
-
-                    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 == "acp_channel") and broadcast shutdown (no module_id or launcher_lost)
-                            if target == "acp_channel" 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"[acp_channel] Debug: Unhandled event: {event_type}")
-
+                    if has_method and has_id:
+                        asyncio.create_task(self._handle_rpc_request(ws, msg))
+                    elif has_method and not has_id:
+                        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"[acp_channel] 事件处理异常（已忽略）: {e}")
+                    print(f"[acp_channel] Slot {slot} 消息处理异常（已忽略）: {e}")
+        except Exception:
+            pass
+        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("[acp_channel] 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": "acp_channel", "estimated_cleanup": 2},
+        await self._publish_event(ws, "module.shutdown.ack", {
+            "module_id": "acp_channel", "estimated_cleanup": 2,
         })
         print("[acp_channel] 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": "acp_channel", "action": "none"},
+        await self._publish_event(ws, "module.exiting", {
+            "module_id": "acp_channel", "action": "none",
+        })
+
+        # Step 3: Send ready（必须在关闭连接之前发送）
+        await self._publish_event(ws, "module.shutdown.ready", {
+            "module_id": "acp_channel",
         })
+        print("[acp_channel] Shutdown ready sent")
+
+        # Step 4: 等待 Kernel 处理 shutdown.ready
+        await asyncio.sleep(0.01)
+
+        # 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()
 
-        # 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": "acp_channel"},
-        })
-        print("[acp_channel] Shutdown ready sent")
+        # Step 6: Close WebSocket connection gracefully
+        try:
+            await ws.close(code=1000, reason="Graceful shutdown")
+            print("[acp_channel] WebSocket closed")
+        except Exception as e:
+            print(f"[acp_channel] Failed to close WebSocket: {e}")
 
-        # Step 4: Exit process
-        import sys
-        sys.exit(0)
+    async def _handle_rpc_request(self, ws, msg: dict):
+        """Handle incoming RPC request."""
+        method = msg.get("method", "")
+        params = msg.get("params", {})
+        req_id = msg.get("id")
 
-    async def _rpc_call(self, ws, method: str, params: dict = None):
+        self._active_handlers += 1
+        try:
+            # 支持带前缀和不带前缀的方法名
+            bare_method = method.split(".")[-1] if "." in method else method
+
+            if bare_method == "echo":
+                result = {"echo": params.get("message", ""), "timestamp": time.time()}
+            elif bare_method == "get_status":
+                result = {
+                    "status": "running",
+                    "uptime": time.time() - self._start_time,
+                    "module_id": "acp_channel",
+                }
+            elif bare_method == "nested_call":
+                # 嵌套调用：调用 phone_channel.echo
+                message = params.get("message", "")
+                print(f"[acp_channel] 嵌套调用: 准备调用 phone_channel.echo")
+                nested_result = await self._call_remote_rpc(
+                    "phone_channel",
+                    "echo",
+                    {"message": f"nested from acp: {message}"}
+                )
+                result = {
+                    "module_id": "acp_channel",
+                    "nested_call_to": "phone_channel.echo",
+                    "nested_result": nested_result,
+                    "original_message": message
+                }
+            else:
+                raise Exception(f"Unknown method: {method}")
+
+            response = {"jsonrpc": "2.0", "id": req_id, "result": result}
+            await ws.send(json.dumps(response))
+            print(f"[acp_channel] RPC 请求处理: {method}")
+        except Exception as e:
+            error_response = {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "error": {"code": -32603, "message": str(e)},
+            }
+            await ws.send(json.dumps(error_response))
+            print(f"[acp_channel] RPC 请求失败: {method} - {e}")
+        finally:
+            self._active_handlers -= 1
+
+    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):
+        """Handle RPC response (resolve pending future with full msg)."""
+        rpc_id = msg.get("id")
+        future = self._pending_rpc.pop(rpc_id, None)
+        if future and not future.done():
+            future.set_result(msg)
+
+    def _cancel_pending_rpc(self):
+        """Cancel all pending RPC futures on disconnect."""
+        for rpc_id, future in self._pending_rpc.items():
+            if not future.done():
+                future.cancel()
+        self._pending_rpc.clear()
+
+    async def _rpc_call(self, ws, method: str, params: dict = None,
+                        wait_response: bool = True, timeout: float = 3.0) -> dict:
         """Send a JSON-RPC 2.0 request."""
-        msg = {"jsonrpc": "2.0", "id": str(uuid.uuid4()), "method": method}
+        req_id = str(uuid.uuid4())
+        msg = {"jsonrpc": "2.0", "id": req_id, "method": method}
         if params:
             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:
+        if wait_response:
+            future = asyncio.Future()
+            self._pending_rpc[req_id] = future
+            await ws.send(json.dumps(msg))
             try:
-                await asyncio.sleep(30)
-                if not self._shutting_down:
-                    await self._rpc_call(ws, "registry.heartbeat", {"module_id": "acp_channel"})
-            except Exception as e:
-                print(f"[acp_channel] Heartbeat error: {e}")
-                break
+                result = await asyncio.wait_for(future, timeout=timeout)
+                return result
+            except asyncio.TimeoutError:
+                self._pending_rpc.pop(req_id, None)
+                raise Exception(f"RPC timeout: {method}")
+        else:
+            await ws.send(json.dumps(msg))
+
+    async def _handle_ping_event(self, data: dict):
+        """Handle system.ping event and reply with system.pong."""
+        t1 = data.get("ping_time")
+        t2 = time.time()
+
+        await self._publish_event(self._ws, "system.pong", {
+            "module_id": "acp_channel",
+            "ping_time": t1,
+            "pong_time": t2,
+        })
+
+    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)
 
-    async def _publish_event(self, event: dict):
-        """Publish an event via JSON-RPC event.publish."""
+    async def _call_remote_rpc(self, target_module: str, method: str, params: dict = None):
+        """Call RPC method on another module via Kernel."""
         if not self._ws:
-            return
+            return None
         try:
-            await self._rpc_call(self._ws, "event.publish", {
-                "event_id": str(uuid.uuid4()),
-                "event": event.get("event", ""),
-                "data": event.get("data", {}),
-            })
+            result = await self._rpc_call(
+                self._ws,
+                f"{target_module}.{method}",
+                params or {},
+                wait_response=True
+            )
+            if result and "error" in result:
+                err = result["error"]
+                print(f"[acp_channel] 跨模块 RPC 调用失败 ({target_module}.{method}): {err.get('message', 'RPC error')}")
+                return None
+            return result.get("result") if result else None
         except Exception as e:
-            print(f"[acp_channel] Failed to publish event: {e}")
+            print(f"[acp_channel] 跨模块 RPC 调用失败 ({target_module}.{method}): {e}")
+            return None
 
     # ── 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": "acp_channel.test",
-                "data": {
-                    "message": "test event from acp_channel",
+        """Publish test events and call cross-module RPC every 12 seconds."""
+        try:
+          await asyncio.sleep(3)  # Wait for initial setup
+          counter = 0
+          while not self._shutting_down:
+            counter += 1
+
+            # Test 1: Publish event
+            if self._ws:
+                await self._publish_event(self._ws, "acp_channel.test", {
+                    "message": f"test event from acp_channel #{counter}",
                     "timestamp": datetime.now(timezone.utc).isoformat(),
-                },
-            })
-            print("[acp_channel] test event published")
+                })
+
+            # Test 2: Call phone_channel RPC (every other cycle)
+            if counter % 2 == 1:
+                print(f"[acp_channel] 测试跨模块 RPC 调用 phone_channel.get_status")
+                result = await self._call_remote_rpc(
+                    "phone_channel",
+                    "get_status",
+                    {}
+                )
+                if result:
+                    print(f"[acp_channel] phone_channel.get_status 返回: {result}")
+
+            # 上报队列压力（供 Kernel 自动缩容决策）
+            if self._ws:
+                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,
+                })
+
+            await asyncio.sleep(12)
+        except Exception:
+            pass

package/extensions/channels/phone_channel/__init__.py

@@ -0,0 +1 @@
+