SourceIndex 0.1.0__py3-none-any.whl

sourceindex/daemon/server.py

@@ -0,0 +1,389 @@
+"""Daemon process: UDS server, RPC dispatch, in-process store ownership."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import signal
+import time
+from pathlib import Path
+from typing import Awaitable, Callable
+
+from . import lifecycle
+from .keyring_store import delete_key, load_or_create_key
+from .protocol import (
+    ADMIN_METHODS,
+    USER_METHODS,
+    E_AUTH_REQUIRED,
+    E_INVALID_PARAMS,
+    E_METHOD_NOT_FOUND,
+    E_NO_INDEX,
+    RpcError,
+    encode_frame,
+    make_error,
+    make_response,
+    read_frame_async,
+)
+from .store import EncryptedStore
+
+
+_log = logging.getLogger(__name__)
+
+
+IDLE_TIMEOUT_S = 600.0
+
+
+class Daemon:
+    watchdog_interval_s = 30.0
+
+    def __init__(self, repo_root: Path, *, allow_dump: bool = False) -> None:
+        self.repo_root = Path(repo_root).resolve()
+        self.allow_dump = allow_dump
+        self.repo_hash = lifecycle.repo_hash_for(self._main_worktree())
+        self.key = load_or_create_key(self.repo_hash)
+        from . import IndexDir
+        self.store = EncryptedStore(
+            IndexDir.for_repo(self.repo_root).path, self.key, self.repo_hash,
+        )
+        self._last_activity = time.monotonic()
+        self._shutdown_evt: asyncio.Event | None = None
+        self._repo_was_deleted = False
+        self._user_handlers: dict[str, Callable[[dict], Awaitable[dict]]] = {}
+        self._admin_handlers: dict[str, Callable[[dict], Awaitable[dict]]] = {}
+        self._register_user_handlers()
+        if allow_dump:
+            self._register_admin_handlers()
+
+    def _main_worktree(self) -> Path:
+        from ..lib.git import resolve_main_worktree
+        return resolve_main_worktree(self.repo_root)
+
+    def run(self) -> None:
+        if self.store.looks_like_plaintext_layout():
+            n = self.store.migrate_plaintext_in_place()
+            _log.warning(
+                "Migrated %d plaintext file(s) into the opaque-chunk store at %s",
+                n, self.store.index_dir,
+                extra={"event": "store.migrated", "files": n},
+            )
+        elif self.store.looks_like_old_encrypted_layout():
+            n = self.store.migrate_old_encrypted_in_place()
+            _log.warning(
+                "Migrated %d legacy encrypted file(s) into the opaque-chunk store at %s",
+                n, self.store.index_dir,
+                extra={"event": "store.chunked_migration", "files": n},
+            )
+        else:
+            try:
+                self.store.ensure_initialized()
+            except Exception as e:
+                _log.error("store init failed: %s", e)
+                raise
+
+        # Load encrypted .env into our env so RPC handlers inherit
+        # SOURCEINDEX_API_KEY without the client having to re-supply it
+        # (post-commit hook can't decrypt .env itself).
+        self._load_env_from_store()
+
+        lifecycle.write_pidfile(self.repo_hash)
+        try:
+            asyncio.run(self._serve())
+        finally:
+            if self._repo_was_deleted:
+                delete_key(self.repo_hash)
+            # Drop our own pidfile before calling cleanup_stale_files so
+            # its liveness guard (which protects a different live daemon's
+            # files from being yanked) doesn't bail on our own PID.
+            lifecycle.pidfile(self.repo_hash).unlink(missing_ok=True)
+            lifecycle.cleanup_stale_files(self.repo_hash)
+
+    def _load_env_from_store(self) -> None:
+        if not self.store.exists(".env"):
+            return
+        try:
+            text = self.store.read_text(".env")
+        except Exception as e:
+            _log.warning("could not load encrypted .env: %s", e)
+            return
+        from ..lib.env import parse_env_text
+        for k, v in parse_env_text(text).items():
+            if k and k not in os.environ:
+                os.environ[k] = v
+
+    async def _serve(self) -> None:
+        self._shutdown_evt = asyncio.Event()
+        loop = asyncio.get_running_loop()
+        for sig in (signal.SIGTERM, signal.SIGINT):
+            try:
+                loop.add_signal_handler(sig, self._request_shutdown)
+            except (NotImplementedError, RuntimeError, ValueError):
+                # Off-main-thread (tests) — stop_daemon's SIGKILL fallback
+                # still cleans up if shutdown never arrives.
+                pass
+
+        user_path = lifecycle.user_socket(self.repo_hash)
+        admin_path = lifecycle.admin_socket(self.repo_hash)
+        for p in (user_path, admin_path):
+            try:
+                p.unlink()
+            except FileNotFoundError:
+                pass
+
+        user_server = await asyncio.start_unix_server(self._handle_user, str(user_path))
+        os.chmod(user_path, 0o600)
+        servers = [user_server]
+        if self.allow_dump:
+            admin_server = await asyncio.start_unix_server(
+                self._handle_admin, str(admin_path)
+            )
+            os.chmod(admin_path, 0o600)
+            servers.append(admin_server)
+
+        watchdog_task = asyncio.create_task(self._watchdog())
+        try:
+            await self._shutdown_evt.wait()
+        finally:
+            watchdog_task.cancel()
+            for s in servers:
+                s.close()
+            for s in servers:
+                await s.wait_closed()
+
+    def _request_shutdown(self) -> None:
+        if self._shutdown_evt is not None:
+            self._shutdown_evt.set()
+
+    async def _watchdog(self) -> None:
+        # Single periodic loop guards two unrelated concerns:
+        #   1. The repo directory got deleted (rm -rf) — tear down and
+        #      drop the encryption key so it doesn't orphan in the keyring.
+        #   2. No client activity for IDLE_TIMEOUT_S — shut down to release
+        #      the socket so a stale daemon doesn't outlive its usefulness.
+        # Repo-gone wins if both fire on the same tick.
+        while True:
+            await asyncio.sleep(self.watchdog_interval_s)
+            if self._detect_repo_deleted():
+                _log.info(
+                    "Repo root no longer exists; deleting key and shutting down",
+                    extra={"event": "daemon.repo_deleted"},
+                )
+                self._repo_was_deleted = True
+                self._request_shutdown()
+                return
+            if time.monotonic() - self._last_activity > IDLE_TIMEOUT_S:
+                _log.info(
+                    "Idle timeout — shutting down",
+                    extra={"event": "daemon.idle_shutdown"},
+                )
+                self._request_shutdown()
+                return
+
+    def _detect_repo_deleted(self) -> bool:
+        # Require the parent to still exist so a temporarily unmounted
+        # drive or unreachable network mount doesn't trigger a wipe.
+        return (
+            not self.repo_root.exists()
+            and self.repo_root.parent.exists()
+        )
+
+    async def _handle_user(self, reader, writer) -> None:
+        await self._handle_conn(reader, writer, self._user_handlers, USER_METHODS)
+
+    async def _handle_admin(self, reader, writer) -> None:
+        await self._handle_conn(reader, writer, self._admin_handlers, ADMIN_METHODS)
+
+    async def _handle_conn(self, reader, writer, handlers, allowed) -> None:
+        try:
+            while True:
+                try:
+                    req = await read_frame_async(reader)
+                except asyncio.IncompleteReadError:
+                    return
+                self._last_activity = time.monotonic()
+                response = await self._dispatch(req, handlers, allowed)
+                writer.write(encode_frame(response))
+                await writer.drain()
+        except (ConnectionResetError, BrokenPipeError):
+            pass
+        finally:
+            try:
+                writer.close()
+                await writer.wait_closed()
+            except Exception:
+                pass
+
+    async def _dispatch(self, req: dict, handlers: dict, allowed) -> dict:
+        rid = req.get("id")
+        method = req.get("method")
+        params = req.get("params") or {}
+        if not isinstance(method, str) or method not in allowed:
+            return make_error(
+                rid, RpcError(E_METHOD_NOT_FOUND, f"method not found: {method!r}")
+            )
+        handler = handlers.get(method)
+        if handler is None:
+            # In allowed set but not registered (e.g. admin method on user socket).
+            return make_error(
+                rid, RpcError(E_METHOD_NOT_FOUND, f"method not found: {method!r}")
+            )
+        try:
+            result = await handler(params)
+            return make_response(rid, result)
+        except RpcError as e:
+            return make_error(rid, e)
+        except Exception as e:
+            _log.exception("handler error in %s", method)
+            return make_error(rid, e)
+
+    def _register_user_handlers(self) -> None:
+        self._user_handlers = {
+            "ping": self._h_ping,
+            "search": self._h_search,
+            "init": self._h_init,
+            "update": self._h_update,
+            "install_hooks": self._h_install_hooks,
+            "usage": self._h_usage,
+        }
+
+    async def _h_ping(self, params: dict) -> dict:
+        return {
+            "ok": True,
+            "repo_hash": self.repo_hash,
+            "allow_dump": self.allow_dump,
+            "idle_for_s": time.monotonic() - self._last_activity,
+        }
+
+    async def _h_search(self, params: dict) -> dict:
+        query = params.get("query")
+        model = params.get("model")
+        api_key = params.get("api_key") or os.environ.get("SOURCEINDEX_API_KEY")
+        if not isinstance(query, str) or not query:
+            raise RpcError(E_INVALID_PARAMS, "query is required (non-empty string)")
+        if not self.store.exists("index.md"):
+            raise RpcError(E_NO_INDEX, "no index found; run `sourceindex init` first")
+        from ..search import search_index_with_store
+        t0 = time.perf_counter()
+        roadmap = await asyncio.to_thread(
+            search_index_with_store,
+            query, self.repo_root, self.store, model, api_key,
+        )
+        return {
+            "roadmap": roadmap,
+            "latency_ms": int((time.perf_counter() - t0) * 1000),
+        }
+
+    def _progress_writer_for_this_rpc(self) -> "Callable[[], None]":
+        """Register a writer that mirrors each progress tick into
+        $XDG_RUNTIME_DIR/sourceindex/<hash>.progress for the CLI to tail.
+        Returns a cleanup callable to invoke in finally."""
+        from ..build import indexer
+        progress_path = lifecycle.progress_file(self.repo_hash)
+        progress_path.unlink(missing_ok=True)
+
+        def write_progress(line: str) -> None:
+            try:
+                progress_path.write_text(line + "\n")
+            except OSError:
+                pass
+
+        indexer.set_progress_writer(write_progress)
+
+        def cleanup() -> None:
+            indexer.set_progress_writer(None)
+            progress_path.unlink(missing_ok=True)
+        return cleanup
+
+    async def _h_init(self, params: dict) -> dict:
+        from ..build import build_index_with_store
+        api_key = params.get("api_key") or os.environ.get("SOURCEINDEX_API_KEY")
+        if not api_key:
+            raise RpcError(E_AUTH_REQUIRED, "api_key is required")
+        cleanup = self._progress_writer_for_this_rpc()
+        try:
+            return await asyncio.to_thread(
+                build_index_with_store,
+                self.repo_root, self.store,
+                model=params.get("model"),
+                api_key=api_key,
+                languages=params.get("languages"),
+                workers=params.get("workers"),
+            )
+        finally:
+            cleanup()
+
+    async def _h_update(self, params: dict) -> dict:
+        from ..build import update_index_with_store
+        api_key = params.get("api_key") or os.environ.get("SOURCEINDEX_API_KEY")
+        if not api_key:
+            raise RpcError(E_AUTH_REQUIRED, "api_key is required")
+        cleanup = self._progress_writer_for_this_rpc()
+        try:
+            return await asyncio.to_thread(
+                update_index_with_store,
+                self.repo_root, self.store,
+                model=params.get("model"),
+                api_key=api_key,
+                languages=params.get("languages"),
+                changed_files=params.get("changed_files"),
+                workers=params.get("workers"),
+            )
+        finally:
+            cleanup()
+
+    async def _h_install_hooks(self, params: dict) -> dict:
+        from ..cli.install import _install_hooks
+        await asyncio.to_thread(_install_hooks, self.repo_root)
+        return {"installed": True}
+
+    async def _h_usage(self, params: dict) -> dict:
+        from ..lib.cost import fetch_budget
+        budget = await asyncio.to_thread(fetch_budget)
+        return {"budget": budget}
+
+    def _register_admin_handlers(self) -> None:
+        self._admin_handlers = {
+            "dump_all": self._h_dump_all,
+            "read_index_md": self._h_read_index_md,
+            "read_detail": self._h_read_detail,
+            "read_env": self._h_read_env,
+            "shutdown": self._h_shutdown,
+        }
+
+    async def _h_dump_all(self, params: dict) -> dict:
+        out_dir = params.get("out_dir")
+        if not isinstance(out_dir, str) or not out_dir:
+            raise RpcError(E_INVALID_PARAMS, "out_dir is required")
+        files, total = await asyncio.to_thread(
+            self.store.plaintext_dump, Path(out_dir)
+        )
+        return {"files_written": files, "bytes": total}
+
+    async def _h_read_index_md(self, params: dict) -> dict:
+        return {"text": self.store.read_text("index.md")}
+
+    async def _h_read_detail(self, params: dict) -> dict:
+        rel = params.get("rel_path")
+        if not isinstance(rel, str) or not rel:
+            raise RpcError(E_INVALID_PARAMS, "rel_path is required")
+        return {"text": self.store.read_text(f"details/{rel}.md")}
+
+    async def _h_read_env(self, params: dict) -> dict:
+        if not self.store.exists(".env"):
+            return {"vars": {}}
+        from ..lib.env import parse_env_text
+        return {"vars": parse_env_text(self.store.read_text(".env"))}
+
+    async def _h_shutdown(self, params: dict) -> dict:
+        self._request_shutdown()
+        return {"ok": True}
+
+
+def run_daemon(repo_root: Path, *, allow_dump: bool = False) -> None:
+    if lifecycle.is_daemon_running(repo_root):
+        _log.info(
+            "Daemon already running for %s; exiting", repo_root,
+            extra={"event": "daemon.already_running"},
+        )
+        return
+    Daemon(repo_root, allow_dump=allow_dump).run()