@agentunion/kite 1.5.0 → 1.6.1

package/kernel/module.md

@@ -9,7 +9,8 @@ version: 1.0.1
 advertise_ip: 127.0.0.1
 monitor: true
 preferred_port: 0
-display_order: 100
+display_order: 99
+manual_control: false
 ---
 # Kernel
 

package/kernel/registry_store.py

@@ -9,6 +9,7 @@ Adapted from core/registry/store.py for the merged Kernel module:
 """
 
 import fnmatch
+import hmac
 import os
 import time
 from typing import Any
@@ -19,20 +20,125 @@ class RegistryStore:
     def __init__(self, launcher_token: str):
         self.launcher_token = launcher_token
         self.token_map: dict[str, str] = {}       # module_id -> token
-        self.modules: dict[str, dict] = {}         # module_id -> registration payload
+        self.modules: dict[str, dict] = {}         # module_id -> registration payload (includes "status" field)
         self.is_debug = os.environ.get("KITE_DEBUG") == "1"
         self.last_update_time = time.time()        # Global registry update timestamp
 
+        # 认证速率限制：module_id → [timestamp, ...]
+        self._auth_rate_limits: dict[str, list[float]] = {}
+        self._auth_rate_max = 20       # 10 秒内最多 20 次
+        self._auth_rate_window = 10.0  # 秒
+
+        # ── 模块初始化时间戳（system.require_init 判断用） ──
+        self._last_init_time: dict[str, float] = {}  # module_id → 最近注册/订阅时间戳
+
+        # ── Slot token 管理（多连接支持） ──
+        self._slot_tokens: dict[str, tuple[str, int]] = {}     # token → (module_id, slot)
+        self._slot_token_expires: dict[str, float] = {}         # token → expire_ts
+
+        # ── 远程模块标记（通过 Relay 申请 token 连入，Launcher 不管理其生命周期） ──
+        self._remote_modules: set[str] = set()  # module_id
+
+    # ── 认证速率限制 ──
+
+    def check_auth_rate_limit(self, module_id_hint: str) -> bool:
+        """检查认证速率限制。
+
+        Args:
+            module_id_hint: 模块 ID 提示（来自 WS query param）
+
+        Returns:
+            True 表示允许，False 表示被限制
+        """
+        now = time.time()
+        cutoff = now - self._auth_rate_window
+
+        timestamps = self._auth_rate_limits.get(module_id_hint, [])
+        timestamps = [t for t in timestamps if t > cutoff]
+
+        if len(timestamps) >= self._auth_rate_max:
+            self._auth_rate_limits[module_id_hint] = timestamps
+            return False
+
+        timestamps.append(now)
+        self._auth_rate_limits[module_id_hint] = timestamps
+
+        # 定期清理不活跃的 key
+        if len(self._auth_rate_limits) > 500:
+            self._cleanup_auth_rate_limits(now)
+
+        return True
+
+    def _cleanup_auth_rate_limits(self, now: float):
+        """清理不活跃的速率限制条目"""
+        cutoff = now - self._auth_rate_window * 10
+        expired = [
+            k for k, ts_list in self._auth_rate_limits.items()
+            if not ts_list or ts_list[-1] < cutoff
+        ]
+        for k in expired:
+            del self._auth_rate_limits[k]
+
+    # ── Slot token 管理 ──
+
+    def generate_slot_tokens(self, module_id: str, slots: list[int]) -> dict[int, str]:
+        """为指定 slot 生成一次性 token，有效期 30 秒。
+
+        Args:
+            module_id: 模块 ID
+            slots: slot 编号列表（如 [1, 2]）
+
+        Returns:
+            {slot: token} 映射
+        """
+        import secrets
+        result = {}
+        expire_ts = time.time() + 30  # 30 秒有效期
+        for slot in slots:
+            token = secrets.token_hex(32)
+            self._slot_tokens[token] = (module_id, slot)
+            self._slot_token_expires[token] = expire_ts
+            result[slot] = token
+        return result
+
+    def verify_slot_token(self, token: str) -> tuple[str, int] | None:
+        """验证 slot token。通过后删除（一次性）。过期返回 None。
+
+        Returns:
+            (module_id, slot) 或 None
+        """
+        entry = self._slot_tokens.get(token)
+        if entry is None:
+            return None
+        expire_ts = self._slot_token_expires.get(token, 0)
+        if time.time() > expire_ts:
+            # 过期，清理
+            self._slot_tokens.pop(token, None)
+            self._slot_token_expires.pop(token, None)
+            return None
+        # 验证通过，删除（一次性）
+        self._slot_tokens.pop(token, None)
+        self._slot_token_expires.pop(token, None)
+        return entry
+
+    def cleanup_expired_slot_tokens(self):
+        """清理过期的 slot token，由 dedup_loop 定期调用。"""
+        now = time.time()
+        expired = [t for t, ts in self._slot_token_expires.items() if now > ts]
+        for t in expired:
+            self._slot_tokens.pop(t, None)
+            self._slot_token_expires.pop(t, None)
+
     # ── Token verification ──
 
     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,
         but known tokens still resolve to their proper module_id first."""
-        if token == self.launcher_token:
+        if hmac.compare_digest(token.encode(), self.launcher_token.encode()):
             return "launcher"
         for mid, tok in self.token_map.items():
-            if token == tok:
+            if hmac.compare_digest(token.encode(), tok.encode()):
                 return mid
         if self.is_debug and token:
             return "debug"
@@ -40,15 +146,30 @@ class RegistryStore:
 
     def is_launcher(self, token: str) -> bool:
         """Check if token belongs to Launcher. In debug mode, also accept any token."""
-        if token == self.launcher_token:
+        if hmac.compare_digest(token.encode(), self.launcher_token.encode()):
             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."""
+    def register_tokens(self, mapping: dict[str, str], remote: bool = False):
+        """Register module_id -> token mapping. Only Launcher may call this.
+
+        Args:
+            mapping: {module_id: token} pairs
+            remote: if True, marks all modules in mapping as remote (Relay clients).
+                    Remote modules are not managed by Launcher — no lifecycle control,
+                    no config file, no instance scaling.
+        """
         self.token_map.update(mapping)
+        if remote:
+            self._remote_modules.update(mapping.keys())
+
+    def is_remote_module(self, module_id: str) -> bool:
+        """Return True if module connected via Relay (not managed by Launcher)."""
+        # Strip instance suffix (e.g. 'kite_console-client-abc#2' -> 'kite_console-client-abc')
+        base = module_id.split("#")[0]
+        return base in self._remote_modules or module_id in self._remote_modules
 
     # ── Module lifecycle ──
 
@@ -109,6 +230,9 @@ class RegistryStore:
         record["registered_at"] = time.time()
         self.modules[mid] = record
 
+        # 更新模块初始化时间戳（用于 system.require_init 判断）
+        self._last_init_time[mid] = time.time()
+
         # Update global timestamp if content changed
         if changed:
             self.last_update_time = time.time()
@@ -129,20 +253,72 @@ class RegistryStore:
             if self.modules[module_id].get("status") == "offline":
                 self.modules[module_id]["status"] = "connected"
 
-    def set_ready(self, module_id: str):
-        """Mark a module as ready (received module.ready event)."""
+    def set_ready(self, module_id: str, data: dict = None):
+        """Mark a module as ready (received module.ready event).
+        Extracts and stores queue_elastic capability if declared."""
         if module_id in self.modules:
             self.modules[module_id]["status"] = "ready"
+            if data and data.get("queue_elastic"):
+                self.modules[module_id]["queue_elastic"] = True
+
+    def is_elastic(self, module_id: str) -> bool:
+        """Check if a module supports queue elastic management."""
+        mod = self.modules.get(module_id)
+        return mod is not None and mod.get("queue_elastic", False) is True
 
     def set_offline(self, module_id: str):
         """Mark a module as offline (called after debounce on WebSocket disconnect)."""
         if module_id in self.modules:
             self.modules[module_id]["status"] = "offline"
+            self.modules[module_id].pop("degraded_info", None)
 
     def is_ready(self, module_id: str) -> bool:
-        """Check if a module is in ready state."""
+        """Check if a module is in ready state (ready or degraded are both considered 'ready')."""
+        mod = self.modules.get(module_id)
+        return mod is not None and mod.get("status") in ("ready", "degraded")
+
+    def get_state(self, module_id: str) -> str:
+        """Get module state. Returns 'offline' if not found."""
         mod = self.modules.get(module_id)
-        return mod is not None and mod.get("status") == "ready"
+        return mod.get("status", "offline") if mod else "offline"
+
+    # ── 降级状态管理 ──
+
+    def set_degraded(self, module_id: str, info: dict):
+        """将模块标记为 degraded 状态，仅 ready 状态可进入。
+        info 包含 level, reason, affected, available, estimated_recovery_ms, detail 等字段。"""
+        mod = self.modules.get(module_id)
+        if not mod:
+            raise ValueError(f"Module not found: {module_id}")
+        if mod.get("status") not in ("ready", "degraded"):
+            raise ValueError(f"Module '{module_id}' is '{mod.get('status')}', only ready/degraded can degrade")
+        mod["status"] = "degraded"
+        mod["degraded_info"] = {
+            "level": info.get("level", "partial"),
+            "reason": info.get("reason", "unknown"),
+            "affected": info.get("affected", []),
+            "available": info.get("available", []),
+            "estimated_recovery_ms": info.get("estimated_recovery_ms"),
+            "detail": info.get("detail"),
+            "degraded_at": time.time(),
+        }
+
+    def clear_degraded(self, module_id: str):
+        """清除降级状态，恢复为 ready。仅 degraded 状态可调用。"""
+        mod = self.modules.get(module_id)
+        if not mod:
+            raise ValueError(f"Module not found: {module_id}")
+        if mod.get("status") != "degraded":
+            return  # 非降级状态，静默忽略
+        mod["status"] = "ready"
+        mod.pop("degraded_info", None)
+
+    def get_degraded_info(self, module_id: str) -> dict | None:
+        """获取模块降级详情，非降级状态返回 None。"""
+        mod = self.modules.get(module_id)
+        if not mod or mod.get("status") != "degraded":
+            return None
+        return mod.get("degraded_info")
 
     # ── Get by dot-path ──
 
@@ -195,7 +371,7 @@ class RegistryStore:
         """
         results = []
         for mid, data in self.modules.items():
-            if data.get("status") not in ("registered", "ready"):
+            if data.get("status") not in ("registered", "ready", "degraded"):
                 continue
             if module and not fnmatch.fnmatch(mid, module):
                 continue
@@ -276,3 +452,24 @@ class RegistryStore:
         if isinstance(val, (dict, list)):
             return False
         return fnmatch.fnmatch(str(val), pattern)
+
+    # ── 模块初始化时间戳管理（system.require_init） ──
+
+    def needs_init(self, module_id: str, timeout: float = 3600) -> bool:
+        """检查模块是否需要发送 system.require_init。
+
+        Args:
+            module_id: 模块 ID
+            timeout: 超时时间（秒），默认 1 小时
+
+        Returns:
+            True 表示需要发送 system.require_init
+        """
+        last_time = self._last_init_time.get(module_id)
+        if last_time is None:
+            return True
+        return (time.time() - last_time) > timeout
+
+    def update_init_time(self, module_id: str):
+        """更新模块初始化时间戳（订阅时调用）。"""
+        self._last_init_time[module_id] = time.time()