@agentunion/kite 1.5.0 → 1.6.1

package/extensions/services/backup/entry.py

@@ -15,6 +15,7 @@ from datetime import datetime, timezone
 import asyncio
 import traceback
 import uuid
+import random
 
 import websockets
 
@@ -25,6 +26,7 @@ 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",
 }
 
 
@@ -129,6 +131,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:
@@ -381,6 +389,14 @@ _ws_global = None
 _shutting_down = False
 _exit_code = 0  # Exit code for main() to use
 
+# 弹性多连接
+_extra_ws: dict = {}           # slot → WebSocket（附加连接）
+_extra_ws_tasks: dict = {}     # slot → recv loop Task
+_kernel_port = ""              # 缓存 kernel_port，供 offer handler 使用
+_pending_rpc: dict[str, asyncio.Future] = {}  # rpc_id → Future（等待响应）
+
+_has_registered = False
+
 
 def _is_auth_failure(e: Exception) -> bool:
     """Check if a WebSocket exception indicates authentication failure."""
@@ -392,212 +408,251 @@ def _is_auth_failure(e: Exception) -> bool:
 
 async def main():
     global _ws_global, _shutting_down
-    # Initialize log file paths
-    global _log_dir, _log_latest_path, _crash_log_path
-    module_data = os.environ.get("KITE_MODULE_DATA")
-    if module_data:
-        _log_dir = os.path.join(module_data, "log")
-        os.makedirs(_log_dir, exist_ok=True)
-        suffix = os.environ.get("KITE_INSTANCE_SUFFIX", "")
-
-        _log_latest_path = os.path.join(_log_dir, f"latest{suffix}.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, f"crashes{suffix}.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()
-
-    # Read boot_info from stdin (only token)
-    token = ""
-    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"))
-    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("[backup] ERROR: Missing token or kernel_port")
-        sys.exit(1)
-
-    print(f"[backup] Token received ({len(token)} chars), kernel port: {kernel_port} ({_fmt_elapsed(_t0)})")
-
-    # Start reconnect loop
-    await _ws_loop(token, kernel_port, _t0)
+    if _startup_config.get("mode") == "remote":
+        await _remote_mode_loop(_startup_config["gateway_url"], _startup_config["kite_token"], _startup_config["t0"])
+    else:
+        await _ws_loop(_startup_config["token"], _startup_config["kernel_port"], _startup_config["t0"])
 
 
 async def _ws_loop(token: str, kernel_port: int, _t0: float):
     """Connect to Kernel with exponential backoff reconnection."""
     global _shutting_down, _exit_code
-    retry_delay = 0.3
-    max_delay = 5.0
-    max_retries = 10
+    retry_delay = 0.5
+    max_delay = 10.0
     attempt = 0
+    cooldown_attempts = 0
+
     while not _shutting_down:
         try:
             await _ws_connect(token, kernel_port, _t0)
-            retry_delay = 0.3
+            retry_delay = 0.5
             attempt = 0
+            cooldown_attempts = 0
         except asyncio.CancelledError:
             return
         except Exception as e:
-            attempt += 1
-            if _is_auth_failure(e):
-                print(f"[backup] Kernel 认证失败，退出")
-                _exit_code = 1
-                _shutting_down = True
+            if _shutting_down:
                 return
-            if attempt >= max_retries:
-                print(f"[backup] 重连失败 {max_retries} 次，退出")
+
+            code = _get_close_code(e)
+
+            # never: 永不重连
+            if code in (4001, 4003, 4004, 1008, 4010):
+                print(f"[backup] 致命错误 (code {code})，退出")
                 _exit_code = 1
                 _shutting_down = True
                 return
+
+            # cooldown: 速率限制
+            if code == 4020:
+                cooldown_attempts += 1
+                if cooldown_attempts >= 5:
+                    print(f"[backup] 速率限制重试 5 次，退出")
+                    _exit_code = 1
+                    _shutting_down = True
+                    return
+                print(f"[backup] 速率限制，10.0s 后重试 ({cooldown_attempts}/5)")
+                await asyncio.sleep(10.0)
+                continue
+
+            # standard: 指数退避 + jitter
+            attempt += 1
+            jitter = retry_delay * 0.2 * random.random()
+            sleep_time = retry_delay + jitter
             _write_crash(type(e), e, e.__traceback__, severity="error", handled=True)
-            print(f"[backup] 连接错误: {e}, {retry_delay:.1f}s 后重试 ({attempt}/{max_retries})")
+            print(f"[backup] 连接错误: {e}, {sleep_time:.1f}s 后重试 (attempt {attempt})")
+
         _ws_global_clear()
         if _shutting_down:
             return
-        await asyncio.sleep(retry_delay)
+        await asyncio.sleep(sleep_time if 'sleep_time' in locals() else retry_delay)
         retry_delay = min(retry_delay * 2, max_delay)
 
 
+def _get_close_code(e: Exception) -> int:
+    """从 websockets 异常中提取关闭码"""
+    if hasattr(e, 'rcvd') and e.rcvd is not None:
+        return getattr(e.rcvd, 'code', 0)
+    return 0
+
+
 def _ws_global_clear():
     global _ws_global
     _ws_global = None
 
 
+async def _do_init(ws):
+    """收到 system.require_init 后执行：subscribe + register + module.ready"""
+    global _has_registered
+    reason = "startup" if not _has_registered else "recovery"
+    print(f"[backup] Received system.require_init (reason={reason})")
+
+    # 1. Subscribe to events
+    await _rpc_call(ws, "event.subscribe", {
+        "events": [
+            "module.started",
+            "module.stopped",
+            "module.shutdown",
+        ],
+    })
+    print(f"[backup] Subscribed to events")
+
+    # 2. Register to Kernel Registry via RPC
+    await _rpc_call(ws, "registry.register", {
+        "module_id": "backup",
+        "module_type": "service",
+        "launcher_id": "launcher",
+        "tools": {
+            "rpc": {
+                "module": {
+                    "health": {"method": "health", "description": "健康检查"},
+                    "status": {"method": "status", "description": "状态查询"},
+                    "config": {
+                        "get": {"method": "get_settings", "description": "获取配置"},
+                        "update": {"method": "update_settings", "description": "更新配置"},
+                        "reset": {"method": "reset_settings", "description": "恢复默认配置"},
+                    }
+                }
+            }
+        },
+        "events_publish": {
+            "backup": {
+                "test": {"description": "Test event from backup module"},
+            }
+        },
+        "events_subscribe": [
+            "module.started",
+            "module.stopped",
+            "module.shutdown",
+        ],
+    })
+    print(f"[backup] Registered to Kernel")
+
+    # 3. Publish module.ready
+    if not _shutting_down:
+        startup_time = time.monotonic() - _start_ts
+        await _publish_event(ws, "module.ready", {
+            "module_id": "backup",
+            "graceful_shutdown": True,
+            "startup_time": startup_time,
+            "reason": reason,
+        })
+        print(f"[backup] module.ready published (reason={reason})")
+
+    _has_registered = True
+
+    # Start test event loop in background
+    asyncio.create_task(_test_event_loop(ws))
+
+
 async def _ws_connect(token: str, kernel_port: int, _t0: float):
-    """Single WebSocket session: connect → subscribe → register → ready → receive loop."""
-    global _ws_global
+    """Single WebSocket session: connect → auth → start receiver → wait for system.require_init."""
+    global _ws_global, _kernel_port
+    _kernel_port = kernel_port
 
-    ws_url = f"ws://127.0.0.1:{kernel_port}/ws?token={token}&id=backup"
+    ws_url = f"ws://127.0.0.1:{kernel_port}/ws?id=backup"
     print(f"[backup] Connecting to Kernel: {ws_url}")
 
     async with websockets.connect(ws_url, open_timeout=5, ping_interval=None, close_timeout=10) as ws:
+        # Send auth message first
+        auth_req = {
+            "jsonrpc": "2.0",
+            "id": "auth",
+            "method": "auth",
+            "params": {"token": 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']}")
+
         _ws_global = ws
         print(f"[backup] Connected to Kernel ({_fmt_elapsed(_t0)})")
 
-        # Subscribe to events
-        await _rpc_call(ws, "event.subscribe", {
-            "events": [
-                "module.started",
-                "module.stopped",
-                "module.shutdown",
-            ],
-        })
-        print(f"[backup] Subscribed to events ({_fmt_elapsed(_t0)})")
+        # Start receiver task — wait for system.require_init to trigger _do_init
+        receiver_task = asyncio.create_task(_ws_receiver(ws))
+        try:
+            # Wait for receiver to finish (connection closed or error)
+            await receiver_task
+        finally:
+            receiver_task.cancel()
+            for fut in _pending_rpc.values():
+                if not fut.done():
+                    fut.set_result({"error": {"code": -32001, "message": "Connection lost"}})
+            _pending_rpc.clear()
+
+
+async def _ws_receiver(ws):
+    """Main connection receive loop: route events, RPC requests, and RPC responses.
+    Detects system.require_init to trigger _do_init."""
+    async for raw in ws:
+        try:
+            msg = json.loads(raw)
+        except (json.JSONDecodeError, TypeError):
+            continue
 
-        # Register to Kernel Registry via RPC
-        await _rpc_call(ws, "registry.register", {
-            "module_id": "backup",
-            "module_type": "service",
-            "launcher_id": "launcher",  # 声明归属的 Launcher
-            "tools": {
-                "rpc": {
-                    "module": {
-                        "health": {"method": "health", "description": "健康检查"},
-                        "status": {"method": "status", "description": "状态查询"},
-                        "config": {
-                            "get": {"method": "get_settings", "description": "获取配置"},
-                            "update": {"method": "update_settings", "description": "更新配置"},
-                            "reset": {"method": "reset_settings", "description": "恢复默认配置"},
-                        }
-                    }
-                }
-            },
-            "events_publish": {
-                "backup": {
-                    "test": {"description": "Test event from backup module"},
-                }
-            },
-            "events_subscribe": [
-                "module.started",
-                "module.stopped",
-                "module.shutdown",
-            ],
-        })
-        print(f"[backup] Registered to Kernel ({_fmt_elapsed(_t0)})")
-
-        # Publish module.ready (every reconnect)
-        if not _shutting_down:
-            startup_time = time.monotonic() - _t0
-            await _rpc_call(ws, "event.publish", {
-                "event_id": str(uuid.uuid4()),
-                "event": "module.ready",
-                "data": {
-                    "module_id": "backup",
-                    "graceful_shutdown": True,
-                    "startup_time": startup_time,
-                },
-            })
-            print(f"[backup] module.ready published ({_fmt_elapsed(_t0)})")
+        try:
+            has_method = "method" in msg
+            has_id = "id" in msg
+
+            if has_method and not has_id:
+                # 检测 system.require_init 事件
+                params = msg.get("params", {})
+                event_type = params.get("event", "")
+                if event_type == "system.require_init":
+                    asyncio.create_task(_do_init(ws))
+                    continue
+                # Event Notification — run in background to prevent deadlock
+                asyncio.create_task(_handle_event_notification(msg))
+            elif has_method and has_id:
+                # Incoming RPC request — run in background to prevent deadlock
+                asyncio.create_task(_handle_rpc_request(ws, msg))
+            elif not has_method and has_id:
+                _handle_rpc_response(msg)
+        except Exception as e:
+            print(f"[backup] 消息处理异常（已忽略）: {e}")
 
-        # Start test event loop in background
-        test_task = asyncio.create_task(_test_event_loop(ws))
 
-        # Message loop: handle incoming RPC + events
-        # CRITICAL: RPC 死锁防范
-        # - 入站 RPC 请求必须用 create_task() 异步执行，不可 await
-        # - 原因：如果 handler 内部调用 rpc_call() 发出站请求，出站响应需要本接收循环来分发
-        # - 如果接收循环被 await handler 阻塞，出站响应永远收不到 → 超时死锁
-        # - 事件通知和 RPC 响应可以同步处理（它们不会反向调用 rpc_call）
-        async for raw in ws:
-            try:
-                msg = json.loads(raw)
-            except (json.JSONDecodeError, TypeError):
-                continue
+async def _rpc_call(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
 
-            try:
-                has_method = "method" in msg
-                has_id = "id" in msg
+    if not wait_response:
+        await ws.send(json.dumps(msg))
+        return {}
 
-                if has_method and not has_id:
-                    # Event Notification
-                    await _handle_event_notification(msg)
-                elif has_method and has_id:
-                    # Incoming RPC request — run in background to prevent deadlock
-                    asyncio.create_task(_handle_rpc_request(ws, msg))
-                # Ignore RPC responses (we don't await them in this simple impl)
-            except Exception as e:
-                print(f"[backup] 消息处理异常（已忽略）: {e}")
+    future = asyncio.get_event_loop().create_future()
+    _pending_rpc[rpc_id] = future
+    await ws.send(json.dumps(msg))
+    try:
+        return await asyncio.wait_for(future, timeout=timeout)
+    except asyncio.TimeoutError:
+        _pending_rpc.pop(rpc_id, None)
+        return {"error": {"code": -32000, "message": f"RPC timeout: {method} ({timeout}s)"}}
 
 
-async def _rpc_call(ws, method: str, params: dict = None):
-    """Send a JSON-RPC 2.0 request (fire-and-forget, no response awaited)."""
-    msg = {"jsonrpc": "2.0", "id": str(uuid.uuid4()), "method": method}
-    if params:
-        msg["params"] = params
-    await ws.send(json.dumps(msg))
+def _handle_rpc_response(msg: dict):
+    """Route JSON-RPC response to pending future."""
+    rpc_id = msg.get("id")
+    future = _pending_rpc.pop(rpc_id, None)
+    if future and not future.done():
+        future.set_result(msg)
 
 
-async def _publish_event(ws, event: dict):
-    """Publish an event via RPC event.publish."""
+async def _publish_event(ws, event: str, data: dict = None):
+    """Publish event via event.publish RPC (fire-and-forget)."""
     await _rpc_call(ws, "event.publish", {
         "event_id": str(uuid.uuid4()),
-        "event": event.get("event", ""),
-        "data": event.get("data", {}),
-    })
+        "event": event,
+        "data": data or {},
+    }, wait_response=False)
 
 
 async def _handle_ping_event(data: dict):
@@ -605,13 +660,10 @@ async def _handle_ping_event(data: dict):
     t1 = data.get("ping_time")
     t2 = time.time()
 
-    await _publish_event(_ws_global, {
-        "event": "system.pong",
-        "data": {
-            "module_id": MODULE_NAME,
-            "ping_time": t1,
-            "pong_time": t2,
-        },
+    await _publish_event(_ws_global, "system.pong", {
+        "module_id": MODULE_NAME,
+        "ping_time": t1,
+        "pong_time": t2,
     })
 
 
@@ -626,6 +678,14 @@ async def _handle_event_notification(msg: dict):
         await _handle_ping_event(data)
         return
 
+    # 弹性连接 offer/release
+    if event_type == "system.connection.offer":
+        asyncio.create_task(_handle_connection_offer(data))
+        return
+    if event_type == "system.connection.release":
+        asyncio.create_task(_handle_connection_release(data))
+        return
+
     # Special handling for module.shutdown
     if event_type == "module.shutdown":
         target = data.get("module_id", "")
@@ -741,7 +801,7 @@ async def _rpc_get_settings(params: dict) -> dict:
         "events": frontmatter.get("events"),
         "subscriptions": frontmatter.get("subscriptions"),
         "depends_on": frontmatter.get("depends_on"),
-        "has_config": False,  # backup 模块暂无 config.yaml
+        "has_config": False,  # backup 模块暂无业务配置
         "config": None,
     }
 
@@ -758,9 +818,9 @@ async def _rpc_update_settings(params: dict) -> dict:
             frontmatter[key] = value
         _write_module_md(frontmatter, body)
 
-    # 更新 config.yaml（如果需要）
+    # 更新业务配置（如果需要）
     if config:
-        # backup 模块暂无 config.yaml，此处预留接口
+        # backup 模块暂无业务配置，此处预留接口
         pass
 
     # 返回更新后的完整设置
@@ -799,53 +859,290 @@ async def _rpc_reset_settings(params: dict) -> dict:
 
 
 
+async def _handle_connection_offer(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 _extra_ws:
+            continue
+        asyncio.create_task(_connect_slot(slot, token))
+
+
+async def _connect_slot(slot, token):
+    """建立单个 slot 附加连接。"""
+    ws_url = f"ws://127.0.0.1:{_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
+        _extra_ws[slot] = ws
+        _extra_ws_tasks[slot] = asyncio.create_task(_slot_recv_loop(slot, ws))
+        print(f"[backup] Slot {slot} connected")
+    except Exception as e:
+        print(f"[backup] Slot {slot} connect failed: {e}")
+
+
+async def _slot_recv_loop(slot, ws):
+    """附加连接的接收循环：与主连接平等处理所有消息。"""
+    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
+
+                if has_method and not has_id:
+                    asyncio.create_task(_handle_event_notification(msg))
+                elif has_method and has_id:
+                    asyncio.create_task(_handle_rpc_request(ws, msg))
+                elif not has_method and has_id:
+                    _handle_rpc_response(msg)
+            except Exception as e:
+                print(f"[backup] Slot {slot} 消息处理异常（已忽略）: {e}")
+    except Exception as e:
+        print(f"[backup] Slot {slot} 接收循环异常: {e}")
+    finally:
+        _extra_ws.pop(slot, None)
+        _extra_ws_tasks.pop(slot, None)
+
+
+async def _handle_connection_release(data):
+    """Kernel 请求释放 slot，优雅关闭。"""
+    for slot in data.get("slots", []):
+        ws = _extra_ws.pop(slot, None)
+        task = _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():
     """Handle module.shutdown event — ack → exiting → cleanup → ready → exit."""
     global _shutting_down
     print("[backup] Received shutdown request")
     _shutting_down = True
     # Step 1: Send ack (立即确认收到)
-    await _publish_event(_ws_global, {
-        "event": "module.shutdown.ack",
-        "data": {"module_id": "backup"},
-    })
+    await _publish_event(_ws_global, "module.shutdown.ack", {"module_id": "backup"})
     # Step 2: Send module.exiting (开始清理)
-    await _publish_event(_ws_global, {
-        "event": "module.exiting",
-        "data": {
-            "module_id": "backup",
-            "type": "passive",
-            "reason": "shutdown_requested",
-            "restart": "auto",
-            "action": "none",
-            "timeout": 2.0,
-            "restart_delay": 0.0,
-        },
+    await _publish_event(_ws_global, "module.exiting", {
+        "module_id": "backup",
+        "type": "passive",
+        "reason": "shutdown_requested",
+        "restart": "auto",
+        "action": "none",
+        "timeout": 2.0,
+        "restart_delay": 0.0,
     })
     # Step 3: Cleanup (nothing to clean up for backup)
+    # 关闭所有附加连接
+    for _s, _w in list(_extra_ws.items()):
+        try:
+            await _w.close(code=1000, reason="shutdown")
+        except Exception:
+            pass
+    for _t in _extra_ws_tasks.values():
+        _t.cancel()
+    _extra_ws.clear()
+    _extra_ws_tasks.clear()
     # Step 4: Send ready (清理完成)
-    await _publish_event(_ws_global, {
-        "event": "module.shutdown.ready",
-        "data": {"module_id": "backup"},
-    })
+    await _publish_event(_ws_global, "module.shutdown.ready", {"module_id": "backup"})
     print("[backup] Shutdown ready, exiting")
-    # Step 5: Exit
-    sys.exit(_exit_code)
+
+    # 等待 Kernel 处理 shutdown.ready（防止 Close 帧抢先）
+    await asyncio.sleep(0.1)
+
+    # Step 5: Close WebSocket connection gracefully
+    if _ws_global:
+        try:
+            await _ws_global.close(code=1000, reason="Graceful shutdown")
+            print("[backup] WebSocket closed")
+        except Exception as e:
+            print(f"[backup] Failed to close WebSocket: {e}")
+
+    # Note: Do NOT call sys.exit() in async context
+    # Let the event loop naturally complete
 
 
 async def _test_event_loop(ws):
     """Publish a test event every 10 seconds."""
-    while True:
-        await asyncio.sleep(10)
-        await _publish_event(ws, {
-            "event": "backup.test",
-            "data": {
+    try:
+        while True:
+            await asyncio.sleep(10)
+            await _publish_event(ws, "backup.test", {
                 "message": "test event from backup",
                 "timestamp": datetime.now(timezone.utc).isoformat(),
-            },
-        })
-        print("[backup] test event published")
+            })
+    except Exception:
+        pass
+
+
+_startup_config = {}  # populated by _validate_and_prepare()
+
+
+def _validate_and_prepare() -> dict:
+    """同步启动验证：日志初始化、读取 token 和 kernel_port。
+    失败直接 sys.exit(1)，成功返回启动参数。在 asyncio.run() 之前调用。"""
+    t0 = time.monotonic()
+    kernel_port = os.environ.get("KITE_KERNEL_PORT")
+    if kernel_port:
+        return _prepare_local_mode(int(kernel_port), t0)
+    config = _load_module_config()
+    gateway_url = config.get("gateway_url") or os.environ.get("KITE_GATEWAY_URL")
+    if gateway_url:
+        return _prepare_remote_mode(gateway_url, t0)
+    print(f"[{MODULE_NAME}] ERROR: No KITE_KERNEL_PORT and no gateway_url")
+    sys.exit(1)
+
+def _prepare_local_mode(kernel_port: int, t0: float) -> dict:
+    global _log_dir, _log_latest_path, _crash_log_path
+    module_data = os.environ.get("KITE_MODULE_DATA")
+    if module_data:
+        _log_dir = os.path.join(module_data, "log")
+        os.makedirs(_log_dir, exist_ok=True)
+        suffix = os.environ.get("KITE_INSTANCE_SUFFIX", "")
+        _log_latest_path = os.path.join(_log_dir, f"latest{suffix}.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, f"crashes{suffix}.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()
+    line = sys.stdin.readline().strip()
+    if not line:
+        print(f"[{MODULE_NAME}] ERROR: stdin closed")
+        sys.exit(1)
+    try: msg = json.loads(line)
+    except json.JSONDecodeError as e:
+        print(f"[{MODULE_NAME}] ERROR: Invalid JSON: {e}")
+        sys.exit(1)
+    if "error" in msg:
+        print(f"[{MODULE_NAME}] 启动失败: {msg.get('message')}")
+        sys.exit(1)
+    token = msg.get("token", "")
+    if not token:
+        print(f"[{MODULE_NAME}] ERROR: No token")
+        sys.exit(1)
+    print(f"[{MODULE_NAME}] Local mode: port={kernel_port}")
+    return {"mode": "local", "token": token, "kernel_port": kernel_port, "t0": t0}
+
+def _prepare_remote_mode(gateway_url: str, t0: float) -> dict:
+    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()
+    kite_token = _get_kite_token(MODULE_NAME, gateway_url)
+    print(f"[{MODULE_NAME}] Remote mode: gateway={gateway_url}")
+    return {"mode": "remote", "gateway_url": gateway_url, "kite_token": kite_token, "t0": t0}
+
+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))
+    if os.path.exists(token_file):
+        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)
+
+
+async def _remote_mode_loop(gateway_url: str, kite_token: str, _t0: float):
+    global _shutting_down, _exit_code
+    retry_count = 0
+    while not _shutting_down and retry_count < 10:
+        try:
+            await _remote_connect(gateway_url, kite_token, _t0)
+            retry_count = 0
+        except asyncio.CancelledError: break
+        except Exception as e:
+            if _shutting_down: break
+            retry_count += 1
+            await asyncio.sleep(min(2 ** retry_count, 60))
+    if retry_count >= 10: _exit_code = 1
+
+async def _remote_connect(gateway_url: str, kite_token: str, _t0: float):
+    global _ws_global
+    async with websockets.connect(gateway_url, open_timeout=10, ping_interval=30) as ws:
+        challenge = json.loads(await ws.recv())
+        auth_req = {"jsonrpc": "2.0", "id": "auth", "method": "auth.connect",
+                    "params": {"nonce": challenge["params"]["nonce"], "module_id": MODULE_NAME,
+                              "auth": {"method": "kite_token", "token": kite_token},
+                              "client": {"type": "module", "platform": sys.platform}}}
+        await ws.send(json.dumps(auth_req))
+        hello_resp = json.loads(await ws.recv())
+        if "error" in hello_resp:
+            error = hello_resp["error"]
+            if error.get("code") in (4001, 4010, 4011):
+                _clear_token_cache(MODULE_NAME, gateway_url)
+                print(f"[{MODULE_NAME}] Token invalid, re-authenticate")
+                sys.exit(1)
+            raise Exception(f"Connection failed: {error}")
+        print(f"[{MODULE_NAME}] Connected via Gateway")
+        _ws_global = ws
+
+        # Start receiver task — wait for system.require_init to trigger _do_init
+        receiver_task = asyncio.create_task(_ws_receiver(ws))
+        try:
+            await receiver_task
+        finally:
+            receiver_task.cancel()
+            for fut in _pending_rpc.values():
+                if not fut.done():
+                    fut.set_result({"error": {"code": -32001, "message": "Connection lost"}})
+            _pending_rpc.clear()
 
 
 if __name__ == "__main__":
+    _startup_config = _validate_and_prepare()
     asyncio.run(main())
+    sys.exit(_exit_code)