@agentunion/kite 1.0.7 → 1.2.0

package/core/registry/server.py

@@ -1,289 +1,339 @@
-"""
-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
-
-    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."""
-        while True:
-            try:
-                await self._ws_connect()
-            except asyncio.CancelledError:
-                return
-            except Exception as e:
-                print(f"[registry] Event Hub connection error: {e}")
-            self._ws = None
-            self._event_hub_connected = False
-            await asyncio.sleep(5)
-
-    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}"
-        async with websockets.connect(ws_url, 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")
-
-            # Send module.ready on first successful connection (delayed ready mechanism)
-            if not self._ready_sent:
-                self._ready_sent = True
-                await self._send_module_ready()
-
-            async for raw in ws:
-                try:
-                    msg = json.loads(raw)
-                except (json.JSONDecodeError, TypeError):
-                    continue
-                msg_type = msg.get("type", "")
-                if msg_type == "error":
-                    print(f"[registry] Event Hub error: {msg.get('message')}")
-
-    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}",
-            },
-        }
-        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
-
-    # ── 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})
+"""
+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})