@agentunion/kite 1.2.0 → 1.3.0

package/extensions/services/model_service/server.py

@@ -1,236 +0,0 @@
-"""
-Model Service HTTP server.
-Exposes /health and /status endpoints.
-Connects to Event Hub via WebSocket for event publishing and subscription.
-Sends periodic heartbeat to Registry and test events to Event Hub.
-"""
-
-import asyncio
-import json
-import time
-import uuid
-from datetime import datetime, timezone
-
-import httpx
-import websockets
-from fastapi import FastAPI
-
-
-class ModelServiceServer:
-
-    def __init__(self, token: str = "", registry_url: str = "",
-                 event_hub_ws: str = "", boot_t0: float = 0):
-        self.token = token
-        self.registry_url = registry_url
-        self.event_hub_ws = event_hub_ws
-        self.boot_t0 = boot_t0
-        self._ws_task: asyncio.Task | None = None
-        self._heartbeat_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._uvicorn_server = None  # set by entry.py for graceful shutdown
-        self._start_time = time.time()
-        self.app = self._create_app()
-
-    def _create_app(self) -> FastAPI:
-        app = FastAPI(title="Kite Model Service", docs_url=None, redoc_url=None)
-        server = self
-
-        @app.on_event("startup")
-        async def _startup():
-            server._heartbeat_task = asyncio.create_task(server._heartbeat_loop())
-            if server.event_hub_ws:
-                server._ws_task = asyncio.create_task(server._ws_loop())
-            server._test_task = asyncio.create_task(server._test_event_loop())
-
-        @app.on_event("shutdown")
-        async def _shutdown():
-            if server._heartbeat_task:
-                server._heartbeat_task.cancel()
-            if server._ws_task:
-                server._ws_task.cancel()
-            if server._test_task:
-                server._test_task.cancel()
-            if server._ws:
-                await server._ws.close()
-            print("[model_service] Shutdown complete")
-
-        @app.get("/health")
-        async def health():
-            return {
-                "status": "healthy",
-                "details": {
-                    "event_hub_connected": server._ws is not None,
-                    "uptime_seconds": round(time.time() - server._start_time),
-                },
-            }
-
-        @app.get("/status")
-        async def status():
-            return {
-                "module": "model_service",
-                "status": "running",
-                "event_hub_connected": server._ws is not None,
-                "uptime_seconds": round(time.time() - server._start_time),
-            }
-
-        return app
-
-    # ── Event Hub WebSocket client ──
-
-    async def _ws_loop(self):
-        """Connect to Event Hub, subscribe, and listen. Reconnect on failure."""
-        retry_delay = 0.5  # start with 0.5s
-        max_delay = 30  # cap at 30s
-        while not self._shutting_down:
-            try:
-                await self._ws_connect()
-            except asyncio.CancelledError:
-                return
-                retry_delay = 0.5  # reset on successful connection
-            except Exception as e:
-                print(f"[model_service] Event Hub 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)  # exponential backoff
-
-    async def _ws_connect(self):
-        """Single WebSocket session: connect, subscribe, receive loop."""
-        url = f"{self.event_hub_ws}?token={self.token}&id=model_service"
-        print(f"[model_service] WS connecting to {self.event_hub_ws}")
-        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"[model_service] Connected to Event Hub{elapsed_str}")
-
-            # Subscribe to module lifecycle events + shutdown
-            await ws.send(json.dumps({
-                "type": "subscribe",
-                "events": ["module.started", "module.stopped", "module.shutdown"],
-            }))
-
-            # Send module.ready (once) so Launcher knows we're up
-            if not self._ready_sent:
-                ready_msg = {
-                    "type": "event",
-                    "event_id": str(uuid.uuid4()),
-                    "event": "module.ready",
-                    "source": "model_service",
-                    "timestamp": datetime.now(timezone.utc).isoformat(),
-                    "data": {
-                        "module_id": "model_service",
-                        "graceful_shutdown": True,
-                    },
-                }
-                await ws.send(json.dumps(ready_msg))
-                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"[model_service] module.ready sent{elapsed_str}")
-
-            # Receive loop
-            async for raw in ws:
-                try:
-                    msg = json.loads(raw)
-                except (json.JSONDecodeError, TypeError):
-                    continue
-
-                try:
-                    msg_type = msg.get("type", "")
-                    if msg_type == "event":
-                        event_name = msg.get("event", "")
-                        if event_name == "module.shutdown":
-                            target = (msg.get("data") if isinstance(msg.get("data"), dict) else {}).get("module_id", "")
-                            if target == "model_service":
-                                await self._handle_shutdown(ws)
-                                return
-                    elif msg_type == "ack":
-                        pass  # publish confirmed
-                    elif msg_type == "error":
-                        print(f"[model_service] Event Hub error: {msg.get('message')}")
-                except Exception as e:
-                    print(f"[model_service] 事件处理异常（已忽略）: {e}")
-
-    async def _handle_shutdown(self, ws):
-        """Handle module.shutdown: ack → cleanup → ready → exit."""
-        print("[model_service] Received module.shutdown")
-        self._shutting_down = True
-
-        # Step 1: Send ack
-        await self._publish_event({
-            "event": "module.shutdown.ack",
-            "data": {"module_id": "model_service", "estimated_cleanup": 2},
-        })
-        print("[model_service] shutdown ack sent")
-
-        # Step 2: Cleanup (cancel background tasks)
-        if self._heartbeat_task:
-            self._heartbeat_task.cancel()
-        if self._test_task:
-            self._test_task.cancel()
-
-        # Step 3: Send ready (before closing WS!)
-        await self._publish_event({
-            "event": "module.shutdown.ready",
-            "data": {"module_id": "model_service"},
-        })
-        print("[model_service] Shutdown complete")
-
-        # Step 4: Trigger uvicorn exit (WS will close when uvicorn shuts down)
-        if self._uvicorn_server:
-            self._uvicorn_server.should_exit = True
-
-    async def _publish_event(self, event: dict):
-        """Publish an event to Event Hub via WebSocket."""
-        if not self._ws:
-            return
-        msg = {
-            "type": "event",
-            "event_id": str(uuid.uuid4()),
-            "event": event.get("event", ""),
-            "source": "model_service",
-            "timestamp": datetime.now(timezone.utc).isoformat(),
-            "data": event.get("data", {}),
-        }
-        try:
-            await self._ws.send(json.dumps(msg))
-        except Exception as e:
-            print(f"[model_service] Failed to publish event: {e}")
-
-    # ── Heartbeat to Registry ──
-
-    async def _heartbeat_loop(self):
-        """Send heartbeat to Registry every 30 seconds."""
-        while True:
-            await asyncio.sleep(30)
-            try:
-                async with httpx.AsyncClient() as client:
-                    await client.post(
-                        f"{self.registry_url}/modules",
-                        json={"action": "heartbeat", "module_id": "model_service"},
-                        headers={"Authorization": f"Bearer {self.token}"},
-                        timeout=5,
-                    )
-                print("[model_service] heartbeat sent")
-            except Exception:
-                pass
-
-    # ── 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": "model_service.test",
-                "data": {
-                    "message": "test event from model_service",
-                    "timestamp": datetime.now(timezone.utc).isoformat(),
-                },
-            })
-            print("[model_service] test event published")

package/extensions/services/watchdog/server.py

@@ -1,229 +0,0 @@
-"""
-Watchdog HTTP server.
-Exposes /health and /status endpoints. Runs the monitor loop on startup.
-Connects to Event Hub via WebSocket for event publishing and subscription.
-"""
-
-import asyncio
-import json
-import time
-import uuid
-
-import httpx
-import websockets
-from fastapi import FastAPI
-
-from .monitor import HealthMonitor
-
-
-class WatchdogServer:
-
-    def __init__(self, monitor: HealthMonitor, token: str = "",
-                 event_hub_ws: str = ""):
-        self.monitor = monitor
-        self.token = token
-        self.event_hub_ws = event_hub_ws
-        self._monitor_task: asyncio.Task | None = None
-        self._ws_task: asyncio.Task | None = None
-        self._heartbeat_task: asyncio.Task | None = None
-        self._ws: object | None = None
-        self._ready_sent = False
-        self._start_time = time.time()
-        self._uvicorn_server = None  # set by entry.py for graceful shutdown
-        self._shutting_down = False
-        self.app = self._create_app()
-
-        # Wire up publish callback on monitor
-        self.monitor.publish_event = self._publish_event
-
-    def _create_app(self) -> FastAPI:
-        app = FastAPI(title="Kite Watchdog", docs_url=None, redoc_url=None)
-        server = self
-
-        @app.on_event("startup")
-        async def _startup():
-            server._monitor_task = asyncio.create_task(server.monitor.run())
-            server._heartbeat_task = asyncio.create_task(server._heartbeat_loop())
-            if server.event_hub_ws:
-                server._ws_task = asyncio.create_task(server._ws_loop())
-            server._test_task = asyncio.create_task(server._test_event_loop())
-
-        @app.on_event("shutdown")
-        async def _shutdown():
-            server.monitor.stop()
-            if server._monitor_task:
-                server._monitor_task.cancel()
-            if server._heartbeat_task:
-                server._heartbeat_task.cancel()
-            if server._ws_task:
-                server._ws_task.cancel()
-            if hasattr(server, '_test_task') and server._test_task:
-                server._test_task.cancel()
-            if server._ws:
-                await server._ws.close()
-
-        @app.get("/health")
-        async def health():
-            return {
-                "status": "healthy",
-                "details": {
-                    "monitored_modules": len(server.monitor.modules),
-                    "event_hub_connected": server._ws is not None,
-                    "uptime_seconds": round(time.time() - server._start_time),
-                },
-            }
-
-        @app.get("/status")
-        async def status():
-            return server.monitor.get_status()
-
-        return app
-
-    # ── Event Hub WebSocket client ──
-
-    async def _ws_loop(self):
-        """Connect to Event Hub, subscribe, and listen. Reconnect on failure."""
-        retry_delay = 0.5  # start with 0.5s
-        max_delay = 30  # cap at 30s
-        while not self._shutting_down:
-            try:
-                await self._ws_connect()
-                retry_delay = 0.5  # reset on successful connection
-            except asyncio.CancelledError:
-                return
-            except Exception as e:
-                print(f"[watchdog] Event Hub 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)  # exponential backoff
-
-    async def _ws_connect(self):
-        """Single WebSocket session: connect, subscribe, receive loop."""
-        url = f"{self.event_hub_ws}?token={self.token}&id=watchdog"
-        print(f"[watchdog] WS connecting to {self.event_hub_ws}")
-        async with websockets.connect(url, open_timeout=3, ping_interval=None, ping_timeout=None, close_timeout=10) as ws:
-            self._ws = ws
-            print("[watchdog] Connected to Event Hub")
-
-            # Subscribe to module lifecycle events
-            await ws.send(json.dumps({
-                "type": "subscribe",
-                "events": ["system.ready", "module.started", "module.stopped", "module.exiting", "module.ready", "module.shutdown"],
-            }))
-
-            # Send module.ready (once) so Launcher knows we're up
-            if not self._ready_sent:
-                from datetime import datetime, timezone
-                ready_msg = {
-                    "type": "event",
-                    "event_id": str(uuid.uuid4()),
-                    "event": "module.ready",
-                    "source": "watchdog",
-                    "timestamp": datetime.now(timezone.utc).isoformat(),
-                    "data": {
-                        "module_id": "watchdog",
-                        "graceful_shutdown": True,
-                    },
-                }
-                await ws.send(json.dumps(ready_msg))
-                self._ready_sent = True
-
-            # Receive loop
-            async for raw in ws:
-                try:
-                    msg = json.loads(raw)
-                except (json.JSONDecodeError, TypeError):
-                    continue
-
-                try:
-                    msg_type = msg.get("type", "")
-                    if msg_type == "event":
-                        event = msg.get("event", "")
-                        data = msg.get("data") if isinstance(msg.get("data"), dict) else {}
-                        if event == "module.shutdown" and data.get("module_id") == "watchdog":
-                            await self._handle_shutdown(data)
-                            return
-                        await self.monitor.handle_event(msg)
-                    elif msg_type == "ack":
-                        pass  # publish confirmed
-                    elif msg_type == "error":
-                        print(f"[watchdog] Event Hub error: {msg.get('message')}")
-                except Exception as e:
-                    print(f"[watchdog] 事件处理异常（已忽略）: {e}")
-
-    async def _handle_shutdown(self, data: dict):
-        """Handle module.shutdown event — ack, cleanup, ready, exit."""
-        print("[watchdog] Received shutdown request")
-        self._shutting_down = True
-        # Step 1: Send ack
-        await self._publish_event({
-            "event": "module.shutdown.ack",
-            "data": {"module_id": "watchdog", "estimated_cleanup": 2},
-        })
-        # Step 2: Cleanup
-        self.monitor.stop()
-        if self._monitor_task:
-            self._monitor_task.cancel()
-        if self._heartbeat_task:
-            self._heartbeat_task.cancel()
-        if hasattr(self, '_test_task') and self._test_task:
-            self._test_task.cancel()
-        # Step 3: Send ready (before closing WS!)
-        await self._publish_event({
-            "event": "module.shutdown.ready",
-            "data": {"module_id": "watchdog"},
-        })
-        print("[watchdog] Shutdown ready, exiting")
-        # Step 4: Trigger uvicorn exit (WS will close when uvicorn shuts down)
-        if self._uvicorn_server:
-            self._uvicorn_server.should_exit = True
-
-    async def _publish_event(self, event: dict):
-        """Publish an event to Event Hub via WebSocket."""
-        if not self._ws:
-            return
-        from datetime import datetime, timezone
-        msg = {
-            "type": "event",
-            "event_id": str(uuid.uuid4()),
-            "event": event.get("event", ""),
-            "source": "watchdog",
-            "timestamp": datetime.now(timezone.utc).isoformat(),
-            "data": event.get("data", {}),
-        }
-        try:
-            await self._ws.send(json.dumps(msg))
-        except Exception as e:
-            print(f"[watchdog] Failed to publish event: {e}")
-
-    # ── Heartbeat to Registry ──
-
-    async def _heartbeat_loop(self):
-        """Send heartbeat to Registry every 30 seconds."""
-        while True:
-            await asyncio.sleep(30)
-            try:
-                async with httpx.AsyncClient() as client:
-                    await client.post(
-                        f"{self.monitor.registry_url}/modules",
-                        json={"action": "heartbeat", "module_id": "watchdog"},
-                        headers={"Authorization": f"Bearer {self.monitor.own_token}"},
-                        timeout=5,
-                    )
-            except Exception:
-                pass
-
-    async def _test_event_loop(self):
-        """Publish a test event every 5 seconds."""
-        from datetime import datetime, timezone
-        while True:
-            await asyncio.sleep(5)
-            await self._publish_event({
-                "event": "watchdog.test",
-                "data": {
-                    "message": "heartbeat from watchdog",
-                    "timestamp": datetime.now(timezone.utc).isoformat(),
-                },
-            })