@agentunion/kite 1.2.0 → 1.3.0

package/core/registry/server.py

@@ -1,339 +0,0 @@
-"""
-Registry HTTP server.
-7 endpoints: /modules, /lookup, /get/{path}, /tokens, /verify, /query, /health.
-All endpoints except /health require Bearer token auth.
-
-Delayed ready mechanism (mechanism 7):
-  Registry does NOT send module.ready immediately after HTTP starts.
-  When Event Hub registers (POST /modules with metadata.ws_endpoint),
-  Registry connects to Event Hub WS and then sends module.ready.
-"""
-
-import asyncio
-import json
-import uuid
-from typing import Any
-
-import websockets
-from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import JSONResponse
-
-from .store import RegistryStore
-
-
-class RegistryServer:
-    """FastAPI-based Registry HTTP server."""
-
-    def __init__(self, store: RegistryStore, launcher_token: str = "", advertise_ip: str = "127.0.0.1"):
-        self.store = store
-        self.launcher_token = launcher_token
-        self.advertise_ip = advertise_ip
-        self.port: int = 0  # set by entry.py before uvicorn.run
-        self.app = self._create_app()
-        self._ttl_task: asyncio.Task | None = None
-        # Event Hub WebSocket
-        self._event_hub_ws_url: str = ""
-        self._ws: object | None = None
-        self._ws_task: asyncio.Task | None = None
-        self._event_hub_connected = False
-        self._ready_sent = False
-        self._uvicorn_server = None  # set by entry.py for graceful shutdown
-        self._shutting_down = False
-
-    def _extract_token(self, request: Request) -> str:
-        """Extract Bearer token from Authorization header."""
-        auth = request.headers.get("Authorization", "")
-        if auth.startswith("Bearer "):
-            return auth[7:].strip()
-        return ""
-
-    def _require_auth(self, request: Request) -> str:
-        """Verify token, return module_id. Raise 401 on failure."""
-        token = self._extract_token(request)
-        module_id = self.store.verify_token(token)
-        if module_id is None:
-            raise HTTPException(status_code=401, detail="Invalid or missing token")
-        return module_id
-
-    def _require_launcher(self, request: Request):
-        """Verify the caller is Launcher. Raise 403 if not."""
-        token = self._extract_token(request)
-        if not self.store.is_launcher(token):
-            raise HTTPException(status_code=403, detail="Only Launcher may call this endpoint")
-
-    # ── Event Hub connection + delayed ready ──
-
-    async def _try_connect_event_hub(self):
-        """Event Hub just registered — connect to it and send module.ready."""
-        if self._event_hub_connected:
-            return
-        eh = self.store.modules.get("event_hub")
-        if not eh:
-            return
-        ws_url = (eh.get("metadata") or {}).get("ws_endpoint", "")
-        if not ws_url:
-            return
-        self._event_hub_ws_url = ws_url
-        if not self._ws_task:
-            self._ws_task = asyncio.create_task(self._ws_loop())
-
-    async def _ws_loop(self):
-        """Connect to Event Hub, 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"[registry] Event Hub connection error: {e}, retrying in {retry_delay:.1f}s")
-            self._ws = None
-            self._event_hub_connected = False
-            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. On first connect, send module.ready."""
-        # Use registry's own per-module token to avoid conflicting with Launcher's connection
-        token = self.store.token_map.get("registry", "") or self.launcher_token
-        ws_url = f"{self._event_hub_ws_url}?token={token}&id=registry"
-        async with websockets.connect(ws_url, open_timeout=3, ping_interval=None, ping_timeout=None, close_timeout=10) as ws:
-            self._ws = ws
-            self._event_hub_connected = True
-            print("[registry] Connected to Event Hub")
-
-            # Subscribe to shutdown events
-            await ws.send(json.dumps({"type": "subscribe", "events": ["module.shutdown"]}))
-
-            # Send module.ready on first successful connection (delayed ready mechanism)
-            if not self._ready_sent:
-                self._ready_sent = True
-                # Self-register so watchdog (and others) can discover registry
-                # via /lookup. Without this, health checks fail (empty api_endpoint).
-                self.store.register_module({
-                    "module_id": "registry",
-                    "module_type": "infrastructure",
-                    "api_endpoint": f"http://{self.advertise_ip}:{self.port}",
-                    "health_endpoint": "/health",
-                })
-                await self._send_module_ready()
-
-            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") == "registry":
-                            await self._handle_shutdown(data)
-                            return
-                    elif msg_type == "error":
-                        print(f"[registry] Event Hub error: {msg.get('message')}")
-                except Exception as e:
-                    print(f"[registry] 事件处理异常（已忽略）: {e}")
-
-    async def _send_module_ready(self):
-        """Send module.ready event to Event Hub. Launcher is listening for this."""
-        from datetime import datetime, timezone
-        msg = {
-            "type": "event",
-            "event_id": str(uuid.uuid4()),
-            "event": "module.ready",
-            "source": "registry",
-            "timestamp": datetime.now(timezone.utc).isoformat(),
-            "data": {
-                "module_id": "registry",
-                "api_endpoint": f"http://{self.advertise_ip}:{self.port}",
-                "graceful_shutdown": True,
-            },
-        }
-        try:
-            await self._ws.send(json.dumps(msg))
-            print("[registry] Sent module.ready")
-        except Exception as e:
-            print(f"[registry] Failed to send module.ready: {e}")
-
-    async def _publish_event(self, event_type: str, data: dict):
-        """Publish event to Event Hub. Best-effort, no-op if not connected."""
-        if not self._ws:
-            return
-        from datetime import datetime, timezone
-        msg = {
-            "type": "event",
-            "event_id": str(uuid.uuid4()),
-            "event": event_type,
-            "source": "registry",
-            "timestamp": datetime.now(timezone.utc).isoformat(),
-            "data": data,
-        }
-        try:
-            await self._ws.send(json.dumps(msg))
-        except Exception:
-            pass
-
-    async def _handle_shutdown(self, data: dict):
-        """Handle module.shutdown event — ack, cleanup, ready, exit."""
-        print("[registry] Received shutdown request")
-        self._shutting_down = True
-        # Step 1: Send ack
-        await self._publish_event("module.shutdown.ack", {
-            "module_id": "registry",
-            "estimated_cleanup": 2,
-        })
-        # Step 2: Cleanup
-        if self._ttl_task:
-            self._ttl_task.cancel()
-        # Step 3: Send ready (before closing WS!)
-        await self._publish_event("module.shutdown.ready", {
-            "module_id": "registry",
-        })
-        print("[registry] 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
-
-    # ── App factory ──
-
-    def _create_app(self) -> FastAPI:
-        app = FastAPI(title="Kite Registry", docs_url=None, redoc_url=None)
-        server = self
-
-        @app.on_event("startup")
-        async def _startup():
-            server._ttl_task = asyncio.create_task(server._ttl_loop())
-
-        @app.on_event("shutdown")
-        async def _shutdown():
-            if server._ttl_task:
-                server._ttl_task.cancel()
-            if server._ws_task:
-                server._ws_task.cancel()
-            if server._ws:
-                await server._ws.close()
-
-        # ── 1. POST /modules ──
-
-        @app.post("/modules")
-        async def modules(request: Request):
-            caller = server._require_auth(request)
-            body = await request.json()
-            action = body.get("action", "")
-
-            if action == "register":
-                if "module_id" not in body:
-                    raise HTTPException(400, "module_id required")
-                # Only Launcher or the module itself may register
-                if caller != "launcher" and caller != body["module_id"]:
-                    raise HTTPException(403, f"Module '{caller}' cannot register as '{body['module_id']}'")
-                result = server.store.register_module(body)
-                if result.get("ok"):
-                    mid = body["module_id"]
-                    await server._publish_event("module.registered", {"module_id": mid})
-                    # If Event Hub just registered, connect and send module.ready
-                    ws_endpoint = (body.get("metadata") or {}).get("ws_endpoint")
-                    if ws_endpoint and mid == "event_hub":
-                        await server._try_connect_event_hub()
-                return result
-
-            elif action == "deregister":
-                mid = body.get("module_id")
-                if not mid:
-                    raise HTTPException(400, "module_id required")
-                if caller != "launcher" and caller != mid:
-                    raise HTTPException(403, f"Module '{caller}' cannot deregister '{mid}'")
-                result = server.store.deregister_module(mid)
-                if result.get("ok"):
-                    await server._publish_event("module.unregistered", {"module_id": mid})
-                return result
-
-            elif action == "heartbeat":
-                mid = body.get("module_id")
-                if not mid:
-                    raise HTTPException(400, "module_id required")
-                if caller != "launcher" and caller != mid:
-                    raise HTTPException(403, f"Module '{caller}' cannot heartbeat for '{mid}'")
-                result = server.store.heartbeat(mid)
-                if result.get("ok"):
-                    await server._publish_event("module.heartbeat", {"module_id": mid})
-                return result
-
-            else:
-                raise HTTPException(400, f"Unknown action: {action}")
-
-        # ── 2. GET /lookup ──
-
-        @app.get("/lookup")
-        async def lookup(request: Request, field: str = None, module: str = None, value: str = None):
-            server._require_auth(request)
-            return server.store.lookup(field=field, module=module, value=value)
-
-        # ── 3. GET /get/{path} ──
-
-        @app.get("/get/{path:path}")
-        async def get_by_path(request: Request, path: str):
-            server._require_auth(request)
-            val, found = server.store.get_by_path(path)
-            if not found:
-                raise HTTPException(404, f"Path not found: {path}")
-            return val
-
-        # ── 4. POST /tokens ──
-
-        @app.post("/tokens")
-        async def register_tokens(request: Request):
-            server._require_launcher(request)
-            body = await request.json()
-            server.store.register_tokens(body)
-            return {"ok": True}
-
-        # ── 5. POST /verify ──
-
-        @app.post("/verify")
-        async def verify_token(request: Request):
-            server._require_auth(request)
-            body = await request.json()
-            target_token = body.get("token", "")
-            module_id = server.store.verify_token(target_token)
-            if module_id:
-                return {"ok": True, "module_id": module_id}
-            return {"ok": False}
-
-        # ── 6. POST /query (stub) ──
-
-        @app.post("/query")
-        async def query(request: Request):
-            server._require_auth(request)
-            body = await request.json()
-            question = body.get("question", "")
-            return {"ok": False, "error": "LLM query not implemented yet", "question": question}
-
-        # ── 7. GET /health ──
-
-        @app.get("/health")
-        async def health():
-            return {
-                "status": "healthy",
-                "module_count": len(server.store.modules),
-                "online_count": sum(
-                    1 for m in server.store.modules.values()
-                    if m.get("status") == "online"
-                ),
-            }
-
-        return app
-
-    async def _ttl_loop(self):
-        """Periodically check heartbeat TTL and publish offline events."""
-        while True:
-            await asyncio.sleep(10)
-            expired = self.store.check_ttl()
-            for mid in expired:
-                await self._publish_event("module.offline", {"module_id": mid})

package/extensions/services/backup/server.py

@@ -1,244 +0,0 @@
-"""
-Backup 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 BackupServer:
-
-    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 Backup", docs_url=None, redoc_url=None)
-        server = self
-
-        @app.on_event("startup")
-        async def _startup():
-            elapsed = time.monotonic() - server.boot_t0 if server.boot_t0 else 0
-            print(f"[backup] FastAPI startup event triggered ({elapsed:.1f}s)")
-            server._heartbeat_task = asyncio.create_task(server._heartbeat_loop())
-            if server.event_hub_ws:
-                print(f"[backup] Creating Event Hub connection task...")
-                server._ws_task = asyncio.create_task(server._ws_loop())
-            server._test_task = asyncio.create_task(server._test_event_loop())
-            print(f"[backup] All background tasks created")
-
-        @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("[backup] 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": "backup",
-                "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."""
-        elapsed = time.monotonic() - self.boot_t0 if self.boot_t0 else 0
-        print(f"[backup] _ws_loop started ({elapsed:.1f}s)")
-        retry_delay = 0.5  # start with 0.5s
-        max_delay = 30  # cap at 30s
-        while not self._shutting_down:
-            try:
-                elapsed = time.monotonic() - self.boot_t0 if self.boot_t0 else 0
-                print(f"[backup] Attempting Event Hub connection... ({elapsed:.1f}s)")
-                await self._ws_connect()
-                retry_delay = 0.5  # reset on successful connection
-            except asyncio.CancelledError:
-                return
-            except Exception as e:
-                print(f"[backup] 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=backup"
-        print(f"[backup] 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"[backup] 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": "backup",
-                    "timestamp": datetime.now(timezone.utc).isoformat(),
-                    "data": {
-                        "module_id": "backup",
-                        "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"[backup] 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 == "backup":
-                                await self._handle_shutdown(ws)
-                                return
-                    elif msg_type == "ack":
-                        pass  # publish confirmed
-                    elif msg_type == "error":
-                        print(f"[backup] Event Hub error: {msg.get('message')}")
-                except Exception as e:
-                    print(f"[backup] 事件处理异常（已忽略）: {e}")
-
-    async def _handle_shutdown(self, ws):
-        """Handle module.shutdown: ack → cleanup → ready → exit."""
-        print("[backup] Received module.shutdown")
-        self._shutting_down = True
-
-        # Step 1: Send ack
-        await self._publish_event({
-            "event": "module.shutdown.ack",
-            "data": {"module_id": "backup", "estimated_cleanup": 2},
-        })
-        print("[backup] 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": "backup"},
-        })
-        print("[backup] 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": "backup",
-            "timestamp": datetime.now(timezone.utc).isoformat(),
-            "data": event.get("data", {}),
-        }
-        try:
-            await self._ws.send(json.dumps(msg))
-        except Exception as e:
-            print(f"[backup] 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": "backup"},
-                        headers={"Authorization": f"Bearer {self.token}"},
-                        timeout=5,
-                    )
-                print("[backup] 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": "backup.test",
-                "data": {
-                    "message": "test event from backup",
-                    "timestamp": datetime.now(timezone.utc).isoformat(),
-                },
-            })
-            print("[backup] test event published")