@agentunion/kite 1.0.7 → 1.3.0

package/{core/registry/store.py → kernel/registry_store.py}

@@ -1,7 +1,11 @@
 """
-Registry in-memory store.
+Kernel registry in-memory store.
 Manages module records, token verification, heartbeat TTL, lookup and get-by-path.
-No persistence — Registry crash triggers Kite full restart, all modules re-register.
+No persistence — Kernel crash triggers Kite full restart, all modules re-register.
+
+Adapted from core/registry/store.py for the merged Kernel module:
+  - api_endpoint is now optional (modules no longer need HTTP)
+  - Added set_offline() for WebSocket disconnect handling
 """
 
 import fnmatch
@@ -25,18 +29,24 @@ class RegistryStore:
 
     def verify_token(self, token: str) -> str | None:
         """Return module_id if token is valid, None otherwise.
-        In debug mode (KITE_DEBUG=1), any non-empty token is accepted."""
-        if self.is_debug and token:
-            return "debug"
+        In debug mode (KITE_DEBUG=1), any non-empty token is accepted,
+        but known tokens still resolve to their proper module_id first."""
         if token == self.launcher_token:
             return "launcher"
         for mid, tok in self.token_map.items():
             if token == tok:
                 return mid
+        if self.is_debug and token:
+            return "debug"
         return None
 
     def is_launcher(self, token: str) -> bool:
-        return token == self.launcher_token
+        """Check if token belongs to Launcher. In debug mode, also accept any token."""
+        if token == self.launcher_token:
+            return True
+        if self.is_debug and token:
+            return True
+        return False
 
     def register_tokens(self, mapping: dict[str, str]):
         """Register module_id -> token mapping. Only Launcher may call this."""
@@ -44,7 +54,7 @@ class RegistryStore:
 
     # ── Module lifecycle ──
 
-    _REQUIRED_FIELDS = ("module_id", "module_type", "api_endpoint")
+    _REQUIRED_FIELDS = ("module_id", "module_type")  # api_endpoint now optional
 
     def register_module(self, data: dict) -> dict:
         """Register or update a module. Idempotent — same module_id overwrites."""
@@ -63,7 +73,7 @@ class RegistryStore:
                     if other_mid == mid:
                         continue
                     if tool_name in other_data.get("tools", {}):
-                        print(f"[registry] WARNING: tool '{tool_name}' registered by both '{other_mid}' and '{mid}'")
+                        print(f"[kernel] WARNING: tool '{tool_name}' registered by both '{other_mid}' and '{mid}'")
 
         # Strip action field — it's a request verb, not part of the registration payload
         record = {k: v for k, v in data.items() if k != "action"}
@@ -87,6 +97,11 @@ class RegistryStore:
         self.modules[module_id]["status"] = "online"
         return {"ok": True}
 
+    def set_offline(self, module_id: str):
+        """Mark a module as offline (called on WebSocket disconnect)."""
+        if module_id in self.modules:
+            self.modules[module_id]["status"] = "offline"
+
     def check_ttl(self) -> list[str]:
         """Mark modules as offline if heartbeat expired. Returns list of newly-offline module_ids."""
         now = time.time()

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}
+