@agentunion/kite 1.5.0 → 1.6.1

package/extensions/services/audit/server.py

@@ -0,0 +1,542 @@
+"""
+Audit service HTTP server.
+Provides audit query APIs.
+"""
+
+import asyncio
+import json
+import os
+import time
+import uuid
+import random
+
+import websockets
+from fastapi import FastAPI
+
+from collector import AuditCollector
+from storage import AuditStorage
+from routes.routes_audit import router as audit_router
+
+
+class AuditServer:
+    """Audit HTTP 服务器"""
+
+    def __init__(self, token: str = "", kernel_port: int = 0,
+                 host: str = "127.0.0.1", port: int = 0, boot_t0: float = 0,
+                 max_connections: int = 1,
+                 gateway_url: str = "", kite_token: str = ""):
+        self.token = token
+        self.kernel_port = kernel_port
+        self.host = host
+        self.port = port
+        self.boot_t0 = boot_t0
+        self.max_connections = max_connections
+        self.gateway_url = gateway_url
+        self.kite_token = kite_token
+        self._ws_task: asyncio.Task | None = None
+        self._ws: object | None = None
+        self._shutting_down = False
+        self._exit_code = 0
+        self._auth_failed = False
+        self._uvicorn_server = None
+        self._start_time = time.time()
+        self._reconnect_attempt = 0
+        self._pending_rpc: dict[str, asyncio.Future] = {}
+        self._extra_ws: dict = {}           # slot → WebSocket
+        self._extra_ws_tasks: dict = {}     # slot → recv loop Task
+        self._has_registered = False
+
+        # 初始化存储和收集器
+        instance_dir = os.environ.get("KITE_INSTANCE_DIR", "")
+        audit_dir = os.path.join(instance_dir, "audit")
+        self.storage = AuditStorage(audit_dir)
+        self.collector = AuditCollector(on_event=self._on_audit_event)
+
+        self.app = self._create_app()
+
+    def _create_app(self) -> FastAPI:
+        app = FastAPI(title="Kite Audit Service", docs_url="/docs", redoc_url=None)
+        server = self
+
+        @app.on_event("startup")
+        async def _startup():
+            print("[audit] HTTP server started")
+            # 启动 Kernel WebSocket 连接
+            server._ws_task = asyncio.create_task(server._ws_client_loop())
+            # 启动收集器
+            server.collector.start()
+
+        @app.on_event("shutdown")
+        async def _shutdown():
+            print("[audit] HTTP server shutting down")
+            # 停止收集器
+            server.collector.stop()
+            # 关闭 WebSocket
+            if server._ws_task:
+                server._ws_task.cancel()
+
+        # Health check
+        @app.get("/health")
+        async def health():
+            return {
+                "status": "healthy",
+                "uptime_seconds": round(time.time() - server._start_time),
+            }
+
+        # Status
+        @app.get("/status")
+        async def status():
+            return {
+                "status": "running",
+                "uptime_seconds": round(time.time() - server._start_time),
+                "storage_stats": server.storage.get_stats(),
+            }
+
+        # 挂载路由
+        app.include_router(audit_router, prefix="/api")
+
+        return app
+
+    def _on_audit_event(self, event):
+        """审计事件回调"""
+        self.storage.append(event)
+
+    async def _do_init(self, ws):
+        """收到 system.require_init 后执行：订阅 + 注册 + module.ready"""
+        reason = "startup" if not self._has_registered else "recovery"
+        print(f"[audit] Received system.require_init (reason={reason})")
+
+        try:
+            # 订阅事件（必须在注册之前）
+            await self._subscribe_events(ws)
+
+            # 注册到 Registry
+            await self._register_to_kernel(ws)
+
+            # 发布 module.ready
+            await self._publish_ready(ws, reason)
+
+            self._has_registered = True
+        except Exception as e:
+            print(f"[audit] _do_init failed: {e}")
+
+    async def _ws_client_loop(self):
+        """Kernel WebSocket 客户端循环（带指数退避重连）"""
+        kernel_url = f"ws://127.0.0.1:{self.kernel_port}/ws?id=audit"
+        retry_delay = 0.5
+        max_delay = 10.0
+        cooldown_attempts = 0
+
+        while not self._shutting_down:
+            try:
+                async with websockets.connect(kernel_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": self.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']}")
+
+                    self._ws = ws
+                    print(f"[audit] Connected to Kernel")
+
+                    # 重置重连计数
+                    self._reconnect_attempt = 0
+                    retry_delay = 0.5
+                    cooldown_attempts = 0
+
+                    # 启动接收任务（必须在 subscribe/register 之前，以便接收响应）
+                    receiver_task = asyncio.create_task(self._ws_receiver(ws))
+
+                    try:
+                        # 等待接收任务结束（system.require_init 会触发 _do_init）
+                        await receiver_task
+                    finally:
+                        receiver_task.cancel()
+                        self._ws = None
+                        # 清理断连时残留的 pending RPC
+                        for fut in self._pending_rpc.values():
+                            if not fut.done():
+                                fut.cancel()
+                        self._pending_rpc.clear()
+
+            except Exception as e:
+                if self._shutting_down:
+                    return
+
+                code = self._get_close_code(e)
+
+                # never: 永不重连
+                if code in (4001, 4003, 4004, 1008, 4010):
+                    print(f"[audit] 致命错误 (code {code})，退出")
+                    self._exit_code = 1
+                    self._auth_failed = True
+                    self._shutting_down = True
+                    if self._uvicorn_server:
+                        self._uvicorn_server.should_exit = True
+                    return
+
+                # cooldown: 速率限制
+                if code == 4020:
+                    cooldown_attempts += 1
+                    if cooldown_attempts >= 5:
+                        print(f"[audit] 速率限制重试 5 次，退出")
+                        self._exit_code = 1
+                        self._shutting_down = True
+                        if self._uvicorn_server:
+                            self._uvicorn_server.should_exit = True
+                        return
+                    print(f"[audit] 速率限制，10.0s 后重试 ({cooldown_attempts}/5)")
+                    await asyncio.sleep(10.0)
+                    continue
+
+                # standard: 指数退避 + jitter
+                self._reconnect_attempt += 1
+                jitter = retry_delay * 0.2 * random.random()
+                sleep_time = retry_delay + jitter
+                print(f"[audit] WebSocket 连接错误: {e}, {sleep_time:.1f}s 后重试 (attempt {self._reconnect_attempt})")
+
+                # 第 5 次重连失败时给出诊断提示
+                if self._reconnect_attempt == 5:
+                    print(f"\033[33m[audit] 提示: 已连续 {self._reconnect_attempt} 次无法连接 Kernel (端口 {self.kernel_port})")
+                    if self.kernel_port < 1024:
+                        print(f"[audit] ⚠ 端口 {self.kernel_port} 异常偏低，可能是 Kernel 端口绑定失败或配置错误")
+                    print(f"[audit] 请检查: 1) Kernel 进程是否存活  2) kernel/module.md 中 preferred_port 配置是否正确\033[0m")
+
+                self._ws = None
+                await asyncio.sleep(sleep_time)
+                retry_delay = min(retry_delay * 2, max_delay)
+
+    async def _ws_receiver(self, ws):
+        """主连接消息接收循环（三路由：事件通知 / RPC 请求 / RPC 响应）"""
+        try:
+            async for raw in ws:
+                try:
+                    msg = json.loads(raw)
+                except (json.JSONDecodeError, TypeError):
+                    continue
+                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(self._do_init(ws))
+                        continue
+                    asyncio.create_task(self._handle_event_notification(ws, msg))
+                elif has_method and has_id:
+                    asyncio.create_task(self._handle_rpc_request(ws, msg))
+                elif not has_method and has_id:
+                    self._handle_rpc_response(msg)
+        except websockets.exceptions.ConnectionClosed:
+            pass  # Kernel 连接关闭，正常退出接收循环
+
+    def _get_close_code(self, e: Exception) -> int:
+        """从 websockets 异常中提取关闭码"""
+        if hasattr(e, 'rcvd') and e.rcvd is not None:
+            return getattr(e.rcvd, 'code', 0)
+        return 0
+
+    def _is_auth_failure(self, e: Exception) -> bool:
+        """检测是否为认证失败（WebSocket 4001/4003 关闭码）"""
+        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 _register_to_kernel(self, ws):
+        """注册到 Kernel（结构化格式）"""
+        await self._rpc_call(ws, "registry.register", {
+            "module_id": "audit",
+            "module_type": "service",
+            "tools": {
+                "rpc": {
+                    "module": {
+                        "health": {"method": "health", "description": "健康检查"},
+                        "status": {"method": "status", "description": "状态查询"}
+                    }
+                }
+            },
+            "events_publish": {},
+            "events_subscribe": [
+                "module.shutdown",
+            ] + self.collector.AUDIT_EVENTS,
+            "http_port": self.port,
+        })
+
+    async def _subscribe_events(self, ws):
+        """订阅事件"""
+        await self._rpc_call(ws, "event.subscribe", {
+            "events": ["module.shutdown"] + self.collector.AUDIT_EVENTS,
+        })
+
+    async def _publish_ready(self, ws, reason: str = "startup"):
+        """发布 module.ready 事件（带启动时间和优雅退出标记）"""
+        startup_time = time.monotonic() - self.boot_t0
+        await self._publish_event(ws, "module.ready", {
+            "module_id": "audit",
+            "graceful_shutdown": True,
+            "startup_time": startup_time,
+            "reason": reason,
+        })
+
+    async def _handle_message(self, ws, msg: dict):
+        """处理 WebSocket 消息（区分 RPC 请求、RPC 响应、事件通知）"""
+        has_method = "method" in msg
+        has_id = "id" in msg
+        has_result_or_error = "result" in msg or "error" in msg
+
+        # JSON-RPC 2.0 协议分类：
+        # 1. 事件通知：有 method，无 id（或 method == "event"）
+        # 2. RPC 请求：有 method，有 id
+        # 3. RPC 响应：无 method，有 id，有 result/error
+
+        if has_method and not has_id:
+            # 事件通知（Notification）— run in background to prevent deadlock
+            asyncio.create_task(self._handle_event_notification(ws, msg))
+        elif has_method and has_id:
+            # RPC 请求（需要响应）— run in background to prevent deadlock
+            asyncio.create_task(self._handle_rpc_request(ws, msg))
+        elif has_id and has_result_or_error:
+            # RPC 响应 — 解析 pending future
+            self._handle_rpc_response(msg)
+
+    async def _handle_event_notification(self, ws, msg: dict):
+        """处理事件通知"""
+        params = msg.get("params", {})
+        event_type = params.get("event", "")
+        data = params.get("data", {})
+        source = params.get("source", "unknown")
+
+        # 处理 system.ping 事件
+        if event_type == "system.ping":
+            await self._handle_ping_event(ws, data)
+            return
+
+        # 弹性连接 offer/release
+        if event_type == "system.connection.offer":
+            asyncio.create_task(self._handle_connection_offer(data))
+            return
+        if event_type == "system.connection.release":
+            asyncio.create_task(self._handle_connection_release(data))
+            return
+
+        # 处理 shutdown 事件
+        if event_type == "module.shutdown":
+            target = data.get("module_id", "")
+            reason = data.get("reason", "")
+            # 处理目标关闭（module_id == "audit"）或广播关闭（无 module_id 或 launcher_lost）
+            if target == "audit" or not target or reason == "launcher_lost":
+                await self._handle_shutdown()
+                return
+
+        # 传递给收集器
+        await self.collector.handle_kernel_event(event_type, data, source)
+
+    async def _handle_ping_event(self, ws, data: dict):
+        """处理 system.ping 事件并回复 system.pong"""
+        import time
+        t1 = data.get("ping_time")
+        t2 = time.time()
+
+        await self._publish_event(ws, "system.pong", {
+            "module_id": "audit",
+            "ping_time": t1,
+            "pong_time": t2,
+        })
+
+    async def _handle_rpc_request(self, ws, msg: dict):
+        """处理 RPC 请求"""
+        rpc_id = msg.get("id", "")
+        method = msg.get("method", "")
+        params = msg.get("params", {})
+
+        handlers = {
+            "health": self._rpc_health,
+            "status": self._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 _handle_connection_offer(self, 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 self._extra_ws:
+                continue
+            asyncio.create_task(self._connect_slot(slot, token))
+
+    async def _connect_slot(self, slot, token):
+        """建立单个 slot 附加连接。"""
+        ws_url = f"ws://127.0.0.1:{self.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
+            self._extra_ws[slot] = ws
+            self._extra_ws_tasks[slot] = asyncio.create_task(self._slot_recv_loop(slot, ws))
+            print(f"[audit] Slot {slot} connected")
+        except Exception as e:
+            print(f"[audit] Slot {slot} connect failed: {e}")
+
+    async def _slot_recv_loop(self, slot, ws):
+        """附加 slot 接收循环：与主连接平等处理所有消息。"""
+        try:
+            async for raw in ws:
+                try:
+                    msg = json.loads(raw)
+                except (json.JSONDecodeError, TypeError):
+                    continue
+
+                try:
+                    await self._handle_message(ws, msg)
+                except Exception as e:
+                    print(f"[audit] Slot {slot} 消息处理异常: {e}")
+        except Exception as e:
+            print(f"[audit] Slot {slot} recv loop exited: {e}")
+        finally:
+            self._extra_ws.pop(slot, None)
+            self._extra_ws_tasks.pop(slot, None)
+
+    async def _handle_connection_release(self, data):
+        """Kernel 请求释放 slot，优雅关闭。"""
+        for slot in data.get("slots", []):
+            ws = self._extra_ws.pop(slot, None)
+            task = self._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(self):
+        """处理关闭请求（三步流程：ack → exiting → ready）"""
+        print("[audit] Received shutdown request")
+        self._shutting_down = True
+
+        # Step 1: 发送 ack（立即确认收到）
+        if self._ws:
+            await self._publish_event(self._ws, "module.shutdown.ack", {
+                "module_id": "audit",
+            })
+
+        # Step 2: 发送 exiting（开始清理）
+        if self._ws:
+            await self._publish_event(self._ws, "module.exiting", {
+                "module_id": "audit",
+                "type": "passive",
+                "reason": "shutdown_requested",
+                "restart": "auto",
+                "action": "none",
+                "timeout": 2.0,
+                "restart_delay": 0.0,
+            })
+
+        # Step 3: 清理资源
+        self.collector.stop()
+
+        # Step 4: 发送 ready（必须在关闭连接之前发送）
+        if self._ws:
+            await self._publish_event(self._ws, "module.shutdown.ready", {
+                "module_id": "audit",
+            })
+
+        print("[audit] Shutdown ready, exiting")
+
+        # 等待 Kernel 处理 shutdown.ready
+        await asyncio.sleep(0.01)
+
+        # Step 5: 关闭所有附加连接
+        for _s, _w in list(self._extra_ws.items()):
+            try:
+                await _w.close(code=1000, reason="shutdown")
+            except Exception:
+                pass
+        for _t in self._extra_ws_tasks.values():
+            _t.cancel()
+        self._extra_ws.clear()
+        self._extra_ws_tasks.clear()
+
+        # Step 6: 关闭 HTTP 服务器
+        if self._uvicorn_server:
+            self._uvicorn_server.should_exit = True
+
+    async def _rpc_health(self, params: dict) -> dict:
+        """RPC: health"""
+        return {
+            "status": "healthy",
+            "uptime_seconds": round(time.time() - self._start_time),
+        }
+
+    async def _rpc_status(self, params: dict) -> dict:
+        """RPC: status"""
+        return {
+            "status": "running",
+            "uptime_seconds": round(time.time() - self._start_time),
+            "storage_stats": self.storage.get_stats(),
+        }
+
+    async def _rpc_call(self, 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
+        if not wait_response:
+            await ws.send(json.dumps(msg))
+            return {}
+        future = asyncio.get_event_loop().create_future()
+        self._pending_rpc[rpc_id] = future
+        await ws.send(json.dumps(msg))
+        try:
+            return await asyncio.wait_for(future, timeout=timeout)
+        except asyncio.TimeoutError:
+            self._pending_rpc.pop(rpc_id, None)
+            return {"error": {"code": -32000, "message": f"RPC timeout: {method} ({timeout}s)"}}
+
+    def _handle_rpc_response(self, msg: dict):
+        """处理 RPC 响应，解析 pending future。"""
+        rpc_id = msg.get("id")
+        future = self._pending_rpc.pop(rpc_id, None)
+        if future and not future.done():
+            future.set_result(msg)
+
+    async def _publish_event(self, ws, event: str, data: dict = None):
+        """Publish event via event.publish RPC (fire-and-forget)."""
+        await self._rpc_call(ws, "event.publish", {
+            "event_id": str(uuid.uuid4()),
+            "event": event,
+            "data": data or {},
+        }, wait_response=False)

package/extensions/services/audit/storage.py

@@ -0,0 +1,95 @@
+"""
+Audit event storage (append-only).
+"""
+
+import json
+import os
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Optional
+
+from schemas.audit_event import AuditEvent
+
+
+class AuditStorage:
+    """审计事件存储（append-only）"""
+
+    def __init__(self, base_dir: str):
+        self.base_dir = Path(base_dir)
+        self.events_dir = self.base_dir / "events"
+        self.events_dir.mkdir(parents=True, exist_ok=True)
+
+    def _get_daily_file(self, date: Optional[datetime] = None) -> Path:
+        """获取指定日期的审计日志文件路径"""
+        if date is None:
+            date = datetime.now(timezone.utc)
+
+        year_month = date.strftime("%Y-%m")
+        day = date.strftime("%Y-%m-%d")
+
+        month_dir = self.events_dir / year_month
+        month_dir.mkdir(parents=True, exist_ok=True)
+
+        return month_dir / f"{day}.jsonl"
+
+    def append(self, event: AuditEvent):
+        """追加审计事件（append-only）"""
+        log_file = self._get_daily_file()
+
+        with open(log_file, "a", encoding="utf-8") as f:
+            f.write(json.dumps(event.to_dict(), ensure_ascii=False) + "\n")
+
+    def read_events(
+        self,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[dict]:
+        """读取审计事件"""
+        events = []
+
+        # 如果没有指定日期范围，读取今天的日志
+        if start_date is None and end_date is None:
+            log_file = self._get_daily_file()
+            if log_file.exists():
+                events.extend(self._read_file(log_file))
+        else:
+            # TODO: 实现日期范围查询
+            pass
+
+        # 应用分页
+        return events[offset:offset + limit]
+
+    def _read_file(self, file_path: Path) -> list[dict]:
+        """读取单个 JSONL 文件"""
+        events = []
+
+        if not file_path.exists():
+            return events
+
+        with open(file_path, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    try:
+                        events.append(json.loads(line))
+                    except json.JSONDecodeError:
+                        continue
+
+        return events
+
+    def get_stats(self) -> dict:
+        """获取存储统计信息"""
+        total_files = 0
+        total_size = 0
+
+        for file_path in self.events_dir.rglob("*.jsonl"):
+            total_files += 1
+            total_size += file_path.stat().st_size
+
+        return {
+            "total_files": total_files,
+            "total_size_bytes": total_size,
+            "total_size_mb": round(total_size / 1024 / 1024, 2),
+        }