@agentunion/kite 1.0.6 → 1.2.0

package/core/event_hub/server.py

@@ -2,9 +2,17 @@
 Event Hub HTTP + WebSocket server.
 FastAPI app: /ws (WebSocket), /health, /stats.
 30s timer for heartbeat renewal + dedup cleanup.
+
+Launcher bootstrap sequence:
+  Launcher connects with launcher_ws_token → Event Hub verifies locally →
+  registers to Registry → sends module.ready.
 """
 
 import asyncio
+import json
+import os
+import uuid
+from datetime import datetime, timezone
 
 import httpx
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
@@ -14,35 +22,159 @@ from .hub import EventHub
 
 class EventHubServer:
 
-    def __init__(self, hub: EventHub, own_token: str, registry_url: str, test_mode: bool = False):
+    def __init__(self, hub: EventHub, own_token: str, registry_url: str,
+                 launcher_ws_token: str = "",
+                 advertise_ip: str = "127.0.0.1", port: int = 0):
         self.hub = hub
         self.own_token = own_token
         self.registry_url = registry_url
-        self.test_mode = test_mode
+        self.is_debug = os.environ.get("KITE_DEBUG") == "1"
+        self.launcher_ws_token = launcher_ws_token
+        self.advertise_ip = advertise_ip
+        self.port = port
         self._timer_task: asyncio.Task | None = None
+        self._launcher_connected = False
+        self._registered_to_registry = False
+        self._uvicorn_server = None  # set by entry.py for graceful shutdown
+        self._launcher_ws: WebSocket | None = None  # reference to Launcher's WS
+        self._http_client: httpx.AsyncClient | None = None  # reused for Registry calls
         self.app = self._create_app()
+        # Register shutdown callback on hub
+        self.hub._shutdown_cb = self._handle_shutdown
 
-    # ── Token verification via Registry ──
+    # ── Token verification ──
 
     async def _verify_token(self, token: str, module_id_hint: str = "") -> str | None:
-        """Call Registry POST /verify to validate a module token. Returns module_id or None."""
-        if self.test_mode:
-            return module_id_hint or token[:32] or "test_module"
+        """Verify a token. Launcher's launcher_ws_token is verified locally.
+        Other tokens are verified via Registry POST /verify.
+        In debug mode (KITE_DEBUG=1), any non-empty token is accepted."""
+        if self.is_debug and token:
+            return module_id_hint or "debug"
+        # Local verification for Launcher bootstrap (before Registry registration)
+        if self.launcher_ws_token and token == self.launcher_ws_token:
+            return "launcher"
+        # Normal verification via Registry
+        import time as _time
+        _t0 = _time.monotonic()
+        try:
+            resp = await self._http_client.post(
+                f"{self.registry_url}/verify",
+                json={"token": token},
+                headers={"Authorization": f"Bearer {self.own_token}"},
+                timeout=5,
+            )
+            _t1 = _time.monotonic()
+            if resp.status_code == 200:
+                body = resp.json()
+                if body.get("ok"):
+                    mid = body.get("module_id")
+                    print(f"[event_hub] _verify_token({mid}): {(_t1-_t0)*1000:.0f}ms")
+                    return mid
+        except Exception as e:
+            print(f"[event_hub] Token verification failed: {e}")
+        return None
+
+    # ── Launcher bootstrap ──
+
+    async def _on_launcher_connected(self, ws: WebSocket):
+        """Called on first Launcher WS connect. Registers to Registry first, then sends module.ready.
+        This order ensures all startup-phase prints are done before Launcher calls close_stdio."""
+        self._launcher_connected = True
+        self._launcher_ws = ws
+        # Step 1: Register to Registry (may print log lines)
+        await self._register_to_registry()
+        # Step 2: Send module.ready (Launcher will close_stdio after receiving this)
+        msg = {
+            "type": "event",
+            "event_id": str(uuid.uuid4()),
+            "event": "module.ready",
+            "source": "event_hub",
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "data": {
+                "module_id": "event_hub",
+                "ws_endpoint": f"ws://{self.advertise_ip}:{self.port}/ws",
+                "graceful_shutdown": True,
+            },
+        }
+        try:
+            await ws.send_text(json.dumps(msg))
+            print("[event_hub] Sent module.ready to Launcher")
+        except Exception as e:
+            print(f"[event_hub] Failed to send module.ready: {e}")
+
+    async def _handle_shutdown(self, data: dict):
+        """Handle module.shutdown event targeting event_hub."""
+        print("[event_hub] Received shutdown request")
+        # Send ack and ready via Launcher's WS connection
+        ws = self._launcher_ws
+        if ws:
+            try:
+                ack_msg = {
+                    "type": "event",
+                    "event_id": str(uuid.uuid4()),
+                    "event": "module.shutdown.ack",
+                    "source": "event_hub",
+                    "timestamp": datetime.now(timezone.utc).isoformat(),
+                    "data": {"module_id": "event_hub", "estimated_cleanup": 2},
+                }
+                await ws.send_text(json.dumps(ack_msg))
+            except Exception as e:
+                print(f"[event_hub] Failed to send shutdown ack: {e}")
+        # Cleanup
+        if self._timer_task:
+            self._timer_task.cancel()
+        if self._http_client:
+            await self._http_client.aclose()
+        # Send ready
+        if ws:
+            try:
+                ready_msg = {
+                    "type": "event",
+                    "event_id": str(uuid.uuid4()),
+                    "event": "module.shutdown.ready",
+                    "source": "event_hub",
+                    "timestamp": datetime.now(timezone.utc).isoformat(),
+                    "data": {"module_id": "event_hub"},
+                }
+                await ws.send_text(json.dumps(ready_msg))
+            except Exception as e:
+                print(f"[event_hub] Failed to send shutdown ready: {e}")
+        print("[event_hub] Shutdown ready, exiting")
+        # Trigger uvicorn exit
+        if self._uvicorn_server:
+            self._uvicorn_server.should_exit = True
+
+    async def _register_to_registry(self):
+        """Register to Registry. Triggered after Launcher connects."""
+        if self._registered_to_registry:
+            return
+        payload = {
+            "action": "register",
+            "module_id": "event_hub",
+            "module_type": "infrastructure",
+            "name": "Event Hub",
+            "api_endpoint": f"http://{self.advertise_ip}:{self.port}",
+            "health_endpoint": "/health",
+            "metadata": {
+                "ws_endpoint": f"ws://{self.advertise_ip}:{self.port}/ws",
+            },
+        }
+        headers = {"Authorization": f"Bearer {self.own_token}"}
         try:
             async with httpx.AsyncClient() as client:
                 resp = await client.post(
-                    f"{self.registry_url}/verify",
-                    json={"token": token},
-                    headers={"Authorization": f"Bearer {self.own_token}"},
+                    f"{self.registry_url}/modules",
+                    json=payload,
+                    headers=headers,
                     timeout=5,
                 )
                 if resp.status_code == 200:
-                    body = resp.json()
-                    if body.get("ok"):
-                        return body.get("module_id")
+                    self._registered_to_registry = True
+                    print("[event_hub] Registered to Registry")
+                else:
+                    print(f"[event_hub] WARNING: Registry returned {resp.status_code}")
         except Exception as e:
-            print(f"[event_hub] Token verification failed: {e}")
-        return None
+            print(f"[event_hub] WARNING: Failed to register to Registry: {e}")
 
     # ── App factory ──
 
@@ -52,15 +184,14 @@ class EventHubServer:
 
         @app.on_event("startup")
         async def _startup():
+            # Initialize HTTP client early to avoid blocking first token verification
+            server._http_client = httpx.AsyncClient()
             server._timer_task = asyncio.create_task(server._timer_loop())
-            server._test_task = asyncio.create_task(server._test_event_loop())
 
         @app.on_event("shutdown")
         async def _shutdown():
             if server._timer_task:
                 server._timer_task.cancel()
-            if hasattr(server, '_test_task') and server._test_task:
-                server._test_task.cancel()
 
         # ── WebSocket endpoint ──
 
@@ -70,12 +201,26 @@ class EventHubServer:
             mid_hint = ws.query_params.get("id", "")
             module_id = await server._verify_token(token, mid_hint)
             if not module_id:
-                await ws.close(code=4001, reason="Authentication failed")
+                # Must accept before close — Starlette drops TCP without close frame otherwise,
+                # causing websockets 15.x clients to get "no close frame received or sent" errors.
+                await ws.accept()
+                print(f"[event_hub] 认证失败: token={token[:8]}... hint={mid_hint}")
+                try:
+                    await ws.close(code=4001, reason="Authentication failed")
+                except Exception:
+                    pass
                 return
 
             await ws.accept()
             server.hub.add_connection(module_id, ws)
 
+            # Launcher bootstrap: first connection triggers module.ready + Registry registration
+            if (module_id == "launcher"
+                    and not server._launcher_connected
+                    and server.launcher_ws_token
+                    and token == server.launcher_ws_token):
+                await server._on_launcher_connected(ws)
+
             try:
                 while True:
                     raw = await ws.receive_text()
@@ -83,7 +228,10 @@ class EventHubServer:
             except WebSocketDisconnect:
                 pass
             except Exception as e:
-                print(f"[event_hub] WebSocket error for {module_id}: {e}")
+                # Connection-state errors are normal during shutdown
+                err = str(e).lower()
+                if "not connected" not in err and "closed" not in err:
+                    print(f"[event_hub] WebSocket error for {module_id}: {e}")
             finally:
                 server.hub.remove_connection(module_id)
 
@@ -104,10 +252,9 @@ class EventHubServer:
     async def _timer_loop(self):
         while True:
             await asyncio.sleep(30)
-            # Dedup cleanup (offload to thread to avoid blocking event loop)
             await asyncio.get_event_loop().run_in_executor(None, self.hub.dedup.cleanup)
-            # Heartbeat to Registry
-            await self._heartbeat()
+            if self._registered_to_registry:
+                await self._heartbeat()
 
     async def _heartbeat(self):
         try:
@@ -120,19 +267,3 @@ class EventHubServer:
                 )
         except Exception:
             pass
-
-    async def _test_event_loop(self):
-        """Publish a test event every 5 seconds via internal routing."""
-        import uuid
-        from datetime import datetime, timezone
-        while True:
-            await asyncio.sleep(5)
-            msg = {
-                "type": "event",
-                "event_id": str(uuid.uuid4()),
-                "event": "event_hub.test",
-                "source": "event_hub",
-                "timestamp": datetime.now(timezone.utc).isoformat(),
-                "data": {"message": "heartbeat from event_hub"},
-            }
-            await self.hub._route_event("event_hub", msg)

package/core/kite_log.py

@@ -0,0 +1,241 @@
+"""
+Kite module logging utilities.
+Provides latest.log + crashes.jsonl infrastructure per the 异常处理规范.
+
+Each module calls init_module_log() and setup_exception_hooks() at startup.
+All paths derived from KITE_MODULE_DATA environment variable.
+
+Note: This is a convenience library for Python modules. Node.js/binary modules
+implement the same spec independently (paths + formats are language-agnostic).
+"""
+
+import json
+import os
+import re
+import sys
+import threading
+import traceback
+from datetime import datetime, timezone
+
+
+# ── Module-level state (initialized by init_module_log) ──
+
+_module_name = ""
+_log_dir = None
+_log_latest_path = None
+_log_daily_path = None
+_log_daily_date = ""
+_crash_log_path = None
+_log_lock = threading.Lock()
+
+# Strip ANSI escape sequences for plain-text log files
+_ANSI_RE = re.compile(r"\033\[[0-9;]*m")
+
+
+def _strip_ansi(s: str) -> str:
+    return _ANSI_RE.sub("", s)
+
+
+def _resolve_daily_log_path():
+    """Resolve daily log path based on current date."""
+    global _log_daily_path, _log_daily_date
+    if not _log_dir:
+        return
+    today = datetime.now().strftime("%Y-%m-%d")
+    if today == _log_daily_date and _log_daily_path:
+        return
+    month_dir = os.path.join(_log_dir, today[:7])
+    os.makedirs(month_dir, exist_ok=True)
+    _log_daily_path = os.path.join(month_dir, f"{today}.log")
+    _log_daily_date = today
+
+
+def init_module_log(module_name: str):
+    """Initialize log files for a module. Call once at startup.
+
+    Requires KITE_MODULE_DATA environment variable to be set.
+    Creates: {KITE_MODULE_DATA}/log/latest.log (truncated)
+             {KITE_MODULE_DATA}/log/crashes.jsonl (truncated)
+             {KITE_MODULE_DATA}/log/{YYYY-MM}/ (daily archive directory)
+    """
+    global _module_name, _log_dir, _log_latest_path, _crash_log_path
+
+    _module_name = module_name
+    module_data = os.environ.get("KITE_MODULE_DATA")
+    if not module_data:
+        return
+
+    _log_dir = os.path.join(module_data, "log")
+    os.makedirs(_log_dir, exist_ok=True)
+
+    suffix = os.environ.get("KITE_INSTANCE_SUFFIX", "")
+
+    # latest.log — truncate on each startup
+    _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
+
+    # crashes.jsonl — truncate on each startup
+    _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
+
+    # Ensure daily archive directory exists
+    _resolve_daily_log_path()
+
+
+def write_log(plain_line: str):
+    """Write a plain-text line to both latest.log and daily log (open-write-close)."""
+    with _log_lock:
+        if _log_latest_path:
+            try:
+                with open(_log_latest_path, "a", encoding="utf-8") as f:
+                    f.write(plain_line)
+            except Exception:
+                pass
+        _resolve_daily_log_path()
+        if _log_daily_path:
+            try:
+                with open(_log_daily_path, "a", encoding="utf-8") as f:
+                    f.write(plain_line)
+            except Exception:
+                pass
+
+
+def write_crash(exc_type, exc_value, exc_tb,
+                thread_name=None, severity="critical", handled=False):
+    """Write crash record to crashes.jsonl + daily crash archive."""
+    record = {
+        "timestamp": datetime.now(timezone.utc).isoformat(),
+        "module": _module_name,
+        "thread": thread_name or threading.current_thread().name,
+        "exception_type": exc_type.__name__ if exc_type else "Unknown",
+        "exception_message": str(exc_value),
+        "traceback": "".join(traceback.format_exception(exc_type, exc_value, exc_tb)),
+        "severity": severity,
+        "handled": handled,
+        "process_id": os.getpid(),
+        "platform": sys.platform,
+        "runtime_version": f"Python {sys.version.split()[0]}",
+    }
+
+    if exc_tb:
+        tb_entries = traceback.extract_tb(exc_tb)
+        if tb_entries:
+            last = tb_entries[-1]
+            record["context"] = {
+                "function": last.name,
+                "file": os.path.basename(last.filename),
+                "line": last.lineno,
+            }
+
+    line = json.dumps(record, ensure_ascii=False) + "\n"
+
+    # 1. Write to crashes.jsonl (current run)
+    if _crash_log_path:
+        try:
+            with open(_crash_log_path, "a", encoding="utf-8") as f:
+                f.write(line)
+        except Exception:
+            pass
+
+    # 2. Write to daily crash archive
+    if _log_dir:
+        try:
+            today = datetime.now().strftime("%Y-%m-%d")
+            archive_dir = os.path.join(_log_dir, "crashes", today[:7])
+            os.makedirs(archive_dir, exist_ok=True)
+            archive_path = os.path.join(archive_dir, f"{today}.jsonl")
+            with open(archive_path, "a", encoding="utf-8") as f:
+                f.write(line)
+        except Exception:
+            pass
+
+
+def print_crash_summary(exc_type, exc_tb, thread_name=None):
+    """Print crash summary to console (red highlight)."""
+    RED = "\033[91m"
+    RESET = "\033[0m"
+
+    if exc_tb:
+        tb_entries = traceback.extract_tb(exc_tb)
+        if tb_entries:
+            last = tb_entries[-1]
+            location = f"{os.path.basename(last.filename)}:{last.lineno}"
+        else:
+            location = "unknown"
+    else:
+        location = "unknown"
+
+    prefix = f"[{_module_name}]"
+    if thread_name:
+        print(f"{prefix} {RED}线程 {thread_name} 崩溃: "
+              f"{exc_type.__name__} in {location}{RESET}")
+    else:
+        print(f"{prefix} {RED}崩溃: {exc_type.__name__} in {location}{RESET}")
+    if _crash_log_path:
+        print(f"{prefix} 崩溃日志: {_crash_log_path}")
+
+
+def setup_exception_hooks(module_name: str = ""):
+    """Set up global exception hooks (sys.excepthook + threading.excepthook).
+    Call once at startup after init_module_log().
+    """
+    if module_name:
+        global _module_name
+        _module_name = module_name
+
+    _orig_excepthook = sys.excepthook
+
+    def _excepthook(exc_type, exc_value, exc_tb):
+        write_crash(exc_type, exc_value, exc_tb, severity="critical", handled=False)
+        print_crash_summary(exc_type, exc_tb)
+        _orig_excepthook(exc_type, exc_value, exc_tb)
+
+    sys.excepthook = _excepthook
+
+    if hasattr(threading, "excepthook"):
+        def _thread_excepthook(args):
+            write_crash(args.exc_type, args.exc_value, args.exc_traceback,
+                        thread_name=args.thread.name if args.thread else "unknown",
+                        severity="error", handled=False)
+            print_crash_summary(args.exc_type, args.exc_traceback,
+                                thread_name=args.thread.name if args.thread else None)
+
+        threading.excepthook = _thread_excepthook
+
+
+def install_log_hook():
+    """Install a hook on the _SafeWriter or stdout to capture all print output
+    to latest.log + daily log. Call after init_module_log().
+
+    Wraps the current sys.stdout so that every write() also goes to log files.
+    """
+    if not _log_latest_path:
+        return
+
+    original_stdout = sys.stdout
+
+    class _LoggingWriter:
+        """Wraps stdout to tee all output to log files."""
+        def __init__(self, stream):
+            self._stream = stream
+
+        def write(self, s):
+            self._stream.write(s)
+            if s and s.strip():
+                write_log(_strip_ansi(s) + ("\n" if not s.endswith("\n") else ""))
+
+        def flush(self):
+            self._stream.flush()
+
+        def __getattr__(self, name):
+            return getattr(self._stream, name)
+
+    sys.stdout = _LoggingWriter(original_stdout)