@agentunion/kite 1.2.0 → 1.3.1

package/extensions/services/model_service/entry.py

@@ -6,15 +6,17 @@ Reads boot_info from stdin, registers to Registry, starts model_service service.
 import builtins
 import json
 import os
-import socket
 import sys
 import threading
 import re
 import time
 from datetime import datetime, timezone
 
-import httpx
-import uvicorn
+import asyncio
+import traceback
+import uuid
+
+import websockets
 
 
 # ── Safe stdout/stderr: ignore BrokenPipeError after Launcher closes stdio ──
@@ -228,7 +230,6 @@ _project_root = os.environ.get("KITE_PROJECT") or os.path.dirname(os.path.dirnam
 if _project_root not in sys.path:
     sys.path.insert(0, _project_root)
 
-from extensions.services.model_service.server import ModelServiceServer
 
 
 def _fmt_elapsed(t0: float) -> str:
@@ -240,62 +241,32 @@ def _fmt_elapsed(t0: float) -> str:
     return f"{d:.0f}s"
 
 
-def _get_free_port() -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        return s.getsockname()[1]
-
-
-def _register_to_registry(client: httpx.Client, token: str, registry_url: str, port: int):
-    payload = {
-        "action": "register",
-        "module_id": "model_service",
-        "module_type": "service",
-        "name": "Model Service",
-        "api_endpoint": f"http://127.0.0.1:{port}",
-        "health_endpoint": "/health",
-        "events_publish": {
-            "model_service.test": {"description": "Test event from model_service module"},
-        },
-        "events_subscribe": [
-            "module.started",
-            "module.stopped",
-            "module.shutdown",
-        ],
-    }
-    headers = {"Authorization": f"Bearer {token}"}
-    try:
-        resp = client.post(f"{registry_url}/modules", json=payload, headers=headers)
-        if resp.status_code == 200:
-            pass  # timing printed in main()
-        else:
-            print(f"[model_service] WARNING: Registry returned {resp.status_code}")
-    except Exception as e:
-        print(f"[model_service] WARNING: Registry registration failed: {e}")
-
+def _read_stdin_kite_message(expected_type: str, timeout: float = 10) -> dict | None:
+    """Read a single kite message of expected type from stdin with timeout."""
+    result = [None]
 
-def _get_event_hub_ws(client: httpx.Client, token: str, registry_url: str) -> str:
-    """Discover Event Hub WebSocket endpoint from Registry, with retry."""
-    import time
-    headers = {"Authorization": f"Bearer {token}"}
-    deadline = time.time() + 10
-    while time.time() < deadline:
+    def _read():
         try:
-            resp = client.get(
-                f"{registry_url}/get/event_hub.metadata.ws_endpoint",
-                headers=headers,
-            )
-            if resp.status_code == 200:
-                val = resp.json()
-                if val:
-                    return val
+            line = sys.stdin.readline().strip()
+            if line:
+                msg = json.loads(line)
+                if isinstance(msg, dict) and msg.get("kite") == expected_type:
+                    result[0] = msg
         except Exception:
             pass
-        time.sleep(0.2)
-    return ""
+
+    t = threading.Thread(target=_read, daemon=True)
+    t.start()
+    t.join(timeout=timeout)
+    return result[0]
 
 
-def main():
+# Global WS reference for publish_event callback
+_ws_global = None
+
+
+async def main():
+    global _ws_global
     # Initialize log file paths
     global _log_dir, _log_latest_path, _crash_log_path
     module_data = os.environ.get("KITE_MODULE_DATA")
@@ -334,47 +305,204 @@ def main():
     except Exception:
         pass
 
-    # Read registry_port from environment variable
-    registry_port = int(os.environ.get("KITE_REGISTRY_PORT", "0"))
+    # 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 registry_port:
-        print("[model_service] ERROR: Missing token or KITE_REGISTRY_PORT")
+    if not token or not kernel_port:
+        print("[model_service] ERROR: Missing token or kernel_port")
         sys.exit(1)
 
-    print(f"[model_service] Token received ({len(token)} chars), registry port: {registry_port} ({_fmt_elapsed(_t0)})")
-
-    registry_url = f"http://127.0.0.1:{registry_port}"
-    port = _get_free_port()
-
-    # Register and discover Event Hub synchronously before starting uvicorn
-    client = httpx.Client(timeout=5)
-    _register_to_registry(client, token, registry_url, port)
-    print(f"[model_service] Registered to Registry ({_fmt_elapsed(_t0)})")
-    event_hub_ws = _get_event_hub_ws(client, token, registry_url)
-    if not event_hub_ws:
-        print("[model_service] WARNING: Could not discover Event Hub WS, events disabled")
-    else:
-        print(f"[model_service] Discovered Event Hub: {event_hub_ws}")
-    client.close()
+    print(f"[model_service] Token received ({len(token)} chars), kernel port: {kernel_port} ({_fmt_elapsed(_t0)})")
 
-    server = ModelServiceServer(
-        token=token,
-        registry_url=registry_url,
-        event_hub_ws=event_hub_ws,
-        boot_t0=_t0,
-    )
+    # Connect to Kernel WebSocket
+    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}")
 
-    print(f"[model_service] Starting on port {port} ({_fmt_elapsed(_t0)})")
     try:
-        config = uvicorn.Config(server.app, host="127.0.0.1", port=port, log_level="warning")
-        uvi_server = uvicorn.Server(config)
-        server._uvicorn_server = uvi_server
-        uvi_server.run()
+        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
+            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)})")
+
+            # 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)
 
 
+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))
+
+
+async def _publish_event(ws, event: dict):
+    """Publish an event via RPC event.publish."""
+    await _rpc_call(ws, "event.publish", {
+        "event_id": str(uuid.uuid4()),
+        "event": event.get("event", ""),
+        "data": event.get("data", {}),
+    })
+
+
+async def _handle_event_notification(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", {})
+
+    # Special handling for module.shutdown targeting model_service
+    if event_type == "module.shutdown" and data.get("module_id") == "model_service":
+        await _handle_shutdown()
+        return
+
+    # Log other events
+    print(f"[model_service] Event received: {event_type}")
+
+
+async def _handle_rpc_request(ws, msg: dict):
+    """Handle an incoming RPC request (model_service.* methods)."""
+    rpc_id = msg.get("id", "")
+    method = msg.get("method", "")
+    params = msg.get("params", {})
+
+    handlers = {
+        "health": lambda p: _rpc_health(),
+        "status": lambda p: _rpc_status(),
+    }
+    handler = handlers.get(method)
+    if handler:
+        try:
+            result = await handler(params)
+            await ws.send(json.dumps({"jsonrpc": "2.0", "id": rpc_id, "result": result}))
+        except Exception as e:
+            await ws.send(json.dumps({
+                "jsonrpc": "2.0", "id": rpc_id,
+                "error": {"code": -32603, "message": str(e)},
+            }))
+    else:
+        await ws.send(json.dumps({
+            "jsonrpc": "2.0", "id": rpc_id,
+            "error": {"code": -32601, "message": f"Method not found: {method}"},
+        }))
+
+
+async def _rpc_health() -> dict:
+    """RPC handler for model_service.health."""
+    return {
+        "status": "healthy",
+        "details": {
+            "uptime_seconds": round(time.time() - _start_ts),
+        },
+    }
+
+
+async def _rpc_status() -> dict:
+    """RPC handler for model_service.status."""
+    return {
+        "module": "model_service",
+        "status": "running",
+        "uptime_seconds": round(time.time() - _start_ts),
+    }
+
+
+async def _handle_shutdown():
+    """Handle module.shutdown event — ack, cleanup, ready, exit."""
+    print("[model_service] Received shutdown request")
+    # Step 1: Send ack
+    await _publish_event(_ws_global, {
+        "event": "module.shutdown.ack",
+        "data": {"module_id": "model_service", "estimated_cleanup": 2},
+    })
+    # Step 2: Cleanup (nothing to clean up for model_service)
+    # Step 3: Send ready
+    await _publish_event(_ws_global, {
+        "event": "module.shutdown.ready",
+        "data": {"module_id": "model_service"},
+    })
+    print("[model_service] Shutdown ready, exiting")
+    # Step 4: Exit
+    sys.exit(0)
+
+
+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": {
+                "message": "test event from model_service",
+                "timestamp": datetime.now(timezone.utc).isoformat(),
+            },
+        })
+        print("[model_service] test event published")
+
+
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())