@agentunion/kite 1.4.0 → 1.6.0

package/extensions/services/evolmem/server.py

@@ -0,0 +1,641 @@
+"""
+EvolMem — FastAPI HTTP 服务 + Kite WebSocket 客户端。
+所有 hook 逻辑通过直接调用 hooks/ 子包完成，无子进程。
+"""
+
+import asyncio
+import json
+import os
+import random
+import time
+import uuid
+from pathlib import Path
+
+try:
+    import websockets
+except ImportError:
+    raise ImportError("websockets is required: pip install websockets")
+
+from fastapi import FastAPI, Request, Response
+from fastapi.responses import JSONResponse
+
+# Hook 处理器
+from .hooks import session_start, user_prompt, assistant_stop, pre_tool_use, session_end
+from .hooks.common import EVOLMEM_DATA_DIR
+from .scripts.search_topics import search_topics, format_results
+from .scripts.extract_keywords import extract_hot_keywords
+
+# Claude Code 全局 settings 路径
+_CLAUDE_SETTINGS = Path.home() / ".claude" / "settings.json"
+
+_HOOK_EVENTS = ["SessionStart", "UserPromptSubmit", "PreToolUse", "Stop", "SessionEnd"]
+
+
+class EvolMemServer:
+    def __init__(self, token, kernel_port, host, port, boot_t0,
+                 max_connections=1, max_instances=1, instance_elastic=False,
+                 gateway_url="", kite_token=""):
+        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._max_instances    = max_instances
+        self._instance_elastic = instance_elastic
+
+        self._module_name  = "evolmem"
+        self._module_type  = "service"
+        self._display_name = "EvolMem"
+
+        self._ws               = None
+        self._ws_task          = None
+        self._test_task        = None
+        self._uvicorn_server   = None
+        self._shutting_down    = False
+        self._exit_code        = 0
+        self._start_time       = time.time()
+        self._pending_rpc: dict[str, asyncio.Future] = {}
+        self._extra_ws: dict         = {}
+        self._extra_ws_tasks: dict   = {}
+        self._handler_capacity = 100
+        self._active_handlers  = 0
+
+        self.app = self._create_app()
+
+    # ════════════════════════════════════════════
+    #  FastAPI 应用
+    # ════════════════════════════════════════════
+
+    def _create_app(self) -> FastAPI:
+        app = FastAPI(title="evolmem")
+        server = self
+
+        @app.on_event("startup")
+        async def _startup():
+            server._ws_task   = asyncio.create_task(server._ws_loop())
+            server._test_task = asyncio.create_task(server._test_event_loop())
+
+        @app.on_event("shutdown")
+        async def _shutdown():
+            server._shutting_down = True
+            for t in (server._ws_task, server._test_task):
+                if t: t.cancel()
+            for _t in server._extra_ws_tasks.values():
+                _t.cancel()
+            server._extra_ws_tasks.clear()
+            if server._ws:
+                try: await server._ws.close()
+                except Exception: pass
+            print("[evolmem] Shutdown complete")
+
+        # ── 健康 & 状态 ──
+
+        @app.get("/health")
+        async def health():
+            return {"status": "ok", "version": "2.0.0", "module": "evolmem"}
+
+        @app.get("/api/status")
+        async def api_status(request: Request):
+            try:
+                project_name = Path(request.query_params.get("project", ".")).resolve().name
+                summaries_dir = EVOLMEM_DATA_DIR / project_name / "summaries"
+                return JSONResponse({
+                    "project_name": project_name,
+                    "topic_count":  len(list(summaries_dir.glob("topic-*.md"))) if summaries_dir.exists() else 0,
+                    "recent_topics": [],
+                    "hot_keywords":  [],
+                })
+            except Exception as e:
+                return JSONResponse({"error": str(e)}, status_code=500)
+
+        # ── 搜索 ──
+
+        @app.get("/api/search")
+        async def api_search(request: Request):
+            try:
+                q = request.query_params.get("q", "")
+                patterns = [k for k in q.split(",") if k.strip()]
+                if not patterns:
+                    return Response("请提供搜索关键词，如 ?q=关键词", media_type="text/plain; charset=utf-8")
+
+                project_name  = Path(request.query_params.get("project", ".")).resolve().name
+                summaries_dir = EVOLMEM_DATA_DIR / project_name / "summaries"
+
+                results = await asyncio.to_thread(search_topics, patterns, summaries_dir)
+                return Response(format_results(results), media_type="text/plain; charset=utf-8")
+            except Exception as e:
+                return JSONResponse({"error": str(e)}, status_code=500)
+
+        # ── 项目列表 ──
+
+        @app.get("/api/projects")
+        async def api_projects(request: Request):
+            try:
+                current = request.query_params.get("current")
+                projects = await asyncio.to_thread(_list_projects, current)
+                return JSONResponse({"projects": projects})
+            except Exception as e:
+                return JSONResponse({"error": str(e)}, status_code=500)
+
+        # ── 关键词 ──
+
+        @app.get("/api/keywords")
+        async def api_keywords(request: Request):
+            try:
+                project_name = Path(request.query_params.get("project", ".")).resolve().name
+                index_file = EVOLMEM_DATA_DIR / project_name / "summaries" / "index.jsonl"
+                results = await asyncio.to_thread(extract_hot_keywords, index_file)
+                return JSONResponse([{"keyword": k, "count": c} for k, c in results])
+            except Exception as e:
+                return JSONResponse({"error": str(e)}, status_code=500)
+
+        # ── Claude Code Hook 事件接收 ──
+
+        _HOOK_DISPATCH = {
+            "SessionStart":     session_start.handle,
+            "UserPromptSubmit": user_prompt.handle,
+            "PreToolUse":       pre_tool_use.handle,
+            "Stop":             assistant_stop.handle,
+            "SessionEnd":       session_end.handle,
+        }
+
+        @app.post("/hooks/event/{event_name}")
+        async def hook_event(event_name: str, request: Request):
+            handler = _HOOK_DISPATCH.get(event_name)
+            if not handler:
+                return JSONResponse({})
+            try:
+                body = await request.body()
+                context = json.loads(body) if body else {}
+                result = await asyncio.to_thread(handler, context)
+                # 剥离 _debug_info，不发给 Claude
+                result.pop("_debug_info", None)
+                return JSONResponse(result)
+            except Exception as e:
+                print(f"[evolmem] Hook {event_name} error: {e}")
+                return JSONResponse({})
+
+        # ── Hook 安装 / 卸载 ──
+
+        @app.post("/api/hooks/install")
+        async def hooks_install(request: Request):
+            try:
+                body = await request.json()
+                port = body.get("port", server._port)
+                result = await asyncio.to_thread(_install_hooks, port)
+                return JSONResponse(result)
+            except Exception as e:
+                return JSONResponse({"success": False, "error": str(e)}, status_code=500)
+
+        @app.post("/api/hooks/uninstall")
+        async def hooks_uninstall():
+            try:
+                result = await asyncio.to_thread(_uninstall_hooks)
+                return JSONResponse(result)
+            except Exception as e:
+                return JSONResponse({"success": False, "error": str(e)}, status_code=500)
+
+        @app.get("/api/hooks/status")
+        async def hooks_status():
+            try:
+                return JSONResponse(_hooks_status())
+            except Exception as e:
+                return JSONResponse({"error": str(e)}, status_code=500)
+
+        return app
+
+    # ════════════════════════════════════════════
+    #  Kite WebSocket 生命周期
+    # ════════════════════════════════════════════
+
+    async def _ws_loop(self):
+        retry_delay, max_delay, max_retries = 0.3, 5.0, 10
+        for _ in range(max_retries):
+            try:
+                uri = f"ws://127.0.0.1:{self._kernel_port}/ws"
+                async with websockets.connect(uri, max_size=16 * 1024 * 1024) as ws:
+                    self._ws = ws
+                    retry_delay = 0.3
+                    await self._ws_connect(ws)
+            except websockets.ConnectionClosedError as e:
+                if e.code in (4001, 4003, 4004, 1008, 4010):
+                    print(f"[evolmem] Fatal close code={e.code}")
+                    return
+                if e.code == 4020:
+                    await asyncio.sleep(10)
+                    continue
+            except (ConnectionRefusedError, OSError) as e:
+                if not (isinstance(e, ConnectionRefusedError) or
+                        (hasattr(e, "winerror") and e.winerror == 1225)):
+                    print(f"[evolmem] Connection error: {e}")
+            except Exception as e:
+                print(f"[evolmem] Unexpected: {e}")
+
+            if self._shutting_down:
+                return
+            for fut in self._pending_rpc.values():
+                if not fut.done():
+                    fut.set_exception(ConnectionError("Disconnected"))
+            self._pending_rpc.clear()
+            jitter = retry_delay * 0.2 * random.random()
+            await asyncio.sleep(retry_delay + jitter)
+            retry_delay = min(retry_delay * 2, max_delay)
+
+    async def _ws_connect(self, ws):
+        await ws.send(json.dumps({
+            "jsonrpc": "2.0", "id": str(uuid.uuid4()),
+            "method": "auth", "params": {"token": self.token},
+        }))
+        resp = json.loads(await asyncio.wait_for(ws.recv(), timeout=5.0))
+        if resp.get("error") or not resp.get("result", {}).get("ok"):
+            print(f"[evolmem] Auth failed: {resp}")
+            return
+        print("[evolmem] Authenticated")
+
+        async for raw in ws:
+            msg = json.loads(raw)
+            has_method = "method" in msg
+            has_id     = "id" in msg
+            has_result = "result" in msg or "error" in msg
+
+            if has_method and has_id:
+                asyncio.create_task(self._handle_rpc_request(ws, msg))
+            elif has_method and not has_id:
+                asyncio.create_task(self._handle_event_notification(msg))
+            elif has_id and has_result:
+                self._handle_rpc_response(msg)
+
+    # ── RPC ──
+
+    async def _rpc_call(self, ws, method, params=None, wait_response=True, timeout=3.0):
+        req_id = str(uuid.uuid4())
+        await ws.send(json.dumps({"jsonrpc": "2.0", "id": req_id, "method": method, "params": params or {}}))
+        if not wait_response:
+            return None
+        future = asyncio.get_event_loop().create_future()
+        self._pending_rpc[req_id] = future
+        try:
+            return await asyncio.wait_for(future, timeout=timeout)
+        except asyncio.TimeoutError:
+            self._pending_rpc.pop(req_id, None)
+            return {"error": {"code": -32002, "message": f"Timeout: {method}"}}
+
+    def _handle_rpc_response(self, msg):
+        req_id = msg.get("id")
+        future = self._pending_rpc.pop(req_id, None)
+        if future and not future.done():
+            future.set_result(msg.get("result", msg) if "result" in msg else msg)
+
+    async def _handle_rpc_request(self, ws, msg):
+        method = msg.get("method", "")
+        params = msg.get("params", {})
+        req_id = msg.get("id")
+
+        if self._shutting_down:
+            await ws.send(json.dumps({
+                "jsonrpc": "2.0", "id": req_id,
+                "error": {"code": -32603, "message": "Module shutting down"},
+            }))
+            return
+
+        bare   = method.split(".")[-1] if "." in method else method
+        handlers = {"health": self._rpc_health, "status": self._rpc_status}
+        handler = handlers.get(bare)
+        if not handler:
+            await ws.send(json.dumps({
+                "jsonrpc": "2.0", "id": req_id,
+                "error": {"code": -32601, "message": f"Method not found: {method}"},
+            }))
+            return
+        self._active_handlers += 1
+        try:
+            result = await handler(params)
+            await ws.send(json.dumps({"jsonrpc": "2.0", "id": req_id, "result": result}))
+        except Exception as e:
+            await ws.send(json.dumps({
+                "jsonrpc": "2.0", "id": req_id,
+                "error": {"code": -32603, "message": str(e)},
+            }))
+        finally:
+            self._active_handlers -= 1
+
+    async def _rpc_health(self, params) -> dict:
+        try:
+            import psutil
+            mem_mb = round(psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024, 1)
+        except Exception:
+            mem_mb = -1
+        fw = {
+            "uptime": time.time() - self._start_time,
+            "connections": 1 + len(self._extra_ws),
+            "pending_rpc": len(self._pending_rpc),
+            "active_handlers": self._active_handlers,
+            "handler_capacity": self._handler_capacity,
+            "pressure": self._compute_pressure(),
+            "memory_mb": mem_mb,
+        }
+        biz = {
+            "healthy": EVOLMEM_DATA_DIR.exists() or True,
+            "message": "正常运行中",
+            "details": [
+                {"key": "HTTP 端口", "status": "ok", "value": str(self._port)},
+                {"key": "Hooks 安装", "status": "ok" if _hooks_status().get("installed", False) else "warn",
+                 "value": "已安装" if _hooks_status().get("installed", False) else "未安装"},
+                {"key": "数据目录", "status": "ok" if EVOLMEM_DATA_DIR.exists() else "warn",
+                 "value": str(EVOLMEM_DATA_DIR)},
+            ],
+        }
+        healthy = biz["healthy"] and fw["pressure"] != "critical"
+        return {"status": "healthy" if healthy else "unhealthy", "framework": fw, "business": biz}
+
+    async def _rpc_status(self, params) -> dict:
+        return {
+            "module": self._module_name, "state": "running",
+            "http_port": self._port,
+            "data_dir": str(EVOLMEM_DATA_DIR),
+            "hooks": _hooks_status(),
+        }
+
+    def _compute_pressure(self) -> str:
+        depth = self._active_handlers + len(self._pending_rpc)
+        ratio = depth / self._handler_capacity if self._handler_capacity else 0
+        if ratio >= 0.9: return "critical"
+        if ratio >= 0.5: return "overload"
+        if depth > 0:    return "normal"
+        return "idle"
+
+    # ── 事件处理 ──
+
+    async def _publish_event(self, ws, event, data=None):
+        await self._rpc_call(ws, "event.publish", {
+            "event_id": str(uuid.uuid4()), "event": event, "data": data or {},
+        }, wait_response=False)
+
+    async def _handle_event_notification(self, msg):
+        params     = msg.get("params", {})
+        event_type = params.get("event", "")
+        data       = params.get("data", {})
+
+        if event_type == "system.require_init":
+            asyncio.create_task(self._do_init(self._ws))
+        elif event_type == "system.ping":
+            await self._publish_event(self._ws, "system.pong", {
+                "module_id": self._module_name,
+                "ping_time": data.get("ping_time"),
+                "pong_time": time.time(),
+            })
+        elif event_type == "module.shutdown":
+            await self._handle_shutdown(self._ws)
+
+    # ── 初始化 ──
+
+    async def _do_init(self, ws):
+        await self._rpc_call(ws, "event.subscribe", {"events": [
+            "module.started", "module.stopped", "module.crashed",
+            "module.ready", "module.exiting", "module.shutdown",
+            "module.shutdown.ack", "module.shutdown.ready", "system.ping",
+        ]})
+        await self._rpc_call(ws, "registry.register", {
+            "module_id":   self._module_name,
+            "module_type": self._module_type,
+            "name":        self._display_name,
+            "port":        self._port,
+            "base_url":    f"http://{self._host}:{self._port}",
+            "tools": {"rpc": {"module": {
+                "health": {"description": "健康检查"},
+                "status": {"description": "状态查询（含 hook 安装状态）"},
+            }}},
+            "events_publish":   {self._module_name: {"test": {"description": "测试事件"}}},
+            "events_subscribe": ["module.started", "module.stopped", "module.shutdown"],
+        })
+        await self._publish_event(ws, "module.ready", {
+            "module_id": self._module_name,
+            "graceful_shutdown": True,
+            "startup_time": time.monotonic() - self._boot_t0,
+            "reason": "startup",
+        })
+        print(f"[evolmem] Ready (http://{self._host}:{self._port})")
+
+    # ── 优雅关闭 ──
+
+    async def _handle_shutdown(self, ws):
+        print("[evolmem] ========== SHUTDOWN STARTED ==========")
+        self._shutting_down = True
+
+        # Step 1: shutdown.ack
+        try:
+            await self._publish_event(ws, "module.shutdown.ack", {
+                "module_id": self._module_name,
+                "estimated_cleanup": 1,
+            })
+        except Exception as e:
+            print(f"[evolmem] Failed to send shutdown.ack: {e}")
+
+        # Step 2: module.exiting
+        try:
+            await self._publish_event(ws, "module.exiting", {
+                "module_id": self._module_name,
+                "type": "passive",
+                "reason": "shutdown_requested",
+                "restart": "auto",
+                "action": "none",
+                "timeout": 3.0,
+                "restart_delay": 0.0,
+            })
+        except Exception as e:
+            print(f"[evolmem] Failed to send module.exiting: {e}")
+
+        # Step 3: 清理资源
+        if self._test_task:
+            self._test_task.cancel()
+
+        # Step 4: shutdown.ready
+        try:
+            await self._publish_event(ws, "module.shutdown.ready", {
+                "module_id": self._module_name,
+            })
+        except Exception as e:
+            print(f"[evolmem] Failed to send shutdown.ready: {e}")
+
+        # Step 5: 等待 Kernel 处理
+        await asyncio.sleep(0.01)
+
+        # Step 6: 关闭额外连接
+        for _s, extra_ws in list(self._extra_ws.items()):
+            try:
+                await extra_ws.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 7: 清理 pending RPC + 关闭主连接
+        for fut in self._pending_rpc.values():
+            if not fut.done():
+                fut.cancel()
+        self._pending_rpc.clear()
+
+        if ws:
+            try:
+                await ws.close(code=1000, reason="Graceful shutdown")
+            except Exception as e:
+                print(f"[evolmem] Failed to close Kernel WebSocket: {e}")
+
+        # Step 8: 触发 uvicorn 退出
+        if self._uvicorn_server:
+            self._uvicorn_server.should_exit = True
+
+        print("[evolmem] ========== SHUTDOWN COMPLETE ==========")
+
+    # ── 测试事件 ──
+
+    async def _test_event_loop(self):
+        while not self._shutting_down:
+            await asyncio.sleep(12)
+            if self._shutting_down:
+                break
+            if self._ws:
+                await self._publish_event(self._ws, f"{self._module_name}.test", {"timestamp": time.time()})
+
+
+# ════════════════════════════════════════════
+#  Hook 安装 / 卸载 / 状态（纯函数，线程安全）
+# ════════════════════════════════════════════
+
+def _read_settings() -> dict:
+    if _CLAUDE_SETTINGS.exists():
+        try:
+            return json.loads(_CLAUDE_SETTINGS.read_text(encoding="utf-8"))
+        except Exception:
+            pass
+    return {}
+
+
+def _write_settings(data: dict):
+    _CLAUDE_SETTINGS.parent.mkdir(parents=True, exist_ok=True)
+    _CLAUDE_SETTINGS.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
+
+
+def _build_hook_entry(event: str, port: int) -> dict:
+    """生成单个事件的 hooks 配置块。"""
+    curl_cmd = (
+        f'curl -s -X POST http://localhost:{port}/hooks/event/{event} '
+        f'-H "Content-Type: application/json" -d @-'
+    )
+    return {"matcher": "", "hooks": [{"type": "command", "command": curl_cmd}]}
+
+
+def _install_hooks(port: int) -> dict:
+    settings = _read_settings()
+    hooks = settings.setdefault("hooks", {})
+    installed = []
+    for event in _HOOK_EVENTS:
+        hooks[event] = [_build_hook_entry(event, port)]
+        installed.append(event)
+    _write_settings(settings)
+    return {"success": True, "installed": installed, "settings_file": str(_CLAUDE_SETTINGS)}
+
+
+def _uninstall_hooks() -> dict:
+    settings = _read_settings()
+    hooks = settings.get("hooks", {})
+    removed = []
+    for event in _HOOK_EVENTS:
+        if event in hooks:
+            # 只删除 evolmem 注入的 curl 行，保留用户自定义 hook
+            filtered = [
+                entry for entry in hooks[event]
+                if not any(
+                    "/hooks/event/" in h.get("command", "")
+                    for h in entry.get("hooks", [])
+                )
+            ]
+            if filtered:
+                hooks[event] = filtered
+            else:
+                del hooks[event]
+            removed.append(event)
+    _write_settings(settings)
+    return {"success": True, "removed": removed, "settings_file": str(_CLAUDE_SETTINGS)}
+
+
+def _hooks_status() -> dict:
+    settings = _read_settings()
+    hooks = settings.get("hooks", {})
+    installed_events = []
+    for event in _HOOK_EVENTS:
+        entries = hooks.get(event, [])
+        if any(
+            "/hooks/event/" in h.get("command", "")
+            for entry in entries
+            for h in entry.get("hooks", [])
+        ):
+            installed_events.append(event)
+    return {
+        "installed": len(installed_events) == len(_HOOK_EVENTS),
+        "installed_events": installed_events,
+        "missing_events": [e for e in _HOOK_EVENTS if e not in installed_events],
+        "settings_file": str(_CLAUDE_SETTINGS),
+    }
+
+
+# ════════════════════════════════════════════
+#  项目列表辅助
+# ════════════════════════════════════════════
+
+def _list_projects(current: str | None) -> list:
+    projects = []
+    if not EVOLMEM_DATA_DIR.exists():
+        return projects
+    for project_dir in EVOLMEM_DATA_DIR.iterdir():
+        if not project_dir.is_dir() or project_dir.name in ("sessions", "summaries"):
+            continue
+        sessions_dir  = project_dir / "sessions"
+        summaries_dir = project_dir / "summaries"
+        session_count = len(list(sessions_dir.glob("session-*.jsonl"))) if sessions_dir.exists() else 0
+        topic_count   = len(list(summaries_dir.glob("topic-*.md")))     if summaries_dir.exists() else 0
+
+        message_count  = 0
+        total_duration = 0.0
+        if sessions_dir.exists():
+            for sf in sessions_dir.glob("session-*.jsonl"):
+                try:
+                    from datetime import datetime as _dt
+                    start_t = end_t = None
+                    with open(sf, "r", encoding="utf-8") as f:
+                        for line in f:
+                            try:
+                                rec = json.loads(line)
+                                rt  = rec.get("type", "")
+                                if rt in ("user_message", "assistant_message"):
+                                    message_count += 1
+                                elif rt == "session_start":
+                                    start_t = rec.get("timestamp")
+                                elif rt == "session_end":
+                                    end_t = rec.get("timestamp")
+                            except Exception:
+                                pass
+                    if start_t and end_t:
+                        s = _dt.strptime(start_t, "%Y-%m-%d %H:%M:%S")
+                        e = _dt.strptime(end_t,   "%Y-%m-%d %H:%M:%S")
+                        total_duration += (e - s).total_seconds()
+                except Exception:
+                    pass
+
+        h = int(total_duration // 3600)
+        m = int((total_duration % 3600) // 60)
+        is_current = bool(current and project_dir.name == Path(current).resolve().name)
+
+        projects.append({
+            "name":          project_dir.name,
+            "path":          str(project_dir),
+            "session_count": session_count,
+            "topic_count":   topic_count,
+            "message_count": message_count,
+            "work_duration": f"{h}h {m}m" if h else f"{m}m",
+            "is_current":    is_current,
+        })
+    return projects