@agentunion/kite 1.0.7 → 1.3.0

package/extensions/channels/acp_channel/server.py

@@ -0,0 +1,197 @@
+"""
+ACP Channel WebSocket client.
+Connects to Kernel via WebSocket JSON-RPC 2.0 for event publishing and subscription.
+"""
+
+import asyncio
+import json
+import time
+import uuid
+from datetime import datetime, timezone
+
+import websockets
+
+
+class AcpChannelServer:
+
+    def __init__(self, token: str = "", kernel_port: int = 0, boot_t0: float = 0):
+        self.token = token
+        self.kernel_port = kernel_port
+        self.boot_t0 = boot_t0
+        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._start_time = time.time()
+
+    async def run(self):
+        """Main entry point: start WebSocket loop and test event loop."""
+        if self.kernel_port:
+            self._ws_task = asyncio.create_task(self._ws_loop())
+        self._test_task = asyncio.create_task(self._test_event_loop())
+
+        # Wait for tasks to complete
+        tasks = [t for t in [self._ws_task, self._test_task] if t]
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+        print("[acp_channel] Shutdown complete")
+
+    # ── Kernel WebSocket client ──
+
+    async def _ws_loop(self):
+        """Connect to Kernel, subscribe, register, and listen. Reconnect on failure."""
+        retry_delay = 0.5
+        max_delay = 30
+        while not self._shutting_down:
+            try:
+                await self._ws_connect()
+            except asyncio.CancelledError:
+                return
+            except Exception as e:
+                print(f"[acp_channel] Kernel connection error: {e}, retrying in {retry_delay:.1f}s")
+            self._ws = None
+            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"
+        print(f"[acp_channel] Connecting to Kernel (port {self.kernel_port})")
+        async with websockets.connect(url, open_timeout=3, 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"[acp_channel] Connected to Kernel{elapsed_str}")
+
+            # Step 1: Subscribe to events (先订阅)
+            await self._rpc_call(ws, "event.subscribe", {
+                "events": ["module.started", "module.stopped", "module.shutdown"],
+            })
+
+            # Step 2: Register to Kernel (再注册)
+            await self._rpc_call(ws, "registry.register", {
+                "module_id": "acp_channel",
+                "module_type": "channel",
+                "name": "ACP Channel",
+                "events_publish": {
+                    "acp_channel.test": {},
+                },
+                "events_subscribe": [
+                    "module.started",
+                    "module.stopped",
+                    "module.shutdown",
+                ],
+            })
+
+            # Step 3: Publish module.ready (once)
+            if not self._ready_sent:
+                await self._rpc_call(ws, "event.publish", {
+                    "event_id": str(uuid.uuid4()),
+                    "event": "module.ready",
+                    "data": {
+                        "module_id": "acp_channel",
+                        "graceful_shutdown": True,
+                    },
+                })
+                self._ready_sent = True
+                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}")
+
+            # Reset retry delay on successful connection
+            retry_delay = 0.5
+
+            # Receive loop
+            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_name = params.get("event", "")
+                        if event_name == "module.shutdown":
+                            data = params.get("data", {})
+                            target = data.get("module_id", "")
+                            if target == "acp_channel":
+                                await self._handle_shutdown(ws)
+                                return
+                    elif not has_method and has_id:
+                        # JSON-RPC Response (to our RPC calls)
+                        pass
+                except Exception as e:
+                    print(f"[acp_channel] 事件处理异常（已忽略）: {e}")
+
+    async def _handle_shutdown(self, ws):
+        """Handle module.shutdown: ack → cleanup → ready → 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},
+        })
+        print("[acp_channel] shutdown ack sent")
+
+        # Step 2: Cleanup (cancel background tasks)
+        if self._test_task:
+            self._test_task.cancel()
+
+        # Step 3: Send ready (before closing WS!)
+        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 4: Exit process
+        import sys
+        sys.exit(0)
+
+    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}
+        if params:
+            msg["params"] = params
+        await ws.send(json.dumps(msg))
+
+    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"[acp_channel] Failed to publish event: {e}")
+
+    # ── 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",
+                    "timestamp": datetime.now(timezone.utc).isoformat(),
+                },
+            })
+            print("[acp_channel] test event published")