LLeMbas-CLI/src/session/store.ts

import { Database } from "bun:sqlite"
import { mkdirSync } from "node:fs"
import { join } from "node:path"
import { paths } from "../config/paths.ts"
import type { Message, Usage } from "../provider/types.ts"
import { isNote } from "./prefix.ts"

// Sessions and their messages, in ~/.local/share/lembas/sessions.db. Message text is also
// indexed with FTS5 so past sessions can be searched (session_search).

export interface SessionRow {
  id: string
  created: number
  updated: number
  title: string
  root: string
  model: string
  /** A subagent's session: the session whose task started it. Null for a session of its own. */
  parent_id?: string | null
}

/** A background job's row (jobs/registry.ts persists one per job): the record as JSON, with the
 *  columns a list needs beside it. */
export interface StoredJob {
  id: string
  session_id: string
  kind: string
  state: string
  started: number
  ended: number | null
  json: string
}

const SCHEMA = `
CREATE TABLE IF NOT EXISTS sessions (
  id TEXT PRIMARY KEY, created INTEGER NOT NULL, updated INTEGER NOT NULL,
  title TEXT NOT NULL DEFAULT '', root TEXT NOT NULL, model TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS messages (
  id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
  role TEXT NOT NULL, json TEXT NOT NULL, created INTEGER NOT NULL,
  kind TEXT NOT NULL DEFAULT 'message'
);
CREATE INDEX IF NOT EXISTS messages_session_kind ON messages(session_id, kind, id);
CREATE INDEX IF NOT EXISTS messages_session ON messages(session_id, id);
CREATE TABLE IF NOT EXISTS jobs (
  id TEXT NOT NULL, session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
  kind TEXT NOT NULL, json TEXT NOT NULL, state TEXT NOT NULL, started INTEGER NOT NULL, ended INTEGER,
  PRIMARY KEY (session_id, id)
);
CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(text, content='', contentless_delete=1);
CREATE TABLE IF NOT EXISTS results (
  session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, seq INTEGER NOT NULL, id TEXT NOT NULL,
  json TEXT NOT NULL, text TEXT NOT NULL,
  PRIMARY KEY (session_id, id)
);
CREATE TABLE IF NOT EXISTS requests (
  id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
  created INTEGER NOT NULL, model TEXT NOT NULL, kind TEXT NOT NULL,
  input INTEGER NOT NULL, output INTEGER NOT NULL, reasoning INTEGER, cached INTEGER, estimated INTEGER NOT NULL DEFAULT 0
);
CREATE TABLE IF NOT EXISTS session_meta (
  session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, key TEXT NOT NULL, value TEXT NOT NULL,
  PRIMARY KEY (session_id, key)
);
`

/** The session_meta key saying a session has been in the web UI's hands (its rules are in
 *  acp/agent.ts, which re-exports it). Here, so the app can ask it too without importing the
 *  agent, which imports the app. */
export const REMOTE_META = "remote"

/** The session_meta key holding the directory a session was started in: a session started
 *  in a subdirectory of its project is reopened there, not at the project's root. */
export const CWD_META = "cwd"

export function plainText(m: Message): string {
  if (m.role === "tool") return m.content
  // A session note is the harness's, not the person's words: not searched.
  return m.parts.map((p) => (p.type === "text" ? (m.role === "user" && isNote(p.text) ? "" : p.text) : p.type === "reasoning" ? p.text : p.type === "tool_call" ? `${p.name} ${p.args}` : "")).join("\n")
}

/** One request's usage as the store keeps it (harness context.md §10 rule 6): what it was for — a
 *  step of a turn, or a compaction's summary — and the cached input tokens where the server said. */
export interface RequestRow extends Usage {
  model: string
  kind: "step" | "compaction"
  created?: number
}

export interface Row {
  id: number
  role: string
  json: string
  created: number
}
export interface Hit extends Row {
  session_id: string
}

export class Store {
  readonly db: Database

  constructor(file = join(paths.data, "sessions.db")) {
    if (file !== ":memory:") mkdirSync(paths.data, { recursive: true })
    this.db = new Database(file, { create: true })
    this.db.run("PRAGMA journal_mode = WAL")
    // Several writers share this file: every terminal, the service, and in the service one
    // connection per session it holds. Without a timeout a write that meets another's lock fails at
    // once with "database is locked" — a delete from the web UI among them, which the instance then
    // only saw as a refusal. Writes here are short; waiting for one is always right.
    this.db.run("PRAGMA busy_timeout = 5000")
    this.db.run("PRAGMA foreign_keys = ON")
    // Additive migrations: a column is added when missing, never renamed or dropped.
    this.db.exec(SCHEMA.split("CREATE INDEX IF NOT EXISTS messages_session_kind")[0]!)
    const cols = (this.db.query("PRAGMA table_info(messages)").all() as { name: string }[]).map((c) => c.name)
    if (!cols.includes("kind")) this.db.run("ALTER TABLE messages ADD COLUMN kind TEXT NOT NULL DEFAULT 'message'")
    // A subagent's session hangs off its parent's and goes with it (spec 2.0.0).
    const sessionCols = (this.db.query("PRAGMA table_info(sessions)").all() as { name: string }[]).map((c) => c.name)
    if (!sessionCols.includes("parent_id")) this.db.run("ALTER TABLE sessions ADD COLUMN parent_id TEXT REFERENCES sessions(id) ON DELETE CASCADE")
    this.db.exec(SCHEMA)
  }

  /** `parentId`: a subagent's session, which belongs to the one whose task started it — listed
   *  nowhere on its own, deleted with it. */
  createSession(root: string, model: string, parentId?: string): SessionRow {
    const now = Date.now()
    const id = `ses_${now.toString(36)}${Math.random().toString(36).slice(2, 8)}`
    this.db.query("INSERT INTO sessions (id, created, updated, root, model, parent_id) VALUES (?, ?, ?, ?, ?, ?)").run(id, now, now, root, model, parentId ?? null)
    return { id, created: now, updated: now, title: "", root, model, parent_id: parentId ?? null }
  }

  /** A subagent's sessions under this one. */
  children(sessionId: string): string[] {
    return (this.db.query("SELECT id FROM sessions WHERE parent_id = ?").all(sessionId) as { id: string }[]).map((r) => r.id)
  }

  saveJob(j: StoredJob) {
    this.db
      .query("INSERT INTO jobs (id, session_id, kind, json, state, started, ended) VALUES (?, ?, ?, ?, ?, ?, ?) ON CONFLICT (session_id, id) DO UPDATE SET json = excluded.json, state = excluded.state, started = excluded.started, ended = excluded.ended")
      .run(j.id, j.session_id, j.kind, j.json, j.state, j.started, j.ended)
  }

  /** A session's jobs as the store has them, oldest first. */
  jobs(sessionId: string): StoredJob[] {
    return this.db.query("SELECT id, session_id, kind, json, state, started, ended FROM jobs WHERE session_id = ? ORDER BY started, id").all(sessionId) as StoredJob[]
  }

  deleteJob(sessionId: string, id: string) {
    this.db.query("DELETE FROM jobs WHERE session_id = ? AND id = ?").run(sessionId, id)
  }

  append(sessionId: string, m: Message): number {
    const now = Date.now()
    const r = this.db.query("INSERT INTO messages (session_id, role, json, created) VALUES (?, ?, ?, ?)").run(sessionId, m.role, JSON.stringify(m), now)
    const text = plainText(m)
    if (text.trim()) this.db.query("INSERT INTO messages_fts (rowid, text) VALUES (?, ?)").run(r.lastInsertRowid, text)
    this.db.query("UPDATE sessions SET updated = ? WHERE id = ?").run(now, sessionId)
    return Number(r.lastInsertRowid)
  }

  /** Delete a session and everything in it: /sessions delete, /delete, or the web UI
   *  deleting the chat it is. The search index has no content of its own, so its rows go by id. */
  deleteSession(sessionId: string): boolean {
    const del = this.db.transaction(() => {
      // Its subagents' sessions go with it (ON DELETE CASCADE), so their index rows go here too.
      this.db.query("DELETE FROM messages_fts WHERE rowid IN (SELECT id FROM messages WHERE session_id = ? OR session_id IN (SELECT id FROM sessions WHERE parent_id = ?))").run(sessionId, sessionId)
      return this.db.query("DELETE FROM sessions WHERE id = ?").run(sessionId).changes > 0
    })
    // IMMEDIATE: the write lock taken up front, so busy_timeout applies. A deferred
    // transaction starts by reading (the subquery), and one that then meets another process's
    // write lock is refused at once — SQLite will not wait in a way that could deadlock — so a
    // delete from the web UI that met a terminal writing a turn failed as "database is locked".
    return del.immediate()
  }

  /** /undo and /redo: hide or bring back the rows of a turn. They are never deleted. */
  setKind(sessionId: string, fromId: number, toId: number, kind: "message" | "reverted") {
    const other = kind === "message" ? "reverted" : "message"
    this.db.query("UPDATE messages SET kind = ? WHERE session_id = ? AND id BETWEEN ? AND ? AND kind = ?").run(kind, sessionId, fromId, toId, other)
  }

  /** The row of the user message that started turn `turnId` (kept with the message, engine.ts),
   *  or undefined for a turn this session does not have. */
  turnRow(sessionId: string, turnId: string): number | undefined {
    // The LIKE is a prefilter; the message itself says.
    const rows = this.db.query("SELECT id, json FROM messages WHERE session_id = ? AND role = 'user' AND kind = 'message' AND json LIKE ? ORDER BY id").all(sessionId, `%${JSON.stringify(turnId)}%`) as { id: number; json: string }[]
    return rows.find((r) => (JSON.parse(r.json) as Message).role === "user" && (JSON.parse(r.json) as { turnId?: string }).turnId === turnId)?.id
  }

  /** The turns that start at row `fromId` or after it, oldest first: the ids the visible user
   *  messages carry. */
  turnIdsFrom(sessionId: string, fromId: number): string[] {
    const rows = this.db.query("SELECT json FROM messages WHERE session_id = ? AND role = 'user' AND kind = 'message' AND id >= ? ORDER BY id").all(sessionId, fromId) as { json: string }[]
    return rows.map((r) => (JSON.parse(r.json) as { turnId?: string }).turnId).filter((id): id is string => typeof id === "string")
  }

  /** Record a compaction: what the model sees from here on starts with this summary. */
  compaction(sessionId: string, summary: string) {
    this.db
      .query("INSERT INTO messages (session_id, role, json, created, kind) VALUES (?, 'user', ?, ?, 'compaction')")
      .run(sessionId, JSON.stringify({ summary }), Date.now())
  }

  /** The conversation as the model sees it: from the last compaction on, the summary first. */
  context(sessionId: string): Message[] {
    const last = this.db
      .query("SELECT id, json FROM messages WHERE session_id = ? AND kind = 'compaction' ORDER BY id DESC LIMIT 1")
      .get(sessionId) as { id: number; json: string } | null
    const after = last?.id ?? 0
    const rows = this.db
      .query("SELECT json FROM messages WHERE session_id = ? AND kind = 'message' AND id > ? ORDER BY id")
      .all(sessionId, after) as { json: string }[]
    const msgs = rows.map((r) => JSON.parse(r.json) as Message)
    return last ? [...summaryTurns(JSON.parse(last.json).summary as string), ...msgs] : msgs
  }

  /** A session's tool results, each in full (harness spec 3.0.0, context.md): what `recall` shows
   *  and what a stub is decided against. Kept as long as the session is. */
  results(sessionId: string): import("./results.ts").ResultStore {
    const q = {
      count: this.db.query("SELECT COUNT(*) AS n FROM results WHERE session_id = ?"),
      save: this.db.query("INSERT OR REPLACE INTO results (session_id, seq, id, json, text) VALUES (?, ?, ?, ?, ?)"),
      get: this.db.query("SELECT json, text FROM results WHERE session_id = ? AND id = ?"),
      all: this.db.query("SELECT json FROM results WHERE session_id = ? ORDER BY seq"),
    }
    type R = import("./results.ts").StoredResult
    return {
      count: () => (q.count.get(sessionId) as { n: number }).n,
      save: (r) => {
        const { text, ...rest } = r
        q.save.run(sessionId, (q.count.get(sessionId) as { n: number }).n + 1, r.id, JSON.stringify(rest), text)
      },
      get: (id) => {
        const row = q.get.get(sessionId, id) as { json: string; text: string } | null
        return row ? ({ ...JSON.parse(row.json), text: row.text } as R) : undefined
      },
      // Without their text: the table is decided by hashes and ranges.
      all: () => (q.all.all(sessionId) as { json: string }[]).map((r) => ({ ...JSON.parse(r.json), text: "" }) as R),
    }
  }

  /** Record one request's usage. */
  request(sessionId: string, r: RequestRow) {
    this.db
      .query("INSERT INTO requests (session_id, created, model, kind, input, output, reasoning, cached, estimated) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)")
      .run(sessionId, Date.now(), r.model, r.kind, r.input, r.output, r.reasoning ?? null, r.cached ?? null, r.estimated ? 1 : 0)
  }

  /** A session's requests, oldest first. */
  requests(sessionId: string): RequestRow[] {
    const rows = this.db.query("SELECT created, model, kind, input, output, reasoning, cached, estimated FROM requests WHERE session_id = ? ORDER BY id").all(sessionId) as (Omit<RequestRow, "estimated" | "reasoning" | "cached"> & { estimated: number; reasoning: number | null; cached: number | null })[]
    return rows.map(({ estimated, reasoning, cached, ...r }) => ({ ...r, ...(reasoning !== null ? { reasoning } : {}), ...(cached !== null ? { cached } : {}), ...(estimated ? { estimated: true } : {}) }))
  }

  /** Something a session keeps besides its messages, as JSON. */
  meta<T>(sessionId: string, key: string): T | undefined {
    const r = this.db.query("SELECT value FROM session_meta WHERE session_id = ? AND key = ?").get(sessionId, key) as { value: string } | null
    return r ? (JSON.parse(r.value) as T) : undefined
  }

  setMeta(sessionId: string, key: string, value: unknown) {
    this.db.query("INSERT INTO session_meta (session_id, key, value) VALUES (?, ?, ?) ON CONFLICT (session_id, key) DO UPDATE SET value = excluded.value").run(sessionId, key, JSON.stringify(value))
  }

  setTitle(sessionId: string, title: string) {
    this.db.query("UPDATE sessions SET title = ? WHERE id = ?").run(title, sessionId)
  }

  /** The text of a session's first user message (sessions from before titles have only this). */
  firstPrompt(sessionId: string): string {
    const r = this.db.query("SELECT json FROM messages WHERE session_id = ? AND role = 'user' AND kind = 'message' ORDER BY id LIMIT 1").get(sessionId) as { json: string } | null
    if (!r) return ""
    const m = JSON.parse(r.json) as Message
    return m.role === "user" ? m.parts.map((p) => (p.type === "text" ? p.text : "")).find((t) => t.trim()) ?? "" : ""
  }

  messages(sessionId: string): Message[] {
    return (this.db.query("SELECT json FROM messages WHERE session_id = ? AND kind = 'message' ORDER BY id").all(sessionId) as { json: string }[]).map(
      (r) => JSON.parse(r.json) as Message,
    )
  }

  /** Newest first. `withMessages`: only sessions somebody said something in (every start and /new
   *  makes one, so the empty ones would crowd out the rest). A subagent's session is never listed:
   *  it is reached through the job that ran it. */
  sessions(limit = 50, root?: string, withMessages = false): SessionRow[] {
    const where = ["s.parent_id IS NULL", root ? "s.root = ?" : "", withMessages ? "EXISTS (SELECT 1 FROM messages m WHERE m.session_id = s.id AND m.kind = 'message')" : ""].filter(Boolean)
    return this.db
      .query(`SELECT s.* FROM sessions s ${where.length ? `WHERE ${where.join(" AND ")}` : ""} ORDER BY s.updated DESC LIMIT ?`)
      .all(...(root ? [root] : []), limit) as SessionRow[]
  }

  /** FTS hits in visible user/assistant messages (not reverted, not tool output), best first.
   *  `root` keeps one project's sessions; `perSession` the best hit of each session. Both are done
   *  before the limit, so a busy other project cannot push this one's hits out. */
  hits(match: string, o: { exclude?: string[]; sort?: "newest" | "oldest"; limit?: number; roles?: string[]; root?: string; perSession?: boolean } = {}): Hit[] {
    const roles = o.roles ?? ["user", "assistant"]
    const ex = o.exclude ?? []
    const order = (t: string) => (o.sort === "newest" ? `${t}.created DESC, ${t}.id DESC` : o.sort === "oldest" ? `${t}.created ASC, ${t}.id ASC` : `${t}.r, ${t}.id`)
    const inner = `SELECT m.session_id, m.id, m.role, m.json, m.created, f.rank AS r FROM messages_fts f JOIN messages m ON m.id = f.rowid
         JOIN sessions s ON s.id = m.session_id AND s.parent_id IS NULL${o.root !== undefined ? " AND s.root = ?" : ""}
         WHERE messages_fts MATCH ? AND m.kind = 'message' AND m.role IN (${roles.map(() => "?").join(",")})
         ${ex.length ? `AND m.session_id NOT IN (${ex.map(() => "?").join(",")})` : ""}`
    const sql = o.perSession
      ? `SELECT session_id, id, role, json, created FROM (SELECT h.*, row_number() OVER (PARTITION BY h.session_id ORDER BY ${order("h")}) AS n FROM (${inner}) h) x WHERE x.n = 1 ORDER BY ${order("x")} LIMIT ?`
      : `SELECT session_id, id, role, json, created FROM (${inner}) h ORDER BY ${order("h")} LIMIT ?`
    return this.db.query(sql).all(...(o.root !== undefined ? [o.root] : []), match, ...roles, ...ex, o.limit ?? 300) as Hit[]
  }

  session(id: string): (SessionRow & { count: number }) | undefined {
    return (this.db.query("SELECT s.*, (SELECT count(*) FROM messages m WHERE m.session_id = s.id AND m.kind = 'message') AS count FROM sessions s WHERE s.id = ?").get(id) ?? undefined) as
      | (SessionRow & { count: number })
      | undefined
  }

  /** Nothing was ever said in it: no rows at all — a turn undone still counts as said,
   *  since /redo could bring it back. */
  isEmpty(sessionId: string): boolean {
    return !this.db.query("SELECT 1 FROM messages WHERE session_id = ? LIMIT 1").get(sessionId)
  }

  /** Visible messages with their ids, in order. */
  rows(sessionId: string): Row[] {
    return this.db.query("SELECT id, role, json, created FROM messages WHERE session_id = ? AND kind = 'message' ORDER BY id").all(sessionId) as Row[]
  }

  search(query: string, limit = 20): { session_id: string; id: number }[] {
    return this.db
      .query(
        "SELECT m.session_id, m.id FROM messages_fts f JOIN messages m ON m.id = f.rowid WHERE messages_fts MATCH ? ORDER BY rank LIMIT ?",
      )
      .all(query, limit) as { session_id: string; id: number }[]
  }
}

/** A compaction summary replayed as a user turn and an acknowledgement (LLeMbas compaction). */
export function summaryTurns(summary: string): Message[] {
  return [
    { role: "user", parts: [{ type: "text", text: `The earlier part of this conversation was compacted. Its summary:\n\n${summary}` }] },
    { role: "assistant", parts: [{ type: "text", text: "Understood — I will carry on from that summary." }] },
  ]
}