@agentunion/kite 1.2.0 → 1.3.0

package/kernel/rpc_router.py

@@ -0,0 +1,388 @@
+"""
+JSON-RPC 2.0 message dispatcher for Kernel.
+Routes builtin methods (registry.*, event.*, kernel.*) and forwards
+cross-module RPC requests with timeout tracking.
+"""
+
+import asyncio
+import json
+import time
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+
+from starlette.websockets import WebSocket
+
+from .registry_store import RegistryStore
+from .event_hub import EventHub
+
+
+# ── JSON-RPC 2.0 error codes ──
+
+PARSE_ERROR = -32700
+INVALID_REQUEST = -32600
+METHOD_NOT_FOUND = -32601
+INVALID_PARAMS = -32602
+INTERNAL_ERROR = -32603
+
+# Kite custom error codes
+MODULE_OFFLINE = -32001
+RPC_TIMEOUT = -32002
+AUTH_FAILED = -32003
+PERMISSION_DENIED = -32004
+DUPLICATE_EVENT = -32005
+
+DEFAULT_FORWARD_TIMEOUT = 5.0  # seconds
+
+
+@dataclass
+class PendingForward:
+    """Tracks a forwarded cross-module RPC awaiting response."""
+    caller_ws: WebSocket
+    original_id: str
+    caller_id: str
+    target_id: str
+    created_at: float = field(default_factory=time.time)
+    timeout_handle: asyncio.TimerHandle | None = None
+
+
+def _result_msg(msg_id: str, result: dict) -> str:
+    return json.dumps({"jsonrpc": "2.0", "id": msg_id, "result": result})
+
+
+def _error_msg(msg_id: str, code: int, message: str, data: dict = None) -> str:
+    err = {"code": code, "message": message}
+    if data:
+        err["data"] = data
+    return json.dumps({"jsonrpc": "2.0", "id": msg_id, "error": err})
+
+
+class RpcRouter:
+    """JSON-RPC 2.0 message dispatcher.
+
+    Handles:
+    - 13 builtin methods (registry.*, event.*, kernel.*)
+    - Cross-module RPC forwarding with timeout
+    - RPC response matching for forwarded calls
+    """
+
+    def __init__(self, registry: RegistryStore, event_hub: EventHub,
+                 connections: dict[str, WebSocket], kernel_server):
+        self.registry = registry
+        self.event_hub = event_hub
+        self.connections = connections
+        self.kernel_server = kernel_server  # Direct reference to KernelServer
+
+        # Builtin method dispatch table
+        self.methods: dict[str, callable] = {
+            "registry.register":      self._registry_register,
+            "registry.deregister":    self._registry_deregister,
+            "registry.heartbeat":     self._registry_heartbeat,
+            "registry.lookup":        self._registry_lookup,
+            "registry.get":           self._registry_get,
+            "registry.verify":        self._registry_verify,
+            "event.publish":          self._event_publish,
+            "event.subscribe":        self._event_subscribe,
+            "event.unsubscribe":      self._event_unsubscribe,
+            "kernel.ping":            self._kernel_ping,
+            "kernel.stats":           self._kernel_stats,
+            "kernel.health":          self._kernel_health,
+            "kernel.generate_tokens": self._kernel_generate_tokens,
+            "kernel.register_tokens": self._kernel_register_tokens,
+            "kernel.shutdown":        self._kernel_shutdown,
+        }
+
+        # Pending cross-module forwards: internal_id -> PendingForward
+        self._pending: dict[str, PendingForward] = {}
+
+    # ── Main dispatch ──
+
+    async def dispatch(self, caller_id: str, ws: WebSocket, msg: dict):
+        """Route a parsed JSON-RPC message to builtin handler or cross-module forward."""
+        method = msg.get("method", "")
+        msg_id = msg.get("id")
+
+        # JSON-RPC Notification (no id) — currently not handled from clients
+        if msg_id is None:
+            return
+
+        params = msg.get("params") or {}
+
+        # Builtin method
+        handler = self.methods.get(method)
+        if handler:
+            try:
+                result = await handler(caller_id, params)
+                await ws.send_text(_result_msg(msg_id, result))
+            except Exception as e:
+                print(f"[kernel] RPC handler error ({method}): {e}")
+                await ws.send_text(_error_msg(msg_id, INTERNAL_ERROR, str(e)))
+            return
+
+        # Cross-module forward: method prefix is target module_id
+        dot_idx = method.find(".")
+        if dot_idx > 0:
+            target = method[:dot_idx]
+            if target in self.connections:
+                await self._forward(caller_id, ws, msg_id, target, method, params)
+                return
+            # Target not connected — check if registered but offline
+            if target in self.registry.modules:
+                await ws.send_text(_error_msg(
+                    msg_id, MODULE_OFFLINE, f"Module offline: {target}"))
+                return
+
+        # Method not found
+        await ws.send_text(_error_msg(msg_id, METHOD_NOT_FOUND, f"Method not found: {method}"))
+
+    async def handle_response(self, module_id: str, msg: dict):
+        """Handle an RPC response from a module (matches pending forwards)."""
+        msg_id = msg.get("id")
+        if not msg_id:
+            return
+
+        pending = self._pending.pop(msg_id, None)
+        if not pending:
+            return  # orphan response, ignore
+
+        # Cancel timeout
+        if pending.timeout_handle:
+            pending.timeout_handle.cancel()
+
+        # Forward response to original caller with original ID
+        response = {"jsonrpc": "2.0", "id": pending.original_id}
+        if "result" in msg:
+            response["result"] = msg["result"]
+        elif "error" in msg:
+            response["error"] = msg["error"]
+        else:
+            response["result"] = None
+
+        try:
+            await pending.caller_ws.send_text(json.dumps(response))
+        except Exception:
+            pass  # caller disconnected
+
+    # ── Cross-module forwarding ──
+
+    async def _forward(self, caller_id: str, caller_ws: WebSocket,
+                       original_id: str, target: str, method: str, params: dict):
+        """Forward RPC request to target module with timeout tracking."""
+        internal_id = f"fwd-{uuid.uuid4().hex[:12]}"
+
+        # Extract timeout from params (optional _timeout field)
+        timeout = DEFAULT_FORWARD_TIMEOUT
+        if isinstance(params, dict) and "_timeout" in params:
+            try:
+                timeout = float(params.pop("_timeout"))
+            except (ValueError, TypeError):
+                pass
+
+        # Strip target prefix from method for the forwarded request
+        actual_method = method[len(target) + 1:]  # e.g. "watchdog.get_status" -> "get_status"
+
+        # Record pending forward
+        loop = asyncio.get_event_loop()
+        pending = PendingForward(
+            caller_ws=caller_ws,
+            original_id=original_id,
+            caller_id=caller_id,
+            target_id=target,
+        )
+        pending.timeout_handle = loop.call_later(
+            timeout, lambda iid=internal_id: asyncio.ensure_future(self._handle_timeout(iid))
+        )
+        self._pending[internal_id] = pending
+
+        # Send to target
+        fwd_msg = json.dumps({
+            "jsonrpc": "2.0",
+            "id": internal_id,
+            "method": actual_method,
+            "params": params or {},
+        })
+        target_ws = self.connections.get(target)
+        if target_ws:
+            try:
+                await target_ws.send_text(fwd_msg)
+            except Exception:
+                # Target send failed — clean up and error to caller
+                self._pending.pop(internal_id, None)
+                if pending.timeout_handle:
+                    pending.timeout_handle.cancel()
+                await caller_ws.send_text(_error_msg(
+                    original_id, MODULE_OFFLINE, f"Failed to reach module: {target}"))
+        else:
+            # Target disconnected between check and send
+            self._pending.pop(internal_id, None)
+            if pending.timeout_handle:
+                pending.timeout_handle.cancel()
+            await caller_ws.send_text(_error_msg(
+                original_id, MODULE_OFFLINE, f"Module offline: {target}"))
+
+    async def _handle_timeout(self, internal_id: str):
+        """Called when a forwarded RPC times out."""
+        pending = self._pending.pop(internal_id, None)
+        if not pending:
+            return
+        try:
+            await pending.caller_ws.send_text(_error_msg(
+                pending.original_id, RPC_TIMEOUT,
+                f"RPC timeout waiting for {pending.target_id}"))
+        except Exception:
+            pass
+
+    # ── Builtin handlers: registry.* ──
+
+    async def _registry_register(self, caller_id: str, params: dict) -> dict:
+        mid = params.get("module_id")
+        if not mid:
+            return {"ok": False, "error": "module_id required"}
+        # Permission: only Launcher or the module itself
+        if caller_id != "launcher" and caller_id != mid:
+            return {"ok": False, "error": f"Module '{caller_id}' cannot register as '{mid}'"}
+        result = self.registry.register_module(params)
+        if result.get("ok"):
+            self.event_hub.publish_internal("module.registered", {"module_id": mid})
+
+            # When Launcher registers, Kernel publishes its own module.ready
+            if mid == "launcher" and self.kernel_server:
+                if not self.kernel_server._ready_published:
+                    self.kernel_server.publish_ready()
+                    self.kernel_server._ready_published = True
+                    print(f"[kernel] launcher registered → kernel module.ready published")
+        return result
+
+    async def _registry_deregister(self, caller_id: str, params: dict) -> dict:
+        mid = params.get("module_id")
+        if not mid:
+            return {"ok": False, "error": "module_id required"}
+        if caller_id != "launcher" and caller_id != mid:
+            return {"ok": False, "error": f"Module '{caller_id}' cannot deregister '{mid}'"}
+        result = self.registry.deregister_module(mid)
+        if result.get("ok"):
+            self.event_hub.publish_internal("module.unregistered", {"module_id": mid})
+        return result
+
+    async def _registry_heartbeat(self, caller_id: str, params: dict) -> dict:
+        mid = params.get("module_id")
+        if not mid:
+            return {"ok": False, "error": "module_id required"}
+        if caller_id != "launcher" and caller_id != mid:
+            return {"ok": False, "error": f"Module '{caller_id}' cannot heartbeat for '{mid}'"}
+        return self.registry.heartbeat(mid)
+
+    async def _registry_lookup(self, caller_id: str, params: dict) -> dict:
+        results = self.registry.lookup(
+            field=params.get("field"),
+            module=params.get("module"),
+            value=params.get("value"),
+        )
+        return {"ok": True, "results": results}
+
+    async def _registry_get(self, caller_id: str, params: dict) -> dict:
+        path = params.get("path", "")
+        if not path:
+            return {"ok": False, "error": "path required"}
+        val, found = self.registry.get_by_path(path)
+        if not found:
+            return {"ok": False, "error": f"Path not found: {path}"}
+        return {"ok": True, "value": val}
+
+    async def _registry_verify(self, caller_id: str, params: dict) -> dict:
+        token = params.get("token", "")
+        module_id = self.registry.verify_token(token)
+        if module_id:
+            return {"ok": True, "module_id": module_id}
+        return {"ok": False}
+
+    # ── Builtin handlers: event.* ──
+
+    async def _event_publish(self, caller_id: str, params: dict) -> dict:
+        event_id = params.get("event_id", "")
+        event_type = params.get("event", "")
+        data = params.get("data")
+        echo = params.get("echo", False)
+        return self.event_hub.publish_event(caller_id, event_id, event_type, data, echo)
+
+    async def _event_subscribe(self, caller_id: str, params: dict) -> dict:
+        events = params.get("events", [])
+        if not isinstance(events, list) or not events:
+            return {"ok": False, "error": "events must be a non-empty list"}
+        self.event_hub.handle_subscribe(caller_id, events)
+        return {"ok": True}
+
+    async def _event_unsubscribe(self, caller_id: str, params: dict) -> dict:
+        events = params.get("events", [])
+        if not isinstance(events, list) or not events:
+            return {"ok": False, "error": "events must be a non-empty list"}
+        self.event_hub.handle_unsubscribe(caller_id, events)
+        return {"ok": True}
+
+    # ── Builtin handlers: kernel.* ──
+
+    async def _kernel_ping(self, caller_id: str, params: dict) -> dict:
+        return {"pong": True, "timestamp": datetime.now(timezone.utc).isoformat()}
+
+    async def _kernel_stats(self, caller_id: str, params: dict) -> dict:
+        return self.event_hub.get_stats()
+
+    async def _kernel_health(self, caller_id: str, params: dict) -> dict:
+        eh_health = self.event_hub.get_health()
+        return {
+            "status": "healthy",
+            "module_count": len(self.registry.modules),
+            "online_count": sum(
+                1 for m in self.registry.modules.values()
+                if m.get("status") == "online"
+            ),
+            "event_stats": eh_health.get("details", {}),
+        }
+
+    async def _kernel_generate_tokens(self, caller_id: str, params: dict) -> dict:
+        """Generate tokens for a list of module names.
+
+        Args:
+            params: {"modules": ["mod1", "mod2", ...]}
+
+        Returns:
+            {"ok": True, "tokens": {"mod1": "token1", "mod2": "token2", ...}}
+        """
+        # Only Launcher may request token generation
+        if caller_id != "launcher":
+            return {"ok": False, "error": "Only Launcher may generate tokens"}
+
+        modules = params.get("modules", [])
+        if not isinstance(modules, list):
+            return {"ok": False, "error": "modules must be a list"}
+
+        import secrets
+        tokens = {}
+        for module_name in modules:
+            tokens[module_name] = secrets.token_hex(32)
+
+        # Register tokens in registry
+        self.registry.register_tokens(tokens)
+
+        return {"ok": True, "tokens": tokens}
+
+    async def _kernel_register_tokens(self, caller_id: str, params: dict) -> dict:
+        # Only Launcher may register tokens
+        if caller_id != "launcher":
+            return {"ok": False, "error": "Only Launcher may register tokens"}
+        self.registry.register_tokens(params)
+        return {"ok": True}
+
+    async def _kernel_shutdown(self, caller_id: str, params: dict) -> dict:
+        """Shutdown Kernel. Only Launcher may call this."""
+        if caller_id != "launcher":
+            return {"ok": False, "error": "Only Launcher may shutdown Kernel"}
+
+        print("[kernel] Received shutdown request from Launcher")
+
+        # Schedule shutdown (don't block RPC response)
+        if self.kernel_server:
+            asyncio.create_task(self.kernel_server.shutdown())
+
+        return {"ok": True}
+

package/kernel/server.py

@@ -0,0 +1,267 @@
+"""
+Kernel unified server.
+FastAPI app with WebSocket endpoint (RPC + Event) and minimal HTTP endpoints (/health, /stats).
+Merges Registry + Event Hub into a single process.
+"""
+
+import asyncio
+import json
+import os
+
+from fastapi import FastAPI
+from starlette.websockets import WebSocket, WebSocketDisconnect
+
+from .registry_store import RegistryStore
+from .event_hub import EventHub
+from .rpc_router import RpcRouter
+
+try:
+    import orjson
+    def _loads(raw: str):
+        return orjson.loads(raw)
+except ImportError:
+    def _loads(raw: str):
+        return json.loads(raw)
+
+
+class KernelServer:
+    """Merged Registry + Event Hub server.
+
+    Single WebSocket endpoint handles:
+    - JSON-RPC 2.0 requests (builtin + cross-module forward)
+    - JSON-RPC 2.0 responses (from forwarded calls)
+    - Event notifications (delivered to subscribers)
+    """
+
+    def __init__(self, launcher_token: str = None, advertise_ip: str = "127.0.0.1"):
+        self.advertise_ip = advertise_ip
+        self.port: int = 0  # set by entry.py before uvicorn.run
+
+        # Core components
+        self.registry = RegistryStore(launcher_token)  # Can be None
+        self.event_hub = EventHub()
+
+        # Shared connection table (module_id -> WebSocket)
+        # RpcRouter and EventHub both reference this
+        self.connections: dict[str, WebSocket] = {}
+
+        # RPC router (pass self reference)
+        self.rpc_router = RpcRouter(
+            self.registry,
+            self.event_hub,
+            self.connections,
+            kernel_server=self
+        )
+
+        # Background tasks
+        self._ttl_task: asyncio.Task | None = None
+        self._dedup_task: asyncio.Task | None = None
+        self._uvicorn_server = None  # set by entry.py for graceful shutdown
+        self._shutting_down = False
+
+        # Launcher connection tracking
+        self._launcher_connected = False
+        self._launcher_subscribed = False
+        self._ready_published = False
+
+        # Build FastAPI app
+        self.app = self._create_app()
+
+    # ── App factory ──
+
+    def _create_app(self) -> FastAPI:
+        app = FastAPI(title="Kite Kernel", docs_url=None, redoc_url=None)
+        server = self
+
+        @app.on_event("startup")
+        async def _startup():
+            server._ttl_task = asyncio.create_task(server._ttl_loop())
+            server._dedup_task = asyncio.create_task(server._dedup_loop())
+
+        @app.on_event("shutdown")
+        async def _shutdown():
+            if server._ttl_task:
+                server._ttl_task.cancel()
+            if server._dedup_task:
+                server._dedup_task.cancel()
+
+        # ── WebSocket endpoint ──
+
+        @app.websocket("/ws")
+        async def ws_endpoint(ws: WebSocket):
+            token = ws.query_params.get("token", "")
+            mid_hint = ws.query_params.get("id", "")
+
+            # Token verification (all modules including Launcher need token)
+            module_id = server.registry.verify_token(token)
+            if module_id is None:
+                # Must accept before close (Starlette requirement)
+                await ws.accept()
+                print(f"[kernel] Auth failed: token={token[:8]}... hint={mid_hint}")
+                try:
+                    await ws.close(code=4001, reason="Authentication failed")
+                except Exception:
+                    pass
+                return
+
+            # Use id hint for debug mode
+            if module_id == "debug" and mid_hint:
+                module_id = mid_hint
+
+            await ws.accept()
+
+            # Register connection in both EventHub and shared connections table
+            server.event_hub.add_connection(module_id, ws)
+            server.connections[module_id] = ws
+
+            # Track Launcher connection
+            if module_id == "launcher":
+                server._launcher_connected = True
+                print(f"[kernel] launcher connected")
+
+            # Renew heartbeat on connect
+            if module_id in server.registry.modules:
+                server.registry.heartbeat(module_id)
+
+            try:
+                while True:
+                    raw = await ws.receive_text()
+                    try:
+                        msg = _loads(raw)
+                    except Exception:
+                        # JSON parse error — send error if there's an id
+                        try:
+                            await ws.send_text(json.dumps({
+                                "jsonrpc": "2.0", "id": None,
+                                "error": {"code": -32700, "message": "Parse error"}
+                            }))
+                        except Exception:
+                            pass
+                        continue
+
+                    if not isinstance(msg, dict):
+                        continue
+
+                    # Classify message type:
+                    has_method = "method" in msg
+                    has_id = "id" in msg
+                    has_result = "result" in msg
+                    has_error = "error" in msg
+
+                    if has_method and has_id:
+                        # RPC Request → dispatch to handler
+                        await server.rpc_router.dispatch(module_id, ws, msg)
+                    elif has_id and (has_result or has_error):
+                        # RPC Response → match to pending forward
+                        await server.rpc_router.handle_response(module_id, msg)
+                    # else: notification or unknown — ignore
+
+            except WebSocketDisconnect:
+                pass
+            except Exception as e:
+                err = str(e).lower()
+                if "not connected" not in err and "closed" not in err:
+                    print(f"[kernel] WebSocket error for {module_id}: {e}")
+            finally:
+                # Cleanup
+                server.event_hub.remove_connection(module_id)
+                server.connections.pop(module_id, None)
+                server.registry.set_offline(module_id)
+                server.event_hub.publish_internal(
+                    "module.offline", {"module_id": module_id})
+
+        # ── HTTP endpoints (debug only) ──
+
+        @app.get("/health")
+        async def health():
+            eh_health = server.event_hub.get_health()
+            return {
+                "status": "healthy",
+                "module_count": len(server.registry.modules),
+                "online_count": sum(
+                    1 for m in server.registry.modules.values()
+                    if m.get("status") == "online"
+                ),
+                "event_stats": eh_health.get("details", {}),
+            }
+
+        @app.get("/stats")
+        async def stats():
+            return {
+                "registry": {
+                    "modules": {
+                        mid: {
+                            "status": data.get("status"),
+                            "module_type": data.get("module_type"),
+                            "registered_at": data.get("registered_at"),
+                        }
+                        for mid, data in server.registry.modules.items()
+                    },
+                },
+                "event_hub": server.event_hub.get_stats(),
+            }
+
+        return app
+
+    # ── Background loops ──
+
+    async def _ttl_loop(self):
+        """Check heartbeat TTL every 10s and publish offline events."""
+        while True:
+            await asyncio.sleep(10)
+            try:
+                expired = self.registry.check_ttl()
+                for mid in expired:
+                    self.event_hub.publish_internal(
+                        "module.offline", {"module_id": mid})
+            except Exception as e:
+                print(f"[kernel] TTL loop error: {e}")
+
+    async def _dedup_loop(self):
+        """Clean up dedup table every 30s."""
+        while True:
+            await asyncio.sleep(30)
+            try:
+                await asyncio.get_event_loop().run_in_executor(
+                    None, self.event_hub.dedup.cleanup)
+            except Exception as e:
+                print(f"[kernel] Dedup cleanup error: {e}")
+
+    # ── Self-registration ──
+
+    def self_register(self):
+        """Register Kernel itself in the registry (in-memory, no RPC needed)."""
+        self.registry.register_module({
+            "module_id": "kernel",
+            "module_type": "infrastructure",
+            "api_endpoint": f"http://{self.advertise_ip}:{self.port}",
+            "health_endpoint": "/health",
+            "metadata": {
+                "ws_endpoint": f"ws://{self.advertise_ip}:{self.port}/ws",
+            },
+        })
+
+    def publish_ready(self):
+        """Publish module.ready event for Kernel (internal, no WS needed)."""
+        self.event_hub.publish_internal("module.ready", {
+            "module_id": "kernel",
+            "ws_endpoint": f"ws://{self.advertise_ip}:{self.port}/ws",
+            "graceful_shutdown": True,
+        })
+
+    async def shutdown(self):
+        """Shutdown Kernel gracefully. Called by Launcher via RPC."""
+        if self._shutting_down:
+            return
+
+        self._shutting_down = True
+        print("[kernel] Shutting down...")
+
+        # Brief delay to ensure RPC response is sent
+        await asyncio.sleep(0.1)
+
+        # Trigger uvicorn shutdown
+        if self._uvicorn_server:
+            self._uvicorn_server.should_exit = True
+        else:
+            print("[kernel] Warning: uvicorn server reference not set")

package/launcher/__init__.py

@@ -0,0 +1,10 @@
+"""
+Launcher package — Kite system entry point.
+
+Can be started in two ways:
+1. Via main.py (with code stats): python main.py
+2. Directly (no stats): python -m launcher
+"""
+from .entry import start_launcher
+
+__all__ = ["start_launcher"]

package/launcher/__main__.py

@@ -0,0 +1,6 @@
+"""
+Direct launcher entry point: python -m launcher
+"""
+if __name__ == "__main__":
+    from .entry import start_launcher
+    start_launcher()