agent-orchestrator/ui-electrobun/src/bun/session-db.ts

/**
 * Session persistence — SQLite-backed agent session & message storage.
 * Uses bun:sqlite. DB: ~/.config/agor/settings.db (shared with SettingsDb).
 *
 * Tables: agent_sessions, agent_messages
 */

import { Database } from "bun:sqlite";
import { homedir } from "os";
import { join } from "path";
import { openDb } from "./db-utils.ts";

// ── DB path ──────────────────────────────────────────────────────────────────

const CONFIG_DIR = join(homedir(), ".config", "agor");
const DB_PATH = join(CONFIG_DIR, "settings.db");

// ── Types ────────────────────────────────────────────────────────────────────

export interface StoredSession {
  projectId: string;
  sessionId: string;
  provider: string;
  status: string;
  costUsd: number;
  inputTokens: number;
  outputTokens: number;
  model: string;
  error?: string;
  createdAt: number;
  updatedAt: number;
}

export interface StoredMessage {
  sessionId: string;
  msgId: string;
  role: string;
  content: string;
  toolName?: string;
  toolInput?: string;
  timestamp: number;
  costUsd: number;
  inputTokens: number;
  outputTokens: number;
}

// ── Schema ───────────────────────────────────────────────────────────────────

const SESSION_SCHEMA = `
CREATE TABLE IF NOT EXISTS agent_sessions (
  project_id   TEXT NOT NULL,
  session_id   TEXT PRIMARY KEY,
  provider     TEXT NOT NULL,
  status       TEXT NOT NULL DEFAULT 'idle',
  cost_usd     REAL NOT NULL DEFAULT 0,
  input_tokens INTEGER NOT NULL DEFAULT 0,
  output_tokens INTEGER NOT NULL DEFAULT 0,
  model        TEXT NOT NULL DEFAULT '',
  error        TEXT,
  created_at   INTEGER NOT NULL,
  updated_at   INTEGER NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_agent_sessions_project
  ON agent_sessions(project_id);

CREATE TABLE IF NOT EXISTS agent_messages (
  session_id    TEXT NOT NULL,
  msg_id        TEXT NOT NULL,
  role          TEXT NOT NULL,
  content       TEXT NOT NULL DEFAULT '',
  tool_name     TEXT,
  tool_input    TEXT,
  timestamp     INTEGER NOT NULL,
  cost_usd      REAL NOT NULL DEFAULT 0,
  input_tokens  INTEGER NOT NULL DEFAULT 0,
  output_tokens INTEGER NOT NULL DEFAULT 0,
  PRIMARY KEY (session_id, msg_id),
  FOREIGN KEY (session_id) REFERENCES agent_sessions(session_id) ON DELETE CASCADE
);

CREATE INDEX IF NOT EXISTS idx_agent_messages_session
  ON agent_messages(session_id, timestamp);
`;

// ── SessionDb class ──────────────────────────────────────────────────────────

export class SessionDb {
  private db: Database;

  constructor() {
    this.db = openDb(DB_PATH, { foreignKeys: true });
    this.db.exec(SESSION_SCHEMA);
  }

  // ── Sessions ─────────────────────────────────────────────────────────────

  saveSession(s: StoredSession): void {
    this.db
      .query(
        `INSERT INTO agent_sessions
           (project_id, session_id, provider, status, cost_usd,
            input_tokens, output_tokens, model, error, created_at, updated_at)
         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)
         ON CONFLICT(session_id) DO UPDATE SET
           status = excluded.status,
           cost_usd = excluded.cost_usd,
           input_tokens = excluded.input_tokens,
           output_tokens = excluded.output_tokens,
           model = excluded.model,
           error = excluded.error,
           updated_at = excluded.updated_at`
      )
      .run(
        s.projectId,
        s.sessionId,
        s.provider,
        s.status,
        s.costUsd,
        s.inputTokens,
        s.outputTokens,
        s.model,
        s.error ?? null,
        s.createdAt,
        s.updatedAt
      );
  }

  loadSession(projectId: string): StoredSession | null {
    const row = this.db
      .query<
        {
          project_id: string;
          session_id: string;
          provider: string;
          status: string;
          cost_usd: number;
          input_tokens: number;
          output_tokens: number;
          model: string;
          error: string | null;
          created_at: number;
          updated_at: number;
        },
        [string]
      >(
        `SELECT * FROM agent_sessions
         WHERE project_id = ?
         ORDER BY updated_at DESC LIMIT 1`
      )
      .get(projectId);

    if (!row) return null;
    return {
      projectId: row.project_id,
      sessionId: row.session_id,
      provider: row.provider,
      status: row.status,
      costUsd: row.cost_usd,
      inputTokens: row.input_tokens,
      outputTokens: row.output_tokens,
      model: row.model,
      error: row.error ?? undefined,
      createdAt: row.created_at,
      updatedAt: row.updated_at,
    };
  }

  listSessionsByProject(projectId: string): StoredSession[] {
    const rows = this.db
      .query<
        {
          project_id: string;
          session_id: string;
          provider: string;
          status: string;
          cost_usd: number;
          input_tokens: number;
          output_tokens: number;
          model: string;
          error: string | null;
          created_at: number;
          updated_at: number;
        },
        [string]
      >(
        `SELECT * FROM agent_sessions
         WHERE project_id = ?
         ORDER BY updated_at DESC LIMIT 20`
      )
      .all(projectId);

    return rows.map((r) => ({
      projectId: r.project_id,
      sessionId: r.session_id,
      provider: r.provider,
      status: r.status,
      costUsd: r.cost_usd,
      inputTokens: r.input_tokens,
      outputTokens: r.output_tokens,
      model: r.model,
      error: r.error ?? undefined,
      createdAt: r.created_at,
      updatedAt: r.updated_at,
    }));
  }

  // ── Messages ─────────────────────────────────────────────────────────────

  saveMessage(m: StoredMessage): void {
    this.db
      .query(
        `INSERT INTO agent_messages
           (session_id, msg_id, role, content, tool_name, tool_input,
            timestamp, cost_usd, input_tokens, output_tokens)
         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)
         ON CONFLICT(session_id, msg_id) DO NOTHING`
      )
      .run(
        m.sessionId,
        m.msgId,
        m.role,
        m.content,
        m.toolName ?? null,
        m.toolInput ?? null,
        m.timestamp,
        m.costUsd,
        m.inputTokens,
        m.outputTokens
      );
  }

  saveMessages(msgs: StoredMessage[]): void {
    const stmt = this.db.prepare(
      `INSERT INTO agent_messages
         (session_id, msg_id, role, content, tool_name, tool_input,
          timestamp, cost_usd, input_tokens, output_tokens)
       VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)
       ON CONFLICT(session_id, msg_id) DO NOTHING`
    );

    const tx = this.db.transaction(() => {
      for (const m of msgs) {
        stmt.run(
          m.sessionId,
          m.msgId,
          m.role,
          m.content,
          m.toolName ?? null,
          m.toolInput ?? null,
          m.timestamp,
          m.costUsd,
          m.inputTokens,
          m.outputTokens
        );
      }
    });
    tx();
  }

  loadMessages(sessionId: string): StoredMessage[] {
    const rows = this.db
      .query<
        {
          session_id: string;
          msg_id: string;
          role: string;
          content: string;
          tool_name: string | null;
          tool_input: string | null;
          timestamp: number;
          cost_usd: number;
          input_tokens: number;
          output_tokens: number;
        },
        [string]
      >(
        `SELECT * FROM agent_messages
         WHERE session_id = ?
         ORDER BY timestamp ASC`
      )
      .all(sessionId);

    return rows.map((r) => ({
      sessionId: r.session_id,
      msgId: r.msg_id,
      role: r.role,
      content: r.content,
      toolName: r.tool_name ?? undefined,
      toolInput: r.tool_input ?? undefined,
      timestamp: r.timestamp,
      costUsd: r.cost_usd,
      inputTokens: r.input_tokens,
      outputTokens: r.output_tokens,
    }));
  }

  // ── Cleanup ──────────────────────────────────────────────────────────────

  /** Delete sessions older than maxAgeDays for a project, keeping at most keepCount. */
  pruneOldSessions(projectId: string, keepCount = 10): void {
    this.db
      .query(
        `DELETE FROM agent_sessions
         WHERE project_id = ?1
           AND session_id NOT IN (
             SELECT session_id FROM agent_sessions
             WHERE project_id = ?1
             ORDER BY updated_at DESC
             LIMIT ?2
           )`
      )
      .run(projectId, keepCount);
  }

  close(): void {
    this.db.close();
  }
}

// Singleton
export const sessionDb = new SessionDb();