@agentunion/kite 1.4.0 → 1.6.0

package/extensions/services/model_service/entry.py

@@ -15,6 +15,7 @@ from datetime import datetime, timezone
 import asyncio
 import traceback
 import uuid
+import random
 
 import websockets
 
@@ -120,6 +121,19 @@ 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
+
+        # business_config（从 businesses 数组取）
+        businesses = fm.get("businesses")
+        if isinstance(businesses, list) and businesses:
+            cfg_file = businesses[0].get("config_file") if isinstance(businesses[0], dict) else None
+            if cfg_file:
+                result["business_config"] = str(cfg_file).strip()
+
     except SystemExit:
         raise  # Re-raise exit to prevent catching by outer except
     except Exception as e:
@@ -131,6 +145,20 @@ def _load_module_config() -> dict:
 _module_config = _load_module_config()
 MODULE_NAME = _module_config["name"]
 
+# 读取业务配置
+_business_config = {}
+config_file = _module_config.get("business_config")
+if config_file:
+    config_path = os.path.join(os.path.dirname(__file__), config_file)
+    if os.path.exists(config_path):
+        try:
+            import json5
+            with open(config_path, 'r', encoding='utf-8') as f:
+                _business_config = json5.load(f)
+            print(f"[{MODULE_NAME}] 业务配置已加载: {config_file}")
+        except Exception as e:
+            print(f"[{MODULE_NAME}] 业务配置读取失败: {e}")
+
 
 class _SafeWriter:
     """Wraps a stream to silently swallow BrokenPipeError on write/flush."""
@@ -372,6 +400,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."""
@@ -381,197 +417,317 @@ def _is_auth_failure(e: Exception) -> bool:
     return False
 
 
-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", "")
+# ── 远程模式：Token 管理 ──
 
-        _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
+def _gateway_to_filename(gateway_url: str) -> str:
+    """从 gateway_url 生成缓存文件名"""
+    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"
 
-        _resolve_daily_log_path()
 
-    _setup_exception_hooks()
+def _load_token_cache(module_name: str, gateway_url: str) -> dict | None:
+    """加载 token 缓存"""
+    home = os.environ.get("HOME") or os.environ.get("USERPROFILE") or os.path.expanduser("~")
+    data_dir = os.path.join(home, ".kite", "remote", module_name, "tokens")
+    filename = _gateway_to_filename(gateway_url)
+    token_file = os.path.join(data_dir, filename)
 
-    _t0 = time.monotonic()
+    if not os.path.exists(token_file):
+        return None
 
-    # 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", "")
+        with open(token_file, "r") as f:
+            return json.load(f)
     except Exception:
-        pass
+        return None
 
-    # 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("[model_service] ERROR: Missing token or kernel_port")
-        sys.exit(1)
+def _clear_token_cache(module_name: str, gateway_url: str):
+    """清除 token 缓存"""
+    home = os.environ.get("HOME") or os.environ.get("USERPROFILE") or os.path.expanduser("~")
+    data_dir = os.path.join(home, ".kite", "remote", module_name, "tokens")
+    filename = _gateway_to_filename(gateway_url)
+    token_file = os.path.join(data_dir, filename)
+
+    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:
+    """获取 kite_token（环境变量 > 缓存 > 提示配对）"""
+    # 1. 环境变量优先
+    token = os.environ.get("KITE_TOKEN")
+    if token:
+        return token
+
+    # 2. 从缓存加载
+    cache = _load_token_cache(module_name, gateway_url)
+    if cache:
+        return cache["token"]
 
-    print(f"[model_service] Token received ({len(token)} chars), kernel port: {kernel_port} ({_fmt_elapsed(_t0)})")
+    # 3. 提示首次配对
+    print(f"[{module_name}] No token found for {gateway_url}")
+    print(f"  Please authenticate:")
+    print(f"    export KITE_TOKEN=<your_token>")
+    print(f"    # Or use kite-cli:")
+    print(f"    kite-cli auth --module {module_name} --gateway {gateway_url}")
+    sys.exit(1)
 
-    # Start reconnect loop
-    await _ws_loop(token, kernel_port, _t0)
+
+async def main():
+    global _ws_global, _shutting_down
+    # 根据启动模式选择连接方式
+    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"[model_service] Kernel 认证失败，退出")
-                _exit_code = 1
-                _shutting_down = True
+            if _shutting_down:
                 return
-            if attempt >= max_retries:
-                print(f"[model_service] 重连失败 {max_retries} 次，退出")
+
+            code = _get_close_code(e)
+
+            # never: 永不重连
+            if code in (4001, 4003, 4004, 1008, 4010):
+                print(f"[model_service] 致命错误 (code {code})，退出")
                 _exit_code = 1
                 _shutting_down = True
                 return
+
+            # cooldown: 速率限制
+            if code == 4020:
+                cooldown_attempts += 1
+                if cooldown_attempts >= 5:
+                    print(f"[model_service] 速率限制重试 5 次，退出")
+                    _exit_code = 1
+                    _shutting_down = True
+                    return
+                print(f"[model_service] 速率限制，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"[model_service] 连接错误: {e}, {retry_delay:.1f}s 后重试 ({attempt}/{max_retries})")
+            print(f"[model_service] 连接错误: {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"[model_service] 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"[model_service] Subscribed to events")
+
+    # 2. Register to Kernel Registry via RPC
+    await _rpc_call(ws, "registry.register", {
+        "module_id": "model_service",
+        "module_type": "service",
+        "events_publish": {
+            "model_service": {
+                "test": {"description": "Test event from model_service module"},
+            }
+        },
+        "events_subscribe": [
+            "module.started",
+            "module.stopped",
+            "module.shutdown",
+        ],
+    })
+    print(f"[model_service] 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": "model_service",
+            "graceful_shutdown": True,
+            "startup_time": startup_time,
+            "reason": reason,
+        })
+        print(f"[model_service] 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
 
-    ws_url = f"ws://127.0.0.1:{kernel_port}/ws?token={token}&id=model_service"
+    _kernel_port = kernel_port
+    ws_url = f"ws://127.0.0.1:{kernel_port}/ws?id=model_service"
     print(f"[model_service] Connecting to Kernel: {ws_url}")
 
-    async with websockets.connect(ws_url, open_timeout=5, ping_interval=20, ping_timeout=20, close_timeout=10) as ws:
+    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"[model_service] Connected to Kernel ({_fmt_elapsed(_t0)})")
 
-        # Subscribe to events
-        await _rpc_call(ws, "event.subscribe", {
-            "events": [
-                "module.started",
-                "module.stopped",
-                "module.shutdown",
-            ],
-        })
-        print(f"[model_service] 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": "model_service",
-            "module_type": "service",
-            "events_publish": {
-                "model_service": {
-                    "test": {"description": "Test event from model_service module"},
-                }
-            },
-            "events_subscribe": [
-                "module.started",
-                "module.stopped",
-                "module.shutdown",
-            ],
-        })
-        print(f"[model_service] Registered to Kernel ({_fmt_elapsed(_t0)})")
-
-        # Publish module.ready (every reconnect)
-        if not _shutting_down:
-            await _rpc_call(ws, "event.publish", {
-                "event_id": str(uuid.uuid4()),
-                "event": "module.ready",
-                "data": {
-                    "module_id": "model_service",
-                    "graceful_shutdown": True,
-                },
-            })
-            print(f"[model_service] 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"[model_service] 消息处理异常（已忽略）: {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"[model_service] 消息处理异常（已忽略）: {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):
+    """Handle system.ping event and reply with system.pong."""
+    t1 = data.get("ping_time")
+    t2 = time.time()
+
+    await _publish_event(_ws_global, "system.pong", {
+        "module_id": MODULE_NAME,
+        "ping_time": t1,
+        "pong_time": t2,
     })
 
 
@@ -581,6 +737,19 @@ async def _handle_event_notification(msg: dict):
     event_type = params.get("event", "")
     data = params.get("data", {})
 
+    # Handle system.ping event
+    if event_type == "system.ping":
+        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", "")
@@ -640,53 +809,355 @@ async def _rpc_status() -> 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"[model_service] Slot {slot} connected")
+    except Exception as e:
+        print(f"[model_service] 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"[model_service] Slot {slot} 消息处理异常（已忽略）: {e}")
+    except Exception as e:
+        print(f"[model_service] 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("[model_service] Received shutdown request")
     _shutting_down = True
     # Step 1: Send ack (立即确认收到)
-    await _publish_event(_ws_global, {
-        "event": "module.shutdown.ack",
-        "data": {"module_id": "model_service"},
-    })
+    await _publish_event(_ws_global, "module.shutdown.ack", {"module_id": "model_service"})
     # Step 2: Send module.exiting (开始清理)
-    await _publish_event(_ws_global, {
-        "event": "module.exiting",
-        "data": {
-            "module_id": "model_service",
-            "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": "model_service",
+        "type": "passive",
+        "reason": "shutdown_requested",
+        "restart": "auto",
+        "action": "none",
+        "timeout": 2.0,
+        "restart_delay": 0.0,
     })
     # Step 3: Cleanup (nothing to clean up for model_service)
+    # 关闭所有附加连接
+    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": "model_service"},
-    })
+    await _publish_event(_ws_global, "module.shutdown.ready", {"module_id": "model_service"})
     print("[model_service] 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("[model_service] WebSocket closed")
+        except Exception as e:
+            print(f"[model_service] 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": "model_service.test",
-            "data": {
+    try:
+        counter = 0
+        while True:
+            await asyncio.sleep(10)
+            counter += 1
+
+            await _publish_event(ws, "model_service.test", {
                 "message": "test event from model_service",
                 "timestamp": datetime.now(timezone.utc).isoformat(),
-            },
-        })
-        print("[model_service] test event published")
+            })
+
+            # Redis 测试：每 3 次保存模型调用统计
+            if counter % 3 == 0:
+                try:
+                    await _rpc_call(ws, "redis.hset", {
+                        "key": "model_service:stats",
+                        "field": "last_test",
+                        "value": str(int(time.time()))
+                    })
+                    result = await _rpc_call(ws, "redis.hget", {
+                        "key": "model_service:stats",
+                        "field": "last_test"
+                    })
+                    if result:
+                        value = result.get("result", {}).get("value")
+                        print(f"[model_service] ✓ Redis 统计已更新: {value}")
+                except Exception as e:
+                    print(f"[model_service] Redis 测试失败: {e}")
+    except Exception:
+        pass
+
+
+_startup_config = {}  # populated by _validate_and_prepare()
+
+
+def _validate_and_prepare() -> dict:
+    """同步启动验证：日志初始化、读取 token 和 kernel_port。
+    失败直接 sys.exit(1)，成功返回启动参数。在 asyncio.run() 之前调用。"""
+    global _log_dir, _log_latest_path, _crash_log_path
+
+    t0 = time.monotonic()
+
+    # 优先级1：检查KITE_KERNEL_PORT（本地模式）
+    kernel_port = os.environ.get("KITE_KERNEL_PORT")
+
+    if kernel_port:
+        # 本地模式：从stdin读取boot_info
+        return _prepare_local_mode(int(kernel_port), t0)
+
+    # 优先级2：检查gateway_url（远程模式）
+    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 in module.md")
+    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()
+
+    # Read boot_info from stdin (only token)
+    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 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", "")
+    if not token:
+        print(f"[{MODULE_NAME}] ERROR: No token in stdin message")
+        sys.exit(1)
+
+    print(f"[{MODULE_NAME}] Token received ({len(token)} chars), kernel port: {kernel_port} ({_fmt_elapsed(t0)})")
+
+    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
+    kite_token = _get_kite_token(MODULE_NAME, gateway_url)
+
+    print(f"[{MODULE_NAME}] Remote mode: gateway={gateway_url} ({_fmt_elapsed(t0)})")
+
+    return {"mode": "remote", "gateway_url": gateway_url, "kite_token": kite_token, "t0": t0}
+
+
+# ── 远程模式：连接循环 ──
+
+async def _remote_mode_loop(gateway_url: str, kite_token: str, _t0: float):
+    """通过 Gateway 连接到 Kernel（带重连）"""
+    global _shutting_down, _exit_code
+    retry_count = 0
+    max_retries = 10
+
+    while not _shutting_down and retry_count < max_retries:
+        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
+            backoff = min(2 ** retry_count, 60)
+            print(f"[{MODULE_NAME}] Connection error: {e}, retrying in {backoff}s...")
+            await asyncio.sleep(backoff)
+
+    if retry_count >= max_retries:
+        print(f"[{MODULE_NAME}] Max retries reached, exiting")
+        _exit_code = 1
+
+
+async def _remote_connect(gateway_url: str, kite_token: str, _t0: float):
+    """单次 Gateway 连接"""
+    global _ws_global
+
+    async with websockets.connect(gateway_url, open_timeout=10, ping_interval=30) as ws:
+        # 1. 接收 challenge
+        challenge = json.loads(await ws.recv())
+        nonce = challenge["params"]["nonce"]
+
+        # 2. 发送 auth.connect
+        auth_req = {
+            "jsonrpc": "2.0",
+            "id": "auth",
+            "method": "auth.connect",
+            "params": {
+                "nonce": 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))
+
+        # 3. 等待 hello-ok
+        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, please re-authenticate")
+                sys.exit(1)
+            raise Exception(f"Connection failed: {error}")
+
+        print(f"[{MODULE_NAME}] Connected via Gateway ({_fmt_elapsed(_t0)})")
+        _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)