SourceIndex 0.1.0__py3-none-any.whl

sourceindex/daemon/store.py

@@ -0,0 +1,426 @@
+"""Encrypted-at-rest store backing ``.sourceindex/``.
+
+Both file CONTENTS and PATHS are hidden on disk. Logical paths like
+``index.md`` or ``details/foo/bar.py.md`` are translated to opaque random
+chunk IDs via an encrypted manifest. Outside observers see only
+``manifest.enc`` and a churn of ``chunks/<random_id>.enc`` blobs; the
+filesystem layout no longer betrays the repo's file tree.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import secrets
+import tempfile
+import threading
+from pathlib import Path
+from typing import Iterator
+
+from . import crypto
+
+
+SENTINEL_FILE = ".daemon_key_id"
+ENC_SUFFIX = ".enc"
+TIER1_FILENAME = "index.md"
+MANIFEST_FILENAME = "manifest"
+CHUNKS_DIRNAME = "chunks"
+DETAILS_DIRNAME = "details"
+MANIFEST_VERSION = 1
+_CHUNK_ID_BYTES = 16
+
+
+class PlaintextStore:
+    """Direct read/write of a ``.sourceindex/`` directory without encryption.
+    Mirrors the EncryptedStore API so callers can be agnostic."""
+
+    def __init__(self, index_dir: Path) -> None:
+        self.index_dir = Path(index_dir)
+
+    def _resolve(self, rel: str) -> Path:
+        return self.index_dir / rel.lstrip("/")
+
+    def path_for(self, rel: str) -> Path:
+        return self._resolve(rel)
+
+    def ensure_initialized(self) -> None:
+        self.index_dir.mkdir(parents=True, exist_ok=True)
+
+    def read_bytes(self, rel: str) -> bytes:
+        return self._resolve(rel).read_bytes()
+
+    def read_text(self, rel: str) -> str:
+        return self._resolve(rel).read_text()
+
+    def write_bytes(self, rel: str, data: bytes) -> None:
+        target = self._resolve(rel)
+        target.parent.mkdir(parents=True, exist_ok=True)
+        target.write_bytes(data)
+
+    def write_text(self, rel: str, data: str) -> None:
+        target = self._resolve(rel)
+        target.parent.mkdir(parents=True, exist_ok=True)
+        target.write_text(data)
+
+    def exists(self, rel: str) -> bool:
+        return self._resolve(rel).is_file()
+
+    def unlink(self, rel: str, *, missing_ok: bool = True) -> None:
+        try:
+            self._resolve(rel).unlink()
+        except FileNotFoundError:
+            if not missing_ok:
+                raise
+
+    def iter_relative(self, subdir: str = "") -> "Iterator[str]":
+        root = self.index_dir / subdir if subdir else self.index_dir
+        if not root.exists():
+            return
+        for p in root.rglob("*"):
+            if not p.is_file():
+                continue
+            if p.parent == self.index_dir and p.name == SENTINEL_FILE:
+                continue
+            yield p.relative_to(self.index_dir).as_posix()
+
+    def plaintext_dump(self, out_dir: Path) -> tuple[int, int]:
+        out_dir = Path(out_dir)
+        out_dir.mkdir(parents=True, exist_ok=True)
+        files = 0
+        total = 0
+        for rel in self.iter_relative():
+            data = self.read_bytes(rel)
+            target = out_dir / rel
+            target.parent.mkdir(parents=True, exist_ok=True)
+            target.write_bytes(data)
+            files += 1
+            total += len(data)
+        return files, total
+
+
+class StoreError(RuntimeError):
+    pass
+
+
+class KeyMismatchError(StoreError):
+    pass
+
+
+class EncryptedStore:
+    def __init__(self, index_dir: Path, key: bytes, repo_hash: str) -> None:
+        self.index_dir = Path(index_dir)
+        self._key = key
+        self._repo_hash = repo_hash
+        self._lock = threading.Lock()
+        self._manifest: dict[str, str] | None = None
+
+    # ── path helpers ──────────────────────────────────────────────────────
+
+    def _manifest_path(self) -> Path:
+        return self.index_dir / (MANIFEST_FILENAME + ENC_SUFFIX)
+
+    def _chunks_dir(self) -> Path:
+        return self.index_dir / CHUNKS_DIRNAME
+
+    def _chunk_path(self, chunk_id: str) -> Path:
+        return self._chunks_dir() / (chunk_id + ENC_SUFFIX)
+
+    def _new_chunk_id(self) -> str:
+        return secrets.token_hex(_CHUNK_ID_BYTES)
+
+    def path_for(self, rel: str) -> Path:
+        """Logical path under ``index_dir`` for error-message use only;
+        the actual on-disk file is an opaque chunk under ``chunks/``."""
+        return self.index_dir / rel.lstrip("/")
+
+    # ── manifest plumbing ─────────────────────────────────────────────────
+
+    def _load_manifest_locked(self) -> dict[str, str]:
+        path = self._manifest_path()
+        if not path.is_file():
+            return {}
+        try:
+            blob = path.read_bytes()
+            plaintext = crypto.decrypt(blob, self._key)
+        except crypto.InvalidBlob as e:
+            raise KeyMismatchError(
+                f"Could not decrypt manifest at {path}: {e}. The on-disk "
+                "store and the loaded master key disagree — likely a stale "
+                "key file or a moved repository."
+            ) from e
+        try:
+            data = json.loads(plaintext.decode("utf-8"))
+        except (json.JSONDecodeError, UnicodeDecodeError) as e:
+            raise StoreError(f"Corrupt manifest at {path}: {e}") from e
+        if not isinstance(data, dict):
+            raise StoreError(f"Corrupt manifest at {path}: not a JSON object")
+        entries = data.get("entries", {})
+        if not isinstance(entries, dict):
+            raise StoreError(
+                f"Corrupt manifest at {path}: 'entries' is not a JSON object"
+            )
+        return {str(k): str(v) for k, v in entries.items()}
+
+    def _save_manifest_locked(self) -> None:
+        manifest = self._manifest if self._manifest is not None else {}
+        payload = {
+            "version": MANIFEST_VERSION,
+            "entries": dict(sorted(manifest.items())),
+        }
+        blob = crypto.encrypt(
+            json.dumps(payload, separators=(",", ":")).encode("utf-8"),
+            self._key,
+        )
+        path = self._manifest_path()
+        path.parent.mkdir(parents=True, exist_ok=True)
+        self._atomic_write_bytes(path, blob)
+
+    def _get_manifest_locked(self) -> dict[str, str]:
+        if self._manifest is None:
+            self._manifest = self._load_manifest_locked()
+        return self._manifest
+
+    @staticmethod
+    def _atomic_write_bytes(target: Path, blob: bytes) -> None:
+        fd, tmp = tempfile.mkstemp(prefix=".tmp-", dir=str(target.parent))
+        try:
+            with os.fdopen(fd, "wb") as f:
+                f.write(blob)
+            os.replace(tmp, target)
+        except Exception:
+            try:
+                os.unlink(tmp)
+            except OSError:
+                pass
+            raise
+
+    # ── lifecycle ─────────────────────────────────────────────────────────
+
+    def ensure_initialized(self) -> None:
+        self.index_dir.mkdir(parents=True, exist_ok=True)
+        sentinel = self.index_dir / SENTINEL_FILE
+        if sentinel.exists():
+            existing = sentinel.read_text().strip()
+            if existing and existing != self._repo_hash:
+                raise KeyMismatchError(
+                    f"Store at {self.index_dir} carries repo_hash={existing!r} "
+                    f"but this daemon is initialized for {self._repo_hash!r}. "
+                    "Either you moved the repo (re-init to mint a new key) or "
+                    "the keyring entry was deleted; nothing was modified."
+                )
+        else:
+            sentinel.write_text(self._repo_hash + "\n")
+
+    # ── core read / write API ─────────────────────────────────────────────
+
+    def read_bytes(self, rel: str) -> bytes:
+        rel = rel.lstrip("/")
+        with self._lock:
+            manifest = self._get_manifest_locked()
+            chunk_id = manifest.get(rel)
+            if chunk_id is None:
+                raise FileNotFoundError(
+                    f"{rel!r} is not present in the encrypted index"
+                )
+            chunk_path = self._chunk_path(chunk_id)
+            try:
+                blob = chunk_path.read_bytes()
+            except FileNotFoundError as e:
+                raise StoreError(
+                    f"Manifest references chunk {chunk_id!r} for {rel!r} but "
+                    f"the chunk file is missing at {chunk_path}. The store "
+                    "is corrupt."
+                ) from e
+        try:
+            return crypto.decrypt(blob, self._key)
+        except crypto.InvalidBlob as e:
+            raise KeyMismatchError(
+                f"Could not decrypt {rel!r}: {e}. The on-disk store and the "
+                "loaded master key disagree — likely a stale key file or a "
+                "moved repository."
+            ) from e
+
+    def read_text(self, rel: str) -> str:
+        return self.read_bytes(rel).decode("utf-8")
+
+    def write_bytes(self, rel: str, data: bytes) -> None:
+        rel = rel.lstrip("/")
+        chunks_dir = self._chunks_dir()
+        chunks_dir.mkdir(parents=True, exist_ok=True)
+        blob = crypto.encrypt(data, self._key)
+
+        with self._lock:
+            new_id = self._new_chunk_id()
+            new_chunk = self._chunk_path(new_id)
+            self._atomic_write_bytes(new_chunk, blob)
+
+            manifest = self._get_manifest_locked()
+            old_id = manifest.get(rel)
+            manifest[rel] = new_id
+            self._save_manifest_locked()
+
+            if old_id and old_id != new_id:
+                try:
+                    self._chunk_path(old_id).unlink()
+                except FileNotFoundError:
+                    pass
+
+    def write_text(self, rel: str, data: str) -> None:
+        self.write_bytes(rel, data.encode("utf-8"))
+
+    def exists(self, rel: str) -> bool:
+        rel = rel.lstrip("/")
+        with self._lock:
+            return rel in self._get_manifest_locked()
+
+    def unlink(self, rel: str, *, missing_ok: bool = True) -> None:
+        rel = rel.lstrip("/")
+        with self._lock:
+            manifest = self._get_manifest_locked()
+            chunk_id = manifest.pop(rel, None)
+            if chunk_id is None:
+                if not missing_ok:
+                    raise FileNotFoundError(rel)
+                return
+            self._save_manifest_locked()
+            try:
+                self._chunk_path(chunk_id).unlink()
+            except FileNotFoundError:
+                pass
+
+    def iter_relative(self, subdir: str = "") -> Iterator[str]:
+        with self._lock:
+            keys = sorted(self._get_manifest_locked().keys())
+        if not subdir:
+            yield from keys
+            return
+        prefix = subdir.rstrip("/") + "/"
+        bare = subdir.rstrip("/")
+        for k in keys:
+            if k == bare or k.startswith(prefix):
+                yield k
+
+    def plaintext_dump(self, out_dir: Path) -> tuple[int, int]:
+        out_dir = Path(out_dir)
+        out_dir.mkdir(parents=True, exist_ok=True)
+        files = 0
+        total = 0
+        for rel in self.iter_relative():
+            data = self.read_bytes(rel)
+            target = out_dir / rel
+            target.parent.mkdir(parents=True, exist_ok=True)
+            target.write_bytes(data)
+            files += 1
+            total += len(data)
+        return files, total
+
+    # ── migration: plaintext → chunked ────────────────────────────────────
+
+    def looks_like_plaintext_layout(self) -> bool:
+        """True iff plaintext-named files exist that need migrating into the
+        chunked layout. Catches partial-migration state too (e.g. crash
+        between writing the manifest and unlinking the plaintext)."""
+        if not self.index_dir.exists():
+            return False
+        return bool(self._collect_plaintext_targets())
+
+    def migrate_plaintext_in_place(self) -> int:
+        migrated = 0
+        for path in self._collect_plaintext_targets():
+            rel = path.relative_to(self.index_dir).as_posix()
+            try:
+                data = path.read_bytes()
+            except OSError:
+                continue
+            self.write_bytes(rel, data)
+            try:
+                path.unlink()
+            except OSError:
+                pass
+            migrated += 1
+        self._prune_empty_details_dir()
+        # Sentinel last so a torn migration is still detectable on rerun.
+        self.ensure_initialized()
+        return migrated
+
+    def _collect_plaintext_targets(self) -> list[Path]:
+        owned_top = {
+            TIER1_FILENAME,
+            "_summary_cache.json",
+            "_file_index.json",
+            "build_meta.json",
+            ".env",
+        }
+        out: list[Path] = []
+        for name in owned_top:
+            p = self.index_dir / name
+            if p.is_file() and not p.name.endswith(ENC_SUFFIX):
+                out.append(p)
+        details = self.index_dir / DETAILS_DIRNAME
+        if details.is_dir():
+            for p in details.rglob("*.md"):
+                if p.is_file() and not p.name.endswith(ENC_SUFFIX):
+                    out.append(p)
+        return out
+
+    # ── migration: legacy per-file encrypted → chunked ────────────────────
+
+    def looks_like_old_encrypted_layout(self) -> bool:
+        """True iff legacy per-file ``<rel>.enc`` blobs are present and need
+        migrating. Independent of whether ``manifest.enc`` already exists,
+        so a crashed partial migration resumes on the next start."""
+        if not self.index_dir.exists():
+            return False
+        return any(True for _ in self._iter_legacy_encrypted_files())
+
+    def migrate_old_encrypted_in_place(self) -> int:
+        migrated = 0
+        for path in list(self._iter_legacy_encrypted_files()):
+            rel = path.relative_to(self.index_dir).as_posix()
+            if not rel.endswith(ENC_SUFFIX):
+                continue
+            rel = rel[: -len(ENC_SUFFIX)]
+            try:
+                blob = path.read_bytes()
+                plaintext = crypto.decrypt(blob, self._key)
+            except (OSError, crypto.InvalidBlob):
+                continue
+            self.write_bytes(rel, plaintext)
+            try:
+                path.unlink()
+            except OSError:
+                pass
+            migrated += 1
+        self._prune_empty_details_dir()
+        self.ensure_initialized()
+        return migrated
+
+    def _iter_legacy_encrypted_files(self) -> Iterator[Path]:
+        manifest_path = self._manifest_path()
+        chunks_dir = self._chunks_dir()
+        for p in self.index_dir.rglob("*" + ENC_SUFFIX):
+            if not p.is_file():
+                continue
+            if p == manifest_path:
+                continue
+            try:
+                p.relative_to(chunks_dir)
+            except ValueError:
+                yield p
+
+    def _prune_empty_details_dir(self) -> None:
+        details = self.index_dir / DETAILS_DIRNAME
+        if not details.is_dir():
+            return
+        for sub in sorted(
+            details.rglob("*"), key=lambda p: len(p.parts), reverse=True
+        ):
+            if sub.is_dir():
+                try:
+                    sub.rmdir()
+                except OSError:
+                    pass
+        try:
+            details.rmdir()
+        except OSError:
+            pass

sourceindex/lib/backend.py

@@ -0,0 +1,240 @@
+"""HTTP client for the sourceindex hosted backend.
+
+Used opt-in: only fires when ``SOURCEINDEX_BACKEND_URL`` is set. Three helpers:
+
+- ``fetch_config()`` — pull the central preamble + model selection assigned
+  to this user's cohort. Cached on disk with TTL from the response.
+- ``apply_remote_config()`` — resolve the LLM model + (optionally) re-pin the
+  backend URL from the fetched config. Called once per CLI command.
+- ``upload_telemetry()`` — send anonymized session metrics. Caller is
+  responsible for anonymization (no file paths, no source).
+
+All three are best-effort. On any HTTP, network, or parse error they return
+``None``/``False``/the caller-supplied default and the caller falls back to
+local defaults. We do NOT hard-fail sourceindex when the backend is unreachable.
+
+Stdlib only — no new dependency on requests/httpx so the search subagent and
+post-commit hook keep their lean import footprint.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from pathlib import Path
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+
+from .log import get_logger
+
+_log = get_logger(__name__)
+
+_TIMEOUT_S = 5
+_CONFIG_CACHE_FILENAME = ".config_cache.json"
+
+# Hardcoded default — the operator's hosted deployment. Users who only paste
+# a sourceindex bearer key during ``sourceindex init`` get routed here without
+# ever seeing or configuring the URL. Override at runtime by exporting
+# ``SOURCEINDEX_BACKEND_URL=<other-url>`` (e.g. http://localhost:8787 for
+# local ``wrangler dev``), or ``SOURCEINDEX_BACKEND_URL=`` (empty string) to
+# disable the backend entirely and talk to providers directly.
+DEFAULT_BACKEND_URL = "https://sourceindex-backend.sourceindex.workers.dev"
+
+# Direct-provider key prefixes. The hosted backend expects an ``sk-si-...``
+# bearer and would reject these with a misleading auth error; when the user
+# is holding one of these (e.g. an OpenRouter key for eval), ``backend_base``
+# auto-opts out so the call hits the provider directly. Explicit
+# ``SOURCEINDEX_BACKEND_URL`` (any value, including a custom URL) still wins.
+_PROVIDER_KEY_PREFIXES = ("sk-or-", "sk-ant-", "gsk_", "AIza")
+
+
+def looks_like_provider_key(key: str) -> bool:
+    return bool(key) and any(key.startswith(p) for p in _PROVIDER_KEY_PREFIXES)
+
+# Process-scoped override set by ``apply_remote_config()`` when the backend
+# response includes a ``url`` field. Higher priority than the env var; we
+# avoid mutating ``os.environ`` because that would leak into Claude subagents
+# and the post-commit hook (which we explicitly do NOT want to redirect).
+_pinned_backend_url: str | None = None
+
+
+def _normalize_backend_url(url: str) -> str:
+    """Strip trailing slashes and idempotently append ``/v1``. Accepts a bare
+    host (``https://x.workers.dev``) or a fully qualified base (``…/v1``)."""
+    url = url.rstrip("/")
+    if not url.endswith("/v1"):
+        url = f"{url}/v1"
+    return url
+
+
+def _api_key() -> str | None:
+    # Lazy import: backend.py is imported on the backend code path from
+    # llm.py; resolve_api_key here would create a cycle at module load.
+    from .llm import resolve_api_key
+    return resolve_api_key(None)
+
+
+def backend_base(api_key: str | None = None) -> str | None:
+    """Return ``<URL>/v1`` for the hosted backend, or ``None`` if disabled.
+
+    Resolution order:
+    1. ``SOURCEINDEX_BACKEND_URL=""`` (explicit empty) → ``None`` (opt-out).
+       Checked first so a local opt-out beats admin pins — otherwise admins
+       could re-enable the backend on a user who explicitly disabled it.
+    2. ``_pinned_backend_url`` (set by ``apply_remote_config()``) → use it.
+    3. ``SOURCEINDEX_BACKEND_URL=<other>`` → use that.
+    4. Default to ``DEFAULT_BACKEND_URL`` — unless ``api_key`` (explicit arg,
+       falling back to env-resolved) looks like a direct-provider key, in
+       which case auto-opt-out. The hosted backend would reject such a key;
+       auto-routing direct keeps "paste your OpenRouter key" zero-config.
+    """
+    env = os.environ.get("SOURCEINDEX_BACKEND_URL")
+    if env == "":
+        return None
+    if _pinned_backend_url is not None:
+        return _pinned_backend_url
+    if env is not None:
+        return _normalize_backend_url(env)
+    effective_key = api_key if api_key is not None else _api_key()
+    if effective_key and looks_like_provider_key(effective_key):
+        return None
+    return _normalize_backend_url(DEFAULT_BACKEND_URL)
+
+
+def get_json(path: str, *, event: str) -> dict | None:
+    """Authenticated GET ``<backend>/v1/<path>`` returning the parsed JSON.
+    Returns ``None`` if the backend is disabled, no API key is available,
+    or the request fails — callers fall back to local defaults. ``event``
+    tags the debug log on failure for grep-ability."""
+    base = backend_base()
+    if not base:
+        return None
+    key = _api_key()
+    if not key:
+        return None
+    req = Request(f"{base}/{path}", headers={"authorization": f"Bearer {key}"})
+    try:
+        with urlopen(req, timeout=_TIMEOUT_S) as resp:
+            return json.loads(resp.read())
+    except (URLError, OSError, TimeoutError, ValueError) as e:
+        _log.debug(
+            f"backend {path} fetch failed",
+            extra={"event": event, "error": str(e)},
+        )
+        return None
+
+
+def fetch_config(*, cache_dir: Path | None = None, force: bool = False) -> dict | None:
+    """GET ``/v1/config``. Returns the parsed dict, or ``None`` if no backend
+    is configured / the request fails / no API key is available.
+
+    Response shape (all fields optional, admin-controlled in the Worker
+    dashboard):
+
+    - ``model`` (str): LiteLLM model string to use for indexing + search.
+      Any provider LiteLLM understands works — ``bedrock/...``,
+      ``openrouter/...``, ``vertex_ai/...``, ``anthropic/...``, etc.
+    - ``url`` (str): alternate SourceIndex deployment URL. When set, future
+      config refreshes AND LLM ``api_base`` both redirect to this URL (they
+      stay coupled — the Worker handles provider routing internally).
+    - ``ttl_seconds`` (int): how long the client may reuse this response
+      before re-fetching. Default 60 if omitted.
+
+    When ``cache_dir`` is provided, the response is cached at
+    ``<cache_dir>/.config_cache.json`` and reused for ``ttl_seconds``.
+    The internal ``_fetched_at`` timestamp is added on cache write.
+    """
+    cache_path = cache_dir / _CONFIG_CACHE_FILENAME if cache_dir else None
+    if cache_path and not force:
+        try:
+            payload = json.loads(cache_path.read_text())
+            ttl = float(payload.get("ttl_seconds", 60))
+            if payload.get("_fetched_at", 0) + ttl > time.time():
+                return payload
+        except (OSError, ValueError):
+            pass  # missing or malformed cache → fall through and re-fetch
+
+    data = get_json("config", event="backend.fetch_config_failed")
+    if data is None:
+        return None
+    data["_fetched_at"] = time.time()
+    if cache_path:
+        try:
+            cache_path.parent.mkdir(parents=True, exist_ok=True)
+            cache_path.write_text(json.dumps(data))
+        except OSError:
+            pass
+    return data
+
+
+def apply_remote_config(*, index_dir: Path | None, default_model: str) -> str:
+    """Fetch the admin-controlled config and apply its overrides.
+
+    Returns the model string to use — backend ``model`` if set, else
+    ``default_model``. As a side effect, pins ``_pinned_backend_url`` when
+    the response carries a non-empty ``url`` so subsequent ``backend_base()``
+    calls (both for further config refreshes and for LiteLLM ``api_base``)
+    target the new deployment.
+
+    Always returns ``default_model`` unchanged if the backend is disabled,
+    no API key is available, the request fails, or the response is empty —
+    callers can use the return value unconditionally.
+    """
+    global _pinned_backend_url
+    config = fetch_config(cache_dir=index_dir)
+    if not config:
+        return default_model
+
+    url_override = config.get("url")
+    model_override = config.get("model")
+
+    if isinstance(url_override, str) and url_override.strip():
+        _pinned_backend_url = _normalize_backend_url(url_override.strip())
+
+    resolved_model = (
+        model_override
+        if isinstance(model_override, str) and model_override.strip()
+        else default_model
+    )
+
+    _log.debug(
+        "remote config applied",
+        extra={
+            "event": "backend.config_applied",
+            "model_override": resolved_model if resolved_model != default_model else None,
+            "url_override": _pinned_backend_url,
+        },
+    )
+    return resolved_model
+
+
+def upload_telemetry(payload: dict) -> bool:
+    """POST ``/v1/telemetry``. Returns True on 2xx, False otherwise.
+
+    Caller-supplied ``payload`` must already be anonymized — no file paths,
+    no raw source, no tool args. See plan §Anonymization.
+    """
+    base = backend_base()
+    key = _api_key()
+    if not base or not key:
+        return False
+    body = json.dumps(payload).encode()
+    req = Request(
+        f"{base}/telemetry",
+        data=body,
+        headers={
+            "authorization": f"Bearer {key}",
+            "content-type": "application/json",
+        },
+        method="POST",
+    )
+    try:
+        with urlopen(req, timeout=_TIMEOUT_S) as resp:
+            return 200 <= resp.status < 300
+    except (URLError, OSError, TimeoutError) as e:
+        _log.debug(
+            "telemetry upload failed",
+            extra={"event": "backend.telemetry_failed", "error": str(e)},
+        )
+        return False