ai-agent/backend/db.py

"""
SQLite store for the ai-agent backend.

Holds two kinds of state so we don't recompute them on every request:

* ``files``  — one row per exposed context file: content hash, size, mtime, the
  *real* Claude token count and its dollar cost. Token counts are expensive
  (a network round-trip to the count_tokens endpoint), so they are computed
  lazily and only when a file's content hash actually changes.
* ``skills`` — aggregated skill-usage analytics mined from the Claude Code
  transcripts (how many times each skill ran, and when it was last used).
* ``transcripts`` — per-transcript-file bookkeeping so skill stats can be
  recomputed incrementally (only re-parse the .jsonl files that changed).

A single connection (``check_same_thread=False``) guarded by a lock is plenty
for this low-traffic, single-process service.
"""

import json
import os
import sqlite3
import threading
import time

_SCHEMA = """
CREATE TABLE IF NOT EXISTS files (
    path         TEXT PRIMARY KEY,
    content_hash TEXT,
    size         INTEGER,
    mtime        REAL,
    is_markdown  INTEGER NOT NULL DEFAULT 0,
    tokens       INTEGER,            -- real (or estimated) token count
    cost         REAL,              -- tokens priced at the input rate
    estimated    INTEGER NOT NULL DEFAULT 1,   -- 1 = heuristic, 0 = real API count
    counted_at   REAL,              -- epoch secs the token count was last set
    seen_at      REAL,              -- epoch secs the file was last seen on disk
    words        INTEGER,           -- text stats, stored on content change so
    chars        INTEGER            -- /api/bundle never re-reads files per request
);

CREATE TABLE IF NOT EXISTS skills (
    name      TEXT PRIMARY KEY,
    count     INTEGER NOT NULL DEFAULT 0,
    last_used TEXT
);

CREATE TABLE IF NOT EXISTS transcripts (
    path         TEXT PRIMARY KEY,
    mtime        REAL,
    size         INTEGER,
    skills_json  TEXT,              -- {skill: {"count": n, "last": iso}}
    summary_json TEXT               -- per-conversation rollup for the list view
);

CREATE TABLE IF NOT EXISTS meta (
    key   TEXT PRIMARY KEY,
    value TEXT
);

CREATE TABLE IF NOT EXISTS conversation_diffs (
    session_id   TEXT PRIMARY KEY,
    fingerprint  TEXT,              -- endedAt + repo HEADs; recompute when it moves
    data_json    TEXT,              -- computed conversation-commits payload
    computed_at  REAL
);
"""


class Store:
    def __init__(self, path: str):
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        self._lock = threading.Lock()
        self._db = sqlite3.connect(path, check_same_thread=False)
        self._db.row_factory = sqlite3.Row
        # WAL keeps readers unblocked during the indexer's writes and cuts the
        # fsync-per-commit cost of its frequent small updates.
        self._db.execute("PRAGMA journal_mode=WAL")
        self._db.execute("PRAGMA synchronous=NORMAL")
        self._db.executescript(_SCHEMA)
        self._migrate()
        self._db.commit()
        # Decoded-summary cache: json.loads over every transcript's summary is
        # what every list/analytics request used to pay (~900 blobs, several MB).
        # `summaries_version` bumps on any change so request-level memoizers
        # (e.g. project cost rollups) can key off it.
        self._summaries: dict[str, dict] | None = None
        self.summaries_version = 0

    def _migrate(self) -> None:
        """Add columns introduced after a DB was first created."""
        cols = {r["name"] for r in self._db.execute("PRAGMA table_info(transcripts)")}
        if "summary_json" not in cols:
            self._db.execute("ALTER TABLE transcripts ADD COLUMN summary_json TEXT")
        fcols = {r["name"] for r in self._db.execute("PRAGMA table_info(files)")}
        for col in ("words", "chars"):
            if col not in fcols:
                self._db.execute(f"ALTER TABLE files ADD COLUMN {col} INTEGER")

    # ── files ────────────────────────────────────────────────────────────
    def get_file(self, path: str) -> sqlite3.Row | None:
        with self._lock:
            cur = self._db.execute("SELECT * FROM files WHERE path = ?", (path,))
            return cur.fetchone()

    def all_files(self) -> dict[str, sqlite3.Row]:
        with self._lock:
            cur = self._db.execute("SELECT * FROM files")
            return {r["path"]: r for r in cur.fetchall()}

    def upsert_file_meta(self, path: str, content_hash: str, size: int,
                         mtime: float, is_markdown: bool,
                         words: int | None = None,
                         chars: int | None = None) -> None:
        """Record a file's on-disk metadata. If the hash changed, the token
        count is invalidated (estimated flag flipped on) so the worker recounts."""
        now = time.time()
        with self._lock:
            row = self._db.execute(
                "SELECT content_hash FROM files WHERE path = ?", (path,)
            ).fetchone()
            if row is None:
                self._db.execute(
                    "INSERT INTO files (path, content_hash, size, mtime, "
                    "is_markdown, estimated, seen_at, words, chars) "
                    "VALUES (?,?,?,?,?,1,?,?,?)",
                    (path, content_hash, size, mtime, int(is_markdown), now,
                     words, chars),
                )
            elif row["content_hash"] != content_hash:
                # content changed → mark token count stale (estimated) to recount
                self._db.execute(
                    "UPDATE files SET content_hash=?, size=?, mtime=?, "
                    "is_markdown=?, estimated=1, seen_at=?, words=?, chars=? "
                    "WHERE path=?",
                    (content_hash, size, mtime, int(is_markdown), now,
                     words, chars, path),
                )
            else:
                self._db.execute(
                    "UPDATE files SET size=?, mtime=?, is_markdown=?, seen_at=?, "
                    "words=COALESCE(?, words), chars=COALESCE(?, chars) "
                    "WHERE path=?",
                    (size, mtime, int(is_markdown), now, words, chars, path),
                )
            self._db.commit()

    def set_tokens(self, path: str, tokens: int, cost: float, estimated: bool) -> None:
        with self._lock:
            self._db.execute(
                "UPDATE files SET tokens=?, cost=?, estimated=?, counted_at=? "
                "WHERE path=?",
                (tokens, cost, int(estimated), time.time(), path),
            )
            self._db.commit()

    def prune_files(self, keep_paths: set[str]) -> None:
        with self._lock:
            existing = {r["path"] for r in self._db.execute("SELECT path FROM files")}
            for gone in existing - keep_paths:
                self._db.execute("DELETE FROM files WHERE path=?", (gone,))
            self._db.commit()

    # ── transcripts / skills ─────────────────────────────────────────────
    def get_transcript(self, path: str) -> sqlite3.Row | None:
        with self._lock:
            return self._db.execute(
                "SELECT * FROM transcripts WHERE path=?", (path,)
            ).fetchone()

    def all_transcripts(self) -> dict[str, sqlite3.Row]:
        """Bookkeeping rows only (no JSON blobs) — this is the change-detection
        index, and pulling every summary_json with it would read the whole DB."""
        with self._lock:
            return {r["path"]: r for r in self._db.execute(
                "SELECT path, mtime, size FROM transcripts")}

    def all_transcript_skills(self) -> list[str]:
        """Every transcript's raw ``skills_json`` (for the skills rollup)."""
        with self._lock:
            return [r["skills_json"] for r in self._db.execute(
                "SELECT skills_json FROM transcripts")]

    def upsert_transcript(self, path: str, mtime: float, size: int,
                          skills: dict, summary: dict | None = None) -> None:
        with self._lock:
            self._db.execute(
                "INSERT INTO transcripts (path, mtime, size, skills_json, summary_json) "
                "VALUES (?,?,?,?,?) ON CONFLICT(path) DO UPDATE SET "
                "mtime=excluded.mtime, size=excluded.size, "
                "skills_json=excluded.skills_json, summary_json=excluded.summary_json",
                (path, mtime, size, json.dumps(skills),
                 json.dumps(summary) if summary is not None else None),
            )
            self._db.commit()
            if self._summaries is not None:
                if summary is not None:
                    # round-trip through JSON so the cache holds exactly what a
                    # cold load would (tuples→lists etc.), decoupled from the
                    # parser state's live dicts
                    self._summaries[path] = json.loads(json.dumps(summary))
                else:
                    self._summaries.pop(path, None)
            self.summaries_version += 1

    def _load_summaries(self) -> dict[str, dict]:
        """Decode-and-cache every stored summary (call with the lock held)."""
        if self._summaries is None:
            out: dict[str, dict] = {}
            rows = self._db.execute(
                "SELECT path, summary_json FROM transcripts "
                "WHERE summary_json IS NOT NULL").fetchall()
            for r in rows:
                try:
                    out[r["path"]] = json.loads(r["summary_json"])
                except (ValueError, TypeError):
                    continue
            self._summaries = out
        return self._summaries

    def all_summaries(self) -> list[tuple[str, dict]]:
        """(abs_path, summary dict) for every transcript that has one.

        Served from the decoded cache — callers must treat the dicts as
        **read-only** (copy before mutating)."""
        with self._lock:
            return list(self._load_summaries().items())

    def get_summary(self, path: str) -> dict | None:
        """The parsed summary dict for one transcript (by abs path), or None.

        Read-only — same shared cache as :meth:`all_summaries`."""
        with self._lock:
            return self._load_summaries().get(path)

    def prune_transcripts(self, keep_paths: set[str]) -> None:
        with self._lock:
            existing = {r["path"] for r in
                        self._db.execute("SELECT path FROM transcripts")}
            for gone in existing - keep_paths:
                self._db.execute("DELETE FROM transcripts WHERE path=?", (gone,))
                if self._summaries is not None:
                    self._summaries.pop(gone, None)
            self._db.commit()
            self.summaries_version += 1

    def replace_skills(self, skills: dict[str, dict]) -> None:
        """skills: {name: {"count": int, "last_used": iso|None}}"""
        with self._lock:
            self._db.execute("DELETE FROM skills")
            self._db.executemany(
                "INSERT INTO skills (name, count, last_used) VALUES (?,?,?)",
                [(n, d["count"], d.get("last_used")) for n, d in skills.items()],
            )
            self._db.commit()

    def all_skills(self) -> list[sqlite3.Row]:
        with self._lock:
            return list(self._db.execute(
                "SELECT * FROM skills ORDER BY count DESC, name ASC"))

    # ── conversation diff cache ──────────────────────────────────────────
    def get_conversation_diff(self, session_id: str) -> dict | None:
        """Cached {fingerprint, data} for a conversation, or None."""
        with self._lock:
            row = self._db.execute(
                "SELECT fingerprint, data_json FROM conversation_diffs "
                "WHERE session_id=?", (session_id,)).fetchone()
        if not row or not row["data_json"]:
            return None
        try:
            return {"fingerprint": row["fingerprint"],
                    "data": json.loads(row["data_json"])}
        except (ValueError, TypeError):
            return None

    def set_conversation_diff(self, session_id: str, fingerprint: str,
                              data: dict) -> None:
        with self._lock:
            self._db.execute(
                "INSERT INTO conversation_diffs "
                "(session_id, fingerprint, data_json, computed_at) VALUES (?,?,?,?) "
                "ON CONFLICT(session_id) DO UPDATE SET "
                "fingerprint=excluded.fingerprint, data_json=excluded.data_json, "
                "computed_at=excluded.computed_at",
                (session_id, fingerprint, json.dumps(data), time.time()),
            )
            self._db.commit()

    # ── meta ─────────────────────────────────────────────────────────────
    def get_meta(self, key: str) -> str | None:
        with self._lock:
            row = self._db.execute(
                "SELECT value FROM meta WHERE key=?", (key,)).fetchone()
            return row["value"] if row else None

    def set_meta(self, key: str, value: str) -> None:
        with self._lock:
            self._db.execute(
                "INSERT INTO meta (key, value) VALUES (?,?) "
                "ON CONFLICT(key) DO UPDATE SET value=excluded.value",
                (key, value),
            )
            self._db.commit()