@agentunion/kite 1.5.0 → 1.6.1

package/extensions/agents/assistant/server.py

@@ -19,21 +19,36 @@ SYSTEM_BROADCAST_EVENTS = {
     "module.crashed", "module.exiting", "module.offline",
     "module.shutdown.ack", "module.shutdown.ready",
     "system.ready", "registry.updated",
+    "system.instance.started", "system.instance.stopped",
 }
 
 
 class AssistantServer:
 
-    def __init__(self, token: str = "", kernel_port: int = 0, boot_t0: float = 0):
+    def __init__(self, token: str = "", kernel_port: int = 0, boot_t0: float = 0, max_connections: int = 1,
+                 max_instances: int = 1, instance_elastic: bool = False,
+                 gateway_url: str = "", kite_token: str = ""):
         self.token = token
         self.kernel_port = kernel_port
         self.boot_t0 = boot_t0
+        self.max_connections = max_connections
+        self.max_instances = max_instances
+        self.instance_elastic = instance_elastic
+        self.gateway_url = gateway_url
+        self.kite_token = kite_token
         self._ws_task: asyncio.Task | None = None
         self._test_task: asyncio.Task | None = None
         self._ws: object | None = None
         self._shutting_down = False
         self._exit_code = 0  # Exit code for main() to use
+        self._auth_failed = False
         self._start_time = time.time()
+        self._extra_ws: dict = {}           # slot → WebSocket
+        self._extra_ws_tasks: dict = {}     # slot → recv loop Task
+        self._has_registered = False
+        self._pending_rpc: dict[str, asyncio.Future] = {}
+        self._active_handlers = 0   # 当前正在处理的入站请求数
+        self._handler_capacity = 50 # 并发处理能力上限（用于压力计算）
 
     async def run(self):
         """Main entry point: start WebSocket loop and test event loop."""
@@ -71,6 +86,7 @@ class AssistantServer:
                     if code in (4001, 4003):
                         print(f"[assistant] Kernel 认证失败 (code {code})，退出")
                         self._exit_code = 1
+                        self._auth_failed = True
                         self._shutting_down = True
                         return
                 if attempt >= max_retries:
@@ -87,20 +103,98 @@ class AssistantServer:
 
     async def _ws_connect(self):
         """Single WebSocket session: connect, subscribe, register, ready, receive loop."""
-        url = f"ws://127.0.0.1:{self.kernel_port}/ws?token={self.token}&id=assistant"
+        url = f"ws://127.0.0.1:{self.kernel_port}/ws?id=assistant"
         print(f"[assistant] Connecting to Kernel (port {self.kernel_port})")
         async with websockets.connect(url, open_timeout=3, ping_interval=None, close_timeout=10) as ws:
+            # Send auth message first
+            auth_req = {
+                "jsonrpc": "2.0",
+                "id": "auth",
+                "method": "auth",
+                "params": {"token": self.token}
+            }
+            await ws.send(json.dumps(auth_req))
+
+            # Wait for auth response
+            auth_resp_raw = await asyncio.wait_for(ws.recv(), timeout=5)
+            auth_resp = json.loads(auth_resp_raw)
+            if "error" in auth_resp:
+                raise Exception(f"Auth failed: {auth_resp['error']}")
+
             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"[assistant] Connected to Kernel{elapsed_str}")
 
-            # Step 1: Subscribe to events (先订阅)
+            # Start receiver task before subscribe/register so responses can be routed
+            receiver_task = asyncio.create_task(self._ws_receiver(ws))
+
+            try:
+                # subscribe/register/module.ready 由 _do_init 在收到 system.require_init 后执行
+
+                # Wait for receiver to finish (connection closed or error)
+                await receiver_task
+            finally:
+                receiver_task.cancel()
+                for fut in self._pending_rpc.values():
+                    if not fut.done():
+                        fut.set_result({"error": {"code": -32001, "message": "Connection lost"}})
+                self._pending_rpc.clear()
+
+    async def _ws_receiver(self, ws):
+        """Receive loop for main WebSocket connection.
+
+        CRITICAL: RPC 死锁防范
+        - 入站 RPC 请求必须用 create_task() 异步执行，不可 await
+        - 原因：如果 handler 内部调用 rpc_call() 发出站请求，出站响应需要本接收循环来分发
+        - 如果接收循环被 await handler 阻塞，出站响应永远收不到 → 超时死锁
+        - 事件通知和 RPC 响应可以同步处理（它们不会反向调用 rpc_call）
+        """
+        try:
+            async for raw in ws:
+                try:
+                    msg = json.loads(raw)
+                except (json.JSONDecodeError, TypeError):
+                    continue
+
+                try:
+                    has_method = "method" in msg
+                    has_id = "id" in msg
+
+                    # 检测 system.require_init 事件，触发注册流程
+                    if has_method and not has_id:
+                        params = msg.get("params", {})
+                        event = params.get("event", "")
+                        if event == "system.require_init":
+                            asyncio.create_task(self._do_init(ws))
+                            continue
+
+                    if has_method and not has_id:
+                        # JSON-RPC Notification (event delivery)
+                        # Run in background to prevent deadlock if handler calls RPC
+                        asyncio.create_task(self._handle_event_notification(msg))
+
+                    elif not has_method and has_id:
+                        # JSON-RPC Response (to our RPC calls)
+                        self._handle_rpc_response(msg)
+                except Exception as e:
+                    print(f"[assistant] 事件处理异常: {e}")
+        except websockets.exceptions.ConnectionClosed:
+            pass  # Kernel 连接关闭，正常退出接收循环
+
+    async def _do_init(self, ws):
+        """收到 system.require_init 后执行：订阅 + 注册 + module.ready"""
+        try:
+            reason = "startup" if not self._has_registered else "recovery"
+
+            # Step 1: Subscribe to events
             await self._rpc_call(ws, "event.subscribe", {
-                "events": ["module.started", "module.stopped", "module.shutdown"],
+                "events": [
+                    "module.started", "module.stopped", "module.shutdown",
+                ],
             })
 
-            # Step 2: Register to Kernel (再注册)
+            # Step 2: Register to Kernel
             await self._rpc_call(ws, "registry.register", {
                 "module_id": "assistant",
                 "module_type": "agent",
@@ -117,28 +211,88 @@ class AssistantServer:
                 ],
             })
 
-            # Step 3: Publish module.ready (every reconnect)
+            # Step 3: Publish module.ready
             if not self._shutting_down:
                 startup_time = time.monotonic() - self.boot_t0 if self.boot_t0 else 0
-                await self._rpc_call(ws, "event.publish", {
-                    "event_id": str(uuid.uuid4()),
-                    "event": "module.ready",
-                    "data": {
-                        "module_id": "assistant",
-                        "graceful_shutdown": True,
-                        "startup_time": startup_time,
-                    },
+                await self._publish_event(ws, "module.ready", {
+                    "module_id": "assistant",
+                    "graceful_shutdown": True,
+                    "startup_time": startup_time,
+                    "reason": reason,
                 })
                 elapsed_str = self._fmt_elapsed(self.boot_t0)
-                print(f"[assistant] module.ready published ({elapsed_str})")
+                print(f"[assistant] module.ready published (reason={reason}, {elapsed_str})")
+                self._has_registered = True
+        except Exception as e:
+            print(f"[assistant] _do_init failed: {e}")
+            self._has_registered = False
+
+    async def _handle_event_notification(self, msg: dict):
+        """Handle an event notification (JSON-RPC 2.0 Notification with method='event')."""
+        params = msg.get("params", {})
+        event_type = params.get("event", "")
+        data = params.get("data", {})
+
+        # Handle system.ping event
+        if event_type == "system.ping":
+            await self._handle_ping_event(data)
+            return
 
-            # Receive loop
-            # CRITICAL: RPC 死锁防范
-            # - 入站 RPC 请求必须用 create_task() 异步执行，不可 await
-            # - 原因：如果 handler 内部调用 rpc_call() 发出站请求，出站响应需要本接收循环来分发
-            # - 如果接收循环被 await handler 阻塞，出站响应永远收不到 → 超时死锁
-            # - 事件通知和 RPC 响应可以同步处理（它们不会反向调用 rpc_call）
+        # 弹性连接 offer/release
+        if event_type == "system.connection.offer":
+            asyncio.create_task(self._handle_connection_offer(data))
+            return
+        if event_type == "system.connection.release":
+            asyncio.create_task(self._handle_connection_release(data))
+            return
 
+        # Layer 1: 处理订阅的事件
+        if event_type == "module.shutdown":
+            target = data.get("module_id", "")
+            reason = data.get("reason", "")
+            # Handle both targeted shutdown (module_id == "assistant") and broadcast shutdown (no module_id or launcher_lost)
+            if target.startswith("assistant") or not target or reason == "launcher_lost":
+                await self._handle_shutdown(self._ws)
+                return
+
+        # Layer 2: 忽略系统广播事件
+        if event_type in SYSTEM_BROADCAST_EVENTS:
+            return
+
+        # Layer 3: 警告未知事件（仅开发环境）
+        if os.environ.get("KITE_ENV") == "development":
+            print(f"[assistant] Debug: Unhandled event: {event_type}")
+
+    async def _handle_connection_offer(self, data):
+        """处理 Kernel 下发的 slot token，建立附加连接。"""
+        slots = data.get("slots", {})
+        for slot_str, info in slots.items():
+            slot = int(slot_str)
+            token = info.get("token", "")
+            if not token or slot in self._extra_ws:
+                continue
+            asyncio.create_task(self._connect_slot(slot, token))
+
+    async def _connect_slot(self, slot, token):
+        """建立单个 slot 附加连接。"""
+        ws_url = f"ws://127.0.0.1:{self.kernel_port}/ws"
+        try:
+            ws = await websockets.connect(ws_url, open_timeout=5, ping_interval=None, close_timeout=5)
+            auth_req = {"jsonrpc": "2.0", "id": f"auth-slot-{slot}", "method": "auth", "params": {"token": token}}
+            await ws.send(json.dumps(auth_req))
+            resp = json.loads(await asyncio.wait_for(ws.recv(), timeout=5))
+            if "error" in resp:
+                await ws.close()
+                return
+            self._extra_ws[slot] = ws
+            self._extra_ws_tasks[slot] = asyncio.create_task(self._slot_recv_loop(slot, ws))
+            print(f"[assistant] Slot {slot} connected")
+        except Exception as e:
+            print(f"[assistant] Slot {slot} connect failed: {e}")
+
+    async def _slot_recv_loop(self, slot, ws):
+        """附加 slot 接收循环：与主连接平等处理所有消息。"""
+        try:
             async for raw in ws:
                 try:
                     msg = json.loads(raw)
@@ -150,95 +304,125 @@ class AssistantServer:
                     has_id = "id" in msg
 
                     if has_method and not has_id:
-                        # JSON-RPC Notification (event delivery)
-                        params = msg.get("params", {})
-                        event_type = params.get("event", "")
-                        data = params.get("data", {})
-
-                        # Handle system.ping event
-                        if event_type == "system.ping":
-                            await self._handle_ping_event(data)
-                            continue
-
-                        # Layer 1: 处理订阅的事件
-                        if event_type == "module.shutdown":
-                            target = data.get("module_id", "")
-                            reason = data.get("reason", "")
-                            # Handle both targeted shutdown (module_id == "assistant") and broadcast shutdown (no module_id or launcher_lost)
-                            if target == "assistant" or not target or reason == "launcher_lost":
-                                await self._handle_shutdown(ws)
-                                return
-
-                        # Layer 2: 忽略系统广播事件
-                        if event_type in SYSTEM_BROADCAST_EVENTS:
-                            continue
-
-                        # Layer 3: 警告未知事件（仅开发环境）
-                        if os.environ.get("KITE_ENV") == "development":
-                            print(f"[assistant] Debug: Unhandled event: {event_type}")
-
+                        asyncio.create_task(self._handle_event_notification(msg))
                     elif not has_method and has_id:
-                        # JSON-RPC Response (to our RPC calls)
-                        pass
+                        self._handle_rpc_response(msg)
                 except Exception as e:
-                    print(f"[assistant] 事件处理异常: {e}")
+                    print(f"[assistant] Slot {slot} 消息处理异常（已忽略）: {e}")
+        except Exception as e:
+            print(f"[assistant] Slot {slot} recv loop exited: {e}")
+        finally:
+            self._extra_ws.pop(slot, None)
+            self._extra_ws_tasks.pop(slot, None)
+
+    async def _handle_connection_release(self, data):
+        """Kernel 请求释放 slot，优雅关闭。"""
+        for slot in data.get("slots", []):
+            ws = self._extra_ws.pop(slot, None)
+            task = self._extra_ws_tasks.pop(slot, None)
+            if ws:
+                try:
+                    await ws.close(code=1000, reason="release")
+                except Exception:
+                    pass
+            if task:
+                task.cancel()
 
     async def _handle_shutdown(self, ws):
-        """Handle module.shutdown: ack → exiting → cleanup → ready → exit."""
+        """Handle module.shutdown: ack → exiting → ready → cleanup → exit."""
         print("[assistant] Received module.shutdown")
         self._shutting_down = True
 
         # Step 1: Send ack (立即确认收到)
-        await self._rpc_call(ws, "event.publish", {
-            "event_id": str(uuid.uuid4()),
-            "event": "module.shutdown.ack",
-            "data": {"module_id": "assistant", "estimated_cleanup": 2},
+        await self._publish_event(ws, "module.shutdown.ack", {
+            "module_id": "assistant", "estimated_cleanup": 2,
         })
         print("[assistant] shutdown ack sent")
 
         # Step 2: Send module.exiting (开始清理)
-        await self._rpc_call(ws, "event.publish", {
-            "event_id": str(uuid.uuid4()),
-            "event": "module.exiting",
-            "data": {"module_id": "assistant", "action": "none"},
+        await self._publish_event(ws, "module.exiting", {
+            "module_id": "assistant", "action": "none",
         })
 
-        # Step 3: Cleanup (cancel background tasks)
-        if self._test_task:
-            self._test_task.cancel()
-
-        # Step 4: Send ready (清理完成)
-        await self._rpc_call(ws, "event.publish", {
-            "event_id": str(uuid.uuid4()),
-            "event": "module.shutdown.ready",
-            "data": {"module_id": "assistant"},
+        # Step 3: Send ready（必须在关闭连接之前发送）
+        await self._publish_event(ws, "module.shutdown.ready", {
+            "module_id": "assistant",
         })
         print("[assistant] Shutdown ready sent")
 
-        # Step 4: Exit process
-        import sys
-        sys.exit(0)
+        # Step 4: 等待 Kernel 处理 shutdown.ready
+        await asyncio.sleep(0.01)
+
+        # Step 5: Cleanup - 关闭所有附加连接
+        for _s, _w in list(self._extra_ws.items()):
+            try:
+                await _w.close(code=1000, reason="shutdown")
+            except Exception:
+                pass
+        for _t in self._extra_ws_tasks.values():
+            _t.cancel()
+        self._extra_ws.clear()
+        self._extra_ws_tasks.clear()
+
+        if self._test_task:
+            self._test_task.cancel()
 
-    async def _rpc_call(self, ws, method: str, params: dict = None):
-        """Send a JSON-RPC 2.0 request."""
-        msg = {"jsonrpc": "2.0", "id": str(uuid.uuid4()), "method": method}
+        # Step 6: Close WebSocket connection gracefully
+        try:
+            await ws.close(code=1000, reason="Graceful shutdown")
+            print("[assistant] WebSocket closed")
+        except Exception as e:
+            print(f"[assistant] Failed to close WebSocket: {e}")
+
+    async def _rpc_call(self, ws, method: str, params: dict = None,
+                        wait_response: bool = True, timeout: float = 3.0) -> dict:
+        """JSON-RPC 2.0 request。默认等待响应。"""
+        rpc_id = str(uuid.uuid4())
+        msg = {"jsonrpc": "2.0", "id": rpc_id, "method": method}
         if params:
             msg["params"] = params
+        if not wait_response:
+            await ws.send(json.dumps(msg))
+            return
+        future = asyncio.get_event_loop().create_future()
+        self._pending_rpc[rpc_id] = future
         await ws.send(json.dumps(msg))
+        try:
+            return await asyncio.wait_for(future, timeout=timeout)
+        except asyncio.TimeoutError:
+            self._pending_rpc.pop(rpc_id, None)
+            return {"error": {"code": -32000, "message": f"RPC timeout: {method} ({timeout}s)"}}
+
+    def _compute_pressure_level(self) -> str:
+        """根据并发处理量计算压力等级。"""
+        depth = self._active_handlers + len(self._pending_rpc)
+        ratio = depth / self._handler_capacity if self._handler_capacity > 0 else 0
+        if ratio >= 0.9:
+            return "critical"
+        elif ratio >= 0.5:
+            return "overload"
+        elif depth > 0:
+            return "normal"
+        return "idle"
+
+    def _handle_rpc_response(self, msg: dict):
+        """Resolve pending RPC future with the response."""
+        rpc_id = msg.get("id")
+        future = self._pending_rpc.pop(rpc_id, None)
+        if future and not future.done():
+            future.set_result(msg)
 
     async def _handle_ping_event(self, data: dict):
         """Handle system.ping event and reply with system.pong."""
+        if self._ws is None:
+            return
         t1 = data.get("ping_time")
         t2 = time.time()
 
-        await self._rpc_call(self._ws, "event.publish", {
-            "event_id": str(uuid.uuid4()),
-            "event": "system.pong",
-            "data": {
-                "module_id": "assistant",
-                "ping_time": t1,
-                "pong_time": t2,
-            },
+        await self._publish_event(self._ws, "system.pong", {
+            "module_id": "assistant",
+            "ping_time": t1,
+            "pong_time": t2,
         })
 
     def _fmt_elapsed(self, t0: float) -> str:
@@ -250,30 +434,32 @@ class AssistantServer:
             return f"{d:.1f}s"
         return f"{d:.0f}s"
 
-    async def _publish_event(self, event: dict):
-        """Publish an event via JSON-RPC event.publish."""
-        if not self._ws:
-            return
-        try:
-            await self._rpc_call(self._ws, "event.publish", {
-                "event_id": str(uuid.uuid4()),
-                "event": event.get("event", ""),
-                "data": event.get("data", {}),
-            })
-        except Exception as e:
-            print(f"[assistant] Failed to publish event: {e}")
+    async def _publish_event(self, ws, event: str, data: dict = None):
+        """Publish event via event.publish RPC (fire-and-forget)."""
+        await self._rpc_call(ws, "event.publish", {
+            "event_id": str(uuid.uuid4()),
+            "event": event,
+            "data": data or {},
+        }, wait_response=False)
 
     # ── 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": "assistant.test",
-                "data": {
-                    "message": "test event from assistant",
-                    "timestamp": datetime.now(timezone.utc).isoformat(),
-                },
-            })
-            print("[assistant] test event published")
+        try:
+            while True:
+                await asyncio.sleep(10)
+                if self._ws:
+                    await self._publish_event(self._ws, "assistant.test", {
+                            "message": "test event from assistant",
+                            "timestamp": datetime.now(timezone.utc).isoformat(),
+                        })
+                    # 上报队列压力（供 Kernel 自动缩容决策）
+                    depth = self._active_handlers + len(self._pending_rpc)
+                    await self._publish_event(self._ws, "module.queue_pressure", {
+                        "level": self._compute_pressure_level(),
+                        "depth": depth,
+                        "capacity": self._handler_capacity,
+                    })
+        except Exception:
+            pass

package/extensions/channels/acp_channel/entry.py

@@ -116,6 +116,12 @@ def _load_module_config() -> dict:
         if "advertise_ip" in fm:
             result["advertise_ip"] = str(fm["advertise_ip"])
 
+        # max_connections（弹性连接上限）
+        try:
+            result["max_connections"] = max(1, min(10, int(fm.get("max_connections", 1))))
+        except (ValueError, TypeError):
+            result["max_connections"] = 1
+
     except SystemExit:
         raise  # Re-raise exit to prevent catching by outer except
     except Exception as e:
@@ -364,7 +370,7 @@ def _read_stdin_kite_message(expected_type: str, timeout: float = 10) -> dict |
     return result[0]
 
 
-def main():
+def _run_local_mode(kernel_port: int):
     # Initialize log file paths
     global _log_dir, _log_latest_path, _crash_log_path
     module_data = os.environ.get("KITE_MODULE_DATA")
@@ -394,35 +400,49 @@ def main():
     _t0 = time.monotonic()
 
     # Read boot_info from stdin (only token)
-    token = ""
+    line = sys.stdin.readline().strip()
+    if not line:
+        print(f"[{MODULE_NAME}] ERROR: stdin closed")
+        sys.exit(1)
+
     try:
-        line = sys.stdin.readline().strip()
-        if line:
-            boot_info = json.loads(line)
-            token = boot_info.get("token", "")
-    except Exception:
-        pass
-
-    # Read kernel_port: env first (fast path), stdin fallback (parallel start)
-    kernel_port = int(os.environ.get("KITE_KERNEL_PORT", "0"))
+        msg = json.loads(line)
+    except json.JSONDecodeError as e:
+        print(f"[{MODULE_NAME}] ERROR: Invalid JSON from stdin: {e}")
+        sys.exit(1)
+
+    if "error" in msg:
+        print(f"[{MODULE_NAME}] 启动失败: {msg.get('message', 'unknown')}")
+        sys.exit(1)
+
+    token = msg.get("token", "")
+    max_instances = msg.get("max_instances", 1)
+    instance_elastic = msg.get("instance_elastic", False)
+    if not token:
+        print(f"[{MODULE_NAME}] ERROR: No token in stdin message")
+        sys.exit(1)
+
     if not kernel_port:
         msg = _read_stdin_kite_message("kernel_port", timeout=10)
         if msg:
             kernel_port = int(msg.get("kernel_port", 0))
 
-    if not token or not kernel_port:
-        print("[acp_channel] ERROR: Missing token or kernel_port")
+    if not kernel_port:
+        print(f"[{MODULE_NAME}] ERROR: Missing kernel_port")
         sys.exit(1)
 
-    print(f"[acp_channel] Token received ({len(token)} chars), kernel port: {kernel_port} ({_fmt_elapsed(_t0)})")
+    print(f"[{MODULE_NAME}] Local mode: port={kernel_port} ({_fmt_elapsed(_t0)})")
 
     server = AcpChannelServer(
         token=token,
         kernel_port=kernel_port,
         boot_t0=_t0,
+        max_connections=_module_config.get("max_connections", 1),
+        max_instances=max_instances,
+        instance_elastic=instance_elastic,
     )
 
-    print(f"[acp_channel] Starting ({_fmt_elapsed(_t0)})")
+    print(f"[{MODULE_NAME}] Starting ({_fmt_elapsed(_t0)})")
     try:
         asyncio.run(server.run())
     except Exception as e:
@@ -435,5 +455,83 @@ def main():
         sys.exit(server._exit_code)
 
 
+def _run_remote_mode(gateway_url: str):
+    global _log_dir, _log_latest_path, _crash_log_path
+    home = os.environ.get("HOME") or os.environ.get("USERPROFILE") or os.path.expanduser("~")
+    data_dir = os.path.join(home, ".kite", "remote", MODULE_NAME)
+    os.makedirs(data_dir, exist_ok=True)
+    _log_dir = os.path.join(data_dir, "log")
+    os.makedirs(_log_dir, exist_ok=True)
+    _log_latest_path = os.path.join(_log_dir, "latest.log")
+    try:
+        with open(_log_latest_path, "w", encoding="utf-8") as f: pass
+    except Exception: _log_latest_path = None
+    _crash_log_path = os.path.join(_log_dir, "crashes.jsonl")
+    try:
+        with open(_crash_log_path, "w", encoding="utf-8") as f: pass
+    except Exception: _crash_log_path = None
+    _resolve_daily_log_path()
+    _setup_exception_hooks()
+    _t0 = time.monotonic()
+    kite_token = _get_kite_token(MODULE_NAME, gateway_url)
+    print(f"[{MODULE_NAME}] Remote mode: gateway={gateway_url}")
+    server = AcpChannelServer(token=None, kernel_port=0, boot_t0=_t0,
+                              max_connections=_module_config.get("max_connections", 1),
+                              gateway_url=gateway_url, kite_token=kite_token)
+    try:
+        asyncio.run(server.run())
+    except Exception as e:
+        _write_crash(type(e), e, e.__traceback__, severity="critical", handled=True)
+        _print_crash_summary(type(e), e.__traceback__)
+        sys.exit(1)
+    if hasattr(server, '_auth_failed') and server._auth_failed:
+        _clear_token_cache(MODULE_NAME, gateway_url)
+    if server._exit_code != 0:
+        sys.exit(server._exit_code)
+
+
+def _gateway_to_filename(gateway_url: str) -> str:
+    try:
+        from urllib.parse import urlparse
+        parsed = urlparse(gateway_url)
+        host = parsed.hostname or "unknown"
+        port = parsed.port or (443 if parsed.scheme == "wss" else 80)
+        return f"{host}-{port}.json".replace(":", "-").replace("/", "-")
+    except Exception: return "default.json"
+
+def _load_token_cache(module_name: str, gateway_url: str) -> dict | None:
+    home = os.environ.get("HOME") or os.environ.get("USERPROFILE") or os.path.expanduser("~")
+    token_file = os.path.join(home, ".kite", "remote", module_name, "tokens", _gateway_to_filename(gateway_url))
+    if not os.path.exists(token_file): return None
+    try:
+        with open(token_file, "r") as f: return json.load(f)
+    except Exception: return None
+
+def _clear_token_cache(module_name: str, gateway_url: str):
+    home = os.environ.get("HOME") or os.environ.get("USERPROFILE") or os.path.expanduser("~")
+    token_file = os.path.join(home, ".kite", "remote", module_name, "tokens", _gateway_to_filename(gateway_url))
+    try: os.remove(token_file)
+    except Exception: pass
+
+def _get_kite_token(module_name: str, gateway_url: str) -> str:
+    token = os.environ.get("KITE_TOKEN")
+    if token: return token
+    cache = _load_token_cache(module_name, gateway_url)
+    if cache: return cache["token"]
+    print(f"[{module_name}] No token for {gateway_url}")
+    print(f"  export KITE_TOKEN=<token>")
+    sys.exit(1)
+
+
 if __name__ == "__main__":
-    main()
+    kernel_port = os.environ.get("KITE_KERNEL_PORT")
+    if kernel_port:
+        _run_local_mode(int(kernel_port))
+    else:
+        config = _load_module_config()
+        gateway_url = config.get("gateway_url") or os.environ.get("KITE_GATEWAY_URL")
+        if gateway_url:
+            _run_remote_mode(gateway_url)
+        else:
+            print(f"[{MODULE_NAME}] ERROR: No KITE_KERNEL_PORT and no gateway_url")
+            sys.exit(1)