@agentunion/kite 1.3.1 → 1.3.2

package/extensions/services/model_service/entry.py

@@ -263,10 +263,19 @@ def _read_stdin_kite_message(expected_type: str, timeout: float = 10) -> dict |
 
 # Global WS reference for publish_event callback
 _ws_global = None
+_shutting_down = False
+
+
+def _is_auth_failure(e: Exception) -> bool:
+    """Check if a WebSocket exception indicates authentication failure."""
+    if hasattr(e, 'rcvd') and e.rcvd is not None:
+        code = e.rcvd.code if hasattr(e.rcvd, 'code') else 0
+        return code in (4001, 4003)
+    return False
 
 
 async def main():
-    global _ws_global
+    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")
@@ -318,41 +327,84 @@ async def main():
 
     print(f"[model_service] Token received ({len(token)} chars), kernel port: {kernel_port} ({_fmt_elapsed(_t0)})")
 
-    # Connect to Kernel WebSocket
+    # Start reconnect loop
+    await _ws_loop(token, kernel_port, _t0)
+
+
+async def _ws_loop(token: str, kernel_port: int, _t0: float):
+    """Connect to Kernel with exponential backoff reconnection."""
+    global _shutting_down
+    retry_delay = 0.3
+    max_delay = 5.0
+    max_retries = 10
+    attempt = 0
+    while not _shutting_down:
+        try:
+            await _ws_connect(token, kernel_port, _t0)
+            retry_delay = 0.3
+            attempt = 0
+        except asyncio.CancelledError:
+            return
+        except Exception as e:
+            attempt += 1
+            if _is_auth_failure(e):
+                print(f"[model_service] Kernel 认证失败，退出")
+                sys.exit(1)
+            if attempt >= max_retries:
+                print(f"[model_service] 重连失败 {max_retries} 次，退出")
+                sys.exit(1)
+            _write_crash(type(e), e, e.__traceback__, severity="error", handled=True)
+            print(f"[model_service] 连接错误: {e}, {retry_delay:.1f}s 后重试 ({attempt}/{max_retries})")
+        _ws_global_clear()
+        if _shutting_down:
+            return
+        await asyncio.sleep(retry_delay)
+        retry_delay = min(retry_delay * 2, max_delay)
+
+
+def _ws_global_clear():
+    global _ws_global
+    _ws_global = None
+
+
+async def _ws_connect(token: str, kernel_port: int, _t0: float):
+    """Single WebSocket session: connect → subscribe → register → ready → receive loop."""
+    global _ws_global
+
     ws_url = f"ws://127.0.0.1:{kernel_port}/ws?token={token}&id=model_service"
     print(f"[model_service] Connecting to Kernel: {ws_url}")
 
-    try:
-        async with websockets.connect(ws_url, open_timeout=5, ping_interval=None, ping_timeout=None, close_timeout=10) as ws:
-            _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)})")
-
-            # 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)})")
+    async with websockets.connect(ws_url, open_timeout=5, ping_interval=None, ping_timeout=None, close_timeout=10) as ws:
+        _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)})")
+
+        # 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
+        # Publish module.ready (every reconnect)
+        if not _shutting_down:
             await _rpc_call(ws, "event.publish", {
                 "event_id": str(uuid.uuid4()),
                 "event": "module.ready",
@@ -363,34 +415,29 @@ async def main():
             })
             print(f"[model_service] module.ready published ({_fmt_elapsed(_t0)})")
 
-            # Start test event loop in background
-            test_task = asyncio.create_task(_test_event_loop(ws))
-
-            # Message loop: handle incoming RPC + events
-            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:
-                        # Event Notification
-                        await _handle_event_notification(msg)
-                    elif has_method and has_id:
-                        # Incoming RPC request
-                        await _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}")
-
-    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)
+        # Start test event loop in background
+        test_task = asyncio.create_task(_test_event_loop(ws))
+
+        # Message loop: handle incoming RPC + events
+        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:
+                    # Event Notification
+                    await _handle_event_notification(msg)
+                elif has_method and has_id:
+                    # Incoming RPC request
+                    await _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}")
 
 
 async def _rpc_call(ws, method: str, params: dict = None):
@@ -416,10 +463,14 @@ async def _handle_event_notification(msg: dict):
     event_type = params.get("event", "")
     data = params.get("data", {})
 
-    # Special handling for module.shutdown targeting model_service
-    if event_type == "module.shutdown" and data.get("module_id") == "model_service":
-        await _handle_shutdown()
-        return
+    # Special handling for module.shutdown
+    if event_type == "module.shutdown":
+        target = data.get("module_id", "")
+        reason = data.get("reason", "")
+        # Handle both targeted shutdown (module_id == "model_service") and broadcast shutdown (no module_id or launcher_lost)
+        if target == "model_service" or not target or reason == "launcher_lost":
+            await _handle_shutdown()
+            return
 
     # Log other events
     print(f"[model_service] Event received: {event_type}")
@@ -472,8 +523,15 @@ async def _rpc_status() -> dict:
 
 
 async def _handle_shutdown():
-    """Handle module.shutdown event — ack, cleanup, ready, exit."""
+    """Handle module.shutdown event — exiting → ack → cleanup → ready → exit."""
+    global _shutting_down
     print("[model_service] Received shutdown request")
+    _shutting_down = True
+    # Step 0: Send module.exiting
+    await _publish_event(_ws_global, {
+        "event": "module.exiting",
+        "data": {"module_id": "model_service", "action": "none"},
+    })
     # Step 1: Send ack
     await _publish_event(_ws_global, {
         "event": "module.shutdown.ack",

package/extensions/services/watchdog/entry.py

@@ -265,11 +265,25 @@ def _read_stdin_kite_message(expected_type: str, timeout: float = 10) -> dict |
 
 # Global WS reference for publish_event callback
 _ws_global = None
+_shutting_down = False
+_monitor = None
+_monitor_task = None
 
+# RPC request-response infrastructure
+_rpc_waiters: dict[str, asyncio.Event] = {}   # rpc_id -> Event
+_rpc_results: dict[str, dict] = {}             # rpc_id -> response dict
+
+
+def _is_auth_failure(e: Exception) -> bool:
+    """Check if a WebSocket exception indicates authentication failure."""
+    if hasattr(e, 'rcvd') and e.rcvd is not None:
+        code = e.rcvd.code if hasattr(e.rcvd, 'code') else 0
+        return code in (4001, 4003)
+    return False
 
 
 async def main():
-    global _ws_global
+    global _ws_global, _shutting_down, _monitor
     # Initialize log file paths
     global _log_dir, _log_latest_path, _crash_log_path
     module_data = os.environ.get("KITE_MODULE_DATA")
@@ -321,57 +335,104 @@ async def main():
 
     print(f"[watchdog] Token received ({len(token)} chars), kernel port: {kernel_port} ({_fmt_elapsed(_t0)})")
 
-    # Connect to Kernel WebSocket
-    ws_url = f"ws://127.0.0.1:{kernel_port}/ws?token={token}&id=watchdog"
-    print(f"[watchdog] Connecting to Kernel: {ws_url}")
+    # Create monitor (once, persists across reconnects)
+    _monitor = HealthMonitor(
+        own_token=token,
+        kernel_port=kernel_port,
+    )
 
-    try:
-        async with websockets.connect(ws_url, open_timeout=5, ping_interval=None, ping_timeout=None, close_timeout=10) as ws:
-            _ws_global = ws
-            print(f"[watchdog] Connected to Kernel ({_fmt_elapsed(_t0)})")
-
-            # Subscribe to events
-            await _rpc_call(ws, "event.subscribe", {
-                "events": [
-                    "system.ready",
-                    "module.started",
-                    "module.stopped",
-                    "module.exiting",
-                    "module.ready",
-                    "module.shutdown",
-                ],
-            })
-            print(f"[watchdog] Subscribed to events ({_fmt_elapsed(_t0)})")
-
-            # Register to Kernel Registry via RPC
-            await _rpc_call(ws, "registry.register", {
-                "module_id": "watchdog",
-                "module_type": "service",
-                "events_publish": {
-                    "watchdog.module.unhealthy": {},
-                    "watchdog.module.recovered": {},
-                    "watchdog.alert": {},
-                },
-                "events_subscribe": [
-                    "system.ready",
-                    "module.started",
-                    "module.stopped",
-                    "module.exiting",
-                    "module.ready",
-                    "module.shutdown",
-                ],
-            })
-            print(f"[watchdog] Registered to Kernel ({_fmt_elapsed(_t0)})")
+    # Start reconnect loop
+    await _ws_loop(token, kernel_port, _t0)
 
-            # Create monitor with RPC callback
-            monitor = HealthMonitor(
-                own_token=token,
-                kernel_port=kernel_port,
-            )
-            monitor.publish_event = lambda event: asyncio.create_task(_publish_event(ws, event))
-            monitor.rpc_call = lambda method, params: _rpc_call(ws, method, params)
 
-            # Publish module.ready
+async def _ws_loop(token: str, kernel_port: int, _t0: float):
+    """Connect to Kernel with exponential backoff reconnection."""
+    global _shutting_down
+    retry_delay = 0.3
+    max_delay = 5.0
+    max_retries = 10
+    attempt = 0
+    while not _shutting_down:
+        try:
+            await _ws_connect(token, kernel_port, _t0)
+            retry_delay = 0.3
+            attempt = 0
+        except asyncio.CancelledError:
+            return
+        except Exception as e:
+            attempt += 1
+            if _is_auth_failure(e):
+                print(f"[watchdog] Kernel 认证失败，退出")
+                sys.exit(1)
+            if attempt >= max_retries:
+                print(f"[watchdog] 重连失败 {max_retries} 次，退出")
+                sys.exit(1)
+            _write_crash(type(e), e, e.__traceback__, severity="error", handled=True)
+            print(f"[watchdog] 连接错误: {e}, {retry_delay:.1f}s 后重试 ({attempt}/{max_retries})")
+        _ws_global_clear()
+        if _shutting_down:
+            return
+        await asyncio.sleep(retry_delay)
+        retry_delay = min(retry_delay * 2, max_delay)
+
+
+def _ws_global_clear():
+    global _ws_global
+    _ws_global = None
+
+
+async def _ws_connect(token: str, kernel_port: int, _t0: float):
+    """Single WebSocket session: connect → subscribe → register → ready → receive loop."""
+    global _ws_global, _monitor, _monitor_task
+
+    ws_url = f"ws://127.0.0.1:{kernel_port}/ws?token={token}&id=watchdog"
+    print(f"[watchdog] Connecting to Kernel: {ws_url}")
+
+    async with websockets.connect(ws_url, open_timeout=5, ping_interval=None, ping_timeout=None, close_timeout=10) as ws:
+        _ws_global = ws
+        print(f"[watchdog] Connected to Kernel ({_fmt_elapsed(_t0)})")
+
+        # Subscribe to events
+        await _rpc_call(ws, "event.subscribe", {
+            "events": [
+                "system.ready",
+                "module.started",
+                "module.stopped",
+                "module.exiting",
+                "module.ready",
+                "module.shutdown",
+                "module.offline",
+            ],
+        })
+        print(f"[watchdog] Subscribed to events ({_fmt_elapsed(_t0)})")
+
+        # Register to Kernel Registry via RPC
+        await _rpc_call(ws, "registry.register", {
+            "module_id": "watchdog",
+            "module_type": "service",
+            "events_publish": {
+                "watchdog.module.unhealthy": {},
+                "watchdog.module.recovered": {},
+                "watchdog.alert": {},
+            },
+            "events_subscribe": [
+                "system.ready",
+                "module.started",
+                "module.stopped",
+                "module.exiting",
+                "module.ready",
+                "module.shutdown",
+                "module.offline",
+            ],
+        })
+        print(f"[watchdog] Registered to Kernel ({_fmt_elapsed(_t0)})")
+
+        # Set up monitor callbacks (reconnect-safe)
+        _monitor.publish_event = lambda event: asyncio.create_task(_publish_event(ws, event))
+        _monitor.rpc_call = lambda method, params: _rpc_call_with_response(ws, method, params)
+
+        # Publish module.ready (every reconnect)
+        if not _shutting_down:
             await _rpc_call(ws, "event.publish", {
                 "event_id": str(uuid.uuid4()),
                 "event": "module.ready",
@@ -382,34 +443,35 @@ async def main():
             })
             print(f"[watchdog] module.ready published ({_fmt_elapsed(_t0)})")
 
-            # Start monitor loop in background
-            monitor_task = asyncio.create_task(monitor.run())
-
-            # Message loop: handle incoming RPC + events
-            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:
-                        # Event Notification
-                        await _handle_event_notification(msg, monitor)
-                    elif has_method and has_id:
-                        # Incoming RPC request
-                        await _handle_rpc_request(ws, msg, monitor)
-                    # Ignore RPC responses (we don't await them in this simple impl)
-                except Exception as e:
-                    print(f"[watchdog] 消息处理异常（已忽略）: {e}")
-
-    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)
+        # Start monitor loop if not already running
+        if _monitor_task is None or _monitor_task.done():
+            _monitor_task = asyncio.create_task(_monitor.run())
+
+        # Message loop: handle incoming RPC + events
+        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:
+                    # Event Notification
+                    await _handle_event_notification(msg, _monitor)
+                elif has_method and has_id:
+                    # Incoming RPC request
+                    await _handle_rpc_request(ws, msg, _monitor)
+                elif has_id and not has_method:
+                    # RPC response — route to waiter
+                    msg_id = msg["id"]
+                    if msg_id in _rpc_waiters:
+                        _rpc_results[msg_id] = msg
+                        _rpc_waiters[msg_id].set()
+            except Exception as e:
+                print(f"[watchdog] 消息处理异常（已忽略）: {e}")
 
 
 
@@ -421,6 +483,28 @@ async def _rpc_call(ws, method: str, params: dict = None):
     await ws.send(json.dumps(msg))
 
 
+async def _rpc_call_with_response(ws, method: str, params: dict = None, timeout: float = 5) -> dict:
+    """Send a JSON-RPC 2.0 request and await the response."""
+    rpc_id = str(uuid.uuid4())
+    msg = {"jsonrpc": "2.0", "id": rpc_id, "method": method}
+    if params:
+        msg["params"] = params
+
+    evt = asyncio.Event()
+    _rpc_waiters[rpc_id] = evt
+
+    await ws.send(json.dumps(msg))
+
+    try:
+        await asyncio.wait_for(evt.wait(), timeout=timeout)
+        return _rpc_results.pop(rpc_id, {})
+    except asyncio.TimeoutError:
+        return {"error": {"code": -32000, "message": f"RPC timeout: {method}"}}
+    finally:
+        _rpc_waiters.pop(rpc_id, None)
+        _rpc_results.pop(rpc_id, None)
+
+
 async def _publish_event(ws, event: dict):
     """Publish an event via RPC event.publish."""
     await _rpc_call(ws, "event.publish", {
@@ -441,8 +525,8 @@ async def _handle_event_notification(msg: dict, monitor: HealthMonitor):
         await _handle_shutdown(monitor)
         return
 
-    # Forward to monitor
-    await monitor.handle_event(msg)
+    # Forward to monitor (extract params from JSON-RPC notification)
+    await monitor.handle_event(params)
 
 
 async def _handle_rpc_request(ws, msg: dict, monitor: HealthMonitor):
@@ -489,8 +573,15 @@ async def _rpc_status(monitor: HealthMonitor) -> dict:
 
 
 async def _handle_shutdown(monitor: HealthMonitor):
-    """Handle module.shutdown event — ack, cleanup, ready, exit."""
+    """Handle module.shutdown event — exiting → ack → cleanup → ready → exit."""
+    global _shutting_down
     print("[watchdog] Received shutdown request")
+    _shutting_down = True
+    # Step 0: Send module.exiting
+    await _publish_event(_ws_global, {
+        "event": "module.exiting",
+        "data": {"module_id": "watchdog", "action": "none"},
+    })
     # Step 1: Send ack
     await _publish_event(_ws_global, {
         "event": "module.shutdown.ack",