AgentHub/server/lib/db-retention.js

/**
 * @file Database size upkeep (Settings → Data → Session cleanup): the shared
 * session purge, the OPT-IN automatic retention, the bounded reclaim of free
 * pages, and the manual Compact.
 *
 * - `purgeSessionIds` is THE deletion of whole sessions — events, agents,
 *   token usage (daily and hourly) and the session row — used by both the
 *   manual `POST /api/settings/cleanup` and the automatic retention, so the two
 *   can never leave different leftovers.
 * - Retention is OFF unless `AGENTHUB_RETENTION_DAYS` (settings file, 1–3650)
 *   is set; it never deletes data by default. When set, the hourly
 *   `runScheduledMaintenance` purges finished (completed/error/abandoned)
 *   sessions that STARTED more than that many days ago — the manual purge's
 *   rule — in transactions of `RETENTION_BATCH` sessions, yielding to the event
 *   loop between them, and at most `RETENTION_MAX_PER_RUN` per run, so a first
 *   run on a long history never blocks hooks for long. Each batch deletes by
 *   `idx_events_session`, never a scan of the events table.
 * - With `auto_vacuum = INCREMENTAL` (server/db.js sets it on a new database
 *   and converts a small existing one at startup), the same run returns free
 *   pages to the OS once the freelist passes `RECLAIM_MIN_PAGES`, in slices of
 *   `RECLAIM_SLICE_PAGES`, at most `RECLAIM_MAX_PAGES` per run.
 * - `runCompact` (`POST /api/settings/compact`) is the one full VACUUM: in a
 *   worker thread on its own connection, setting `auto_vacuum = INCREMENTAL`
 *   first so a large existing database is converted, then a TRUNCATE
 *   checkpoint so the file actually shrinks. It holds the write lock for its
 *   whole run (tens of seconds on a GB-size file), so hooks wait or are retried
 *   meanwhile — it only ever runs on a click. This file is also its worker's
 *   entry point.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

"use strict";

const { Worker, isMainThread, workerData, parentPort } = require("node:worker_threads");

const KEYS = { days: "AGENTHUB_RETENTION_DAYS" };
const MAX_DAYS = 3650;
const RETENTION_BATCH = 25;
const RETENTION_MAX_PER_RUN = 5000;
const RECLAIM_MIN_PAGES = 1024;
const RECLAIM_SLICE_PAGES = 1024;
const RECLAIM_MAX_PAGES = 65536;
const PURGE_CHUNK = 500;
const PURGEABLE = "status IN ('completed','error','abandoned')";

const coded = (message, code) => Object.assign(new Error(message), { code });
const yieldLoop = () => new Promise((r) => setImmediate(r));

/** `{ days }` — 0 means off. Anything outside 1–MAX_DAYS reads as off. */
function retentionConfig(env = process.env) {
  const raw = env[KEYS.days];
  const n = Number(raw);
  const days = raw !== undefined && raw !== "" && Number.isInteger(n) && n >= 1 && n <= MAX_DAYS;
  return { days: days ? n : 0 };
}

/**
 * Delete these sessions and everything hung off them, in chunks (a single
 * IN (...) of every id would pass SQLite's bound-variable limit). Runs in the
 * CALLER's transaction when there is one. Returns the rows removed.
 */
function purgeSessionIds(db, ids) {
  const { forgetHourly } = require("./hourly-usage");
  const out = { purged_sessions: 0, purged_events: 0, purged_agents: 0 };
  for (let i = 0; i < ids.length; i += PURGE_CHUNK) {
    const chunk = ids.slice(i, i + PURGE_CHUNK);
    const placeholders = chunk.map(() => "?").join(",");
    // Cascading deletes handle agents/events, but token_usage FK might not cascade on all setups
    out.purged_events += db
      .prepare(`DELETE FROM events WHERE session_id IN (${placeholders})`)
      .run(...chunk).changes;
    out.purged_agents += db
      .prepare(`DELETE FROM agents WHERE session_id IN (${placeholders})`)
      .run(...chunk).changes;
    db.prepare(`DELETE FROM token_usage WHERE session_id IN (${placeholders})`).run(...chunk);
    db.prepare(`DELETE FROM token_usage_hourly WHERE session_id IN (${placeholders})`).run(
      ...chunk
    );
    forgetHourly(chunk);
    out.purged_sessions += db
      .prepare(`DELETE FROM sessions WHERE id IN (${placeholders})`)
      .run(...chunk).changes;
  }
  return out;
}

/** What `purgeSessionIds` would delete, without deleting it (the dry run). */
function countSessionIds(db, ids) {
  const out = { purged_sessions: ids.length, purged_events: 0, purged_agents: 0 };
  for (let i = 0; i < ids.length; i += PURGE_CHUNK) {
    const chunk = ids.slice(i, i + PURGE_CHUNK);
    const placeholders = chunk.map(() => "?").join(",");
    out.purged_events += db
      .prepare(`SELECT COUNT(*) AS n FROM events WHERE session_id IN (${placeholders})`)
      .get(...chunk).n;
    out.purged_agents += db
      .prepare(`SELECT COUNT(*) AS n FROM agents WHERE session_id IN (${placeholders})`)
      .get(...chunk).n;
  }
  return out;
}

/** Finished sessions that started before `cutoff` (ISO), at most `limit`. */
function purgeableBefore(db, cutoff, limit = -1) {
  return db
    .prepare(`SELECT id, started_at FROM sessions WHERE ${PURGEABLE} AND started_at < ? LIMIT ?`)
    .all(cutoff, limit);
}

/**
 * Purge finished sessions older than `days`, a batch per transaction with a
 * yield between batches. Resolves to the totals; never throws.
 */
async function applyRetention({
  db = require("../db").db,
  days,
  now = Date.now(),
  batch = RETENTION_BATCH,
  maxSessions = RETENTION_MAX_PER_RUN,
} = {}) {
  const out = { purged_sessions: 0, purged_events: 0, purged_agents: 0 };
  if (!Number.isInteger(days) || days < 1) return out;
  const cutoff = new Date(now - days * 86400 * 1000).toISOString();
  try {
    while (out.purged_sessions < maxSessions) {
      const ids = purgeableBefore(
        db,
        cutoff,
        Math.min(batch, maxSessions - out.purged_sessions)
      ).map((r) => r.id);
      if (ids.length === 0) break;
      const r = db.transaction(() => purgeSessionIds(db, ids))();
      out.purged_sessions += r.purged_sessions;
      out.purged_events += r.purged_events;
      out.purged_agents += r.purged_agents;
      if (r.purged_sessions === 0) break; // nothing removable: never spin
      await yieldLoop();
    }
  } catch (err) {
    out.error = err.message;
  }
  return out;
}

/** One `incremental_vacuum(n)`, stepped to the end on either driver. */
function incrementalVacuum(db, pages) {
  const sql = `PRAGMA incremental_vacuum(${Math.max(1, Math.floor(pages))})`;
  try {
    db.prepare(sql).all(); // node:sqlite: one row per freed page
  } catch (err) {
    if (!/does not return data/i.test(err.message)) throw err;
    db.pragma(`incremental_vacuum(${Math.max(1, Math.floor(pages))})`); // better-sqlite3
  }
}

/**
 * Hand free pages back to the OS, a slice at a time. A no-op unless the
 * database is in INCREMENTAL auto_vacuum and its freelist is over `minPages`.
 */
async function reclaimFreePages({
  db = require("../db").db,
  minPages = RECLAIM_MIN_PAGES,
  slice = RECLAIM_SLICE_PAGES,
  maxPages = RECLAIM_MAX_PAGES,
} = {}) {
  const out = { reclaimed_pages: 0 };
  try {
    if (db.pragma("auto_vacuum", { simple: true }) !== 2) return out;
    let free = db.pragma("freelist_count", { simple: true });
    if (free < minPages) return out;
    while (free > 0 && out.reclaimed_pages < maxPages) {
      incrementalVacuum(db, Math.min(slice, free, maxPages - out.reclaimed_pages));
      const after = db.pragma("freelist_count", { simple: true });
      if (after >= free) break;
      out.reclaimed_pages += free - after;
      free = after;
      await yieldLoop();
    }
  } catch (err) {
    out.error = err.message;
  }
  return out;
}

let lastRun = null;
let maintaining = null;

/**
 * The hourly upkeep: `maintainDatabase` (checkpoint + optimize), then the
 * opt-in retention, then the bounded free-page reclaim. A call while one runs
 * joins it.
 */
function runScheduledMaintenance(opts = {}) {
  if (maintaining) return maintaining;
  maintaining = maintain(opts).finally(() => {
    maintaining = null;
  });
  return maintaining;
}

async function maintain({ db } = {}) {
  const dbm = require("../db");
  const conn = db || dbm.db;
  const started = new Date();
  const out = { at: started.toISOString() };
  try {
    if (!db) dbm.maintainDatabase();
  } catch {
    /* maintainDatabase reports its own errors in its result */
  }
  const { days } = retentionConfig();
  if (days) {
    out.retention = await applyRetention({ db: conn, days });
    if (out.retention.purged_sessions > 0) {
      console.log(
        `[retention] purged ${out.retention.purged_sessions} session(s) older than ${days} days`
      );
    }
  }
  out.reclaim = await reclaimFreePages({ db: conn });
  out.ms = Date.now() - started;
  lastRun = out;
  return out;
}

/** Page counts for the Settings card. */
function spaceInfo(db = require("../db").db) {
  const pageSize = db.pragma("page_size", { simple: true });
  const pageCount = db.pragma("page_count", { simple: true });
  const free = db.pragma("freelist_count", { simple: true });
  const mode = db.pragma("auto_vacuum", { simple: true });
  return {
    auto_vacuum: mode === 2 ? "incremental" : mode === 1 ? "full" : "none",
    page_size: pageSize,
    page_count: pageCount,
    freelist_pages: free,
    free_bytes: free * pageSize,
  };
}

let compacting = null;
let lastCompact = null;

function vacuumInWorker(dbPath) {
  return new Promise((resolve, reject) => {
    const w = new Worker(__filename, { workerData: { agenthubCompactWorker: true, dbPath } });
    w.once("message", (m) => (m.error ? reject(coded(m.error.message, m.error.code)) : resolve(m)));
    w.once("error", reject);
    w.once(
      "exit",
      (code) => code !== 0 && reject(coded(`compact worker exited (${code})`, "ECOMPACT"))
    );
  });
}

/** Start (or join) the manual Compact. The outcome is `lastCompactResult()`. */
function runCompact() {
  if (compacting) return compacting;
  const fs = require("node:fs");
  const { DB_PATH } = require("../db");
  const size = () => {
    try {
      return fs.statSync(DB_PATH).size;
    } catch {
      return 0;
    }
  };
  compacting = (async () => {
    const started = Date.now();
    const before = size();
    try {
      await vacuumInWorker(DB_PATH);
      lastCompact = {
        at: new Date(started).toISOString(),
        ok: true,
        size_before: before,
        size_after: size(),
        ms: Date.now() - started,
      };
      return lastCompact;
    } catch (err) {
      lastCompact = { at: new Date(started).toISOString(), ok: false, error: err.message };
      throw err;
    }
  })().finally(() => {
    compacting = null;
  });
  return compacting;
}

const isCompacting = () => compacting !== null;
const lastCompactResult = () => lastCompact;
const lastMaintenance = () => lastRun;

if (!isMainThread && workerData?.agenthubCompactWorker) {
  // Same driver order as server/db.js: better-sqlite3, else node:sqlite.
  let conn;
  try {
    try {
      const Database = require("better-sqlite3");
      conn = new Database(workerData.dbPath);
    } catch {
      const Database = require("../compat-sqlite");
      conn = new Database(workerData.dbPath);
    }
    conn.pragma("busy_timeout = 30000");
    conn.pragma("auto_vacuum = INCREMENTAL"); // stored by the VACUUM below
    conn.exec("VACUUM");
    try {
      conn.pragma("wal_checkpoint(TRUNCATE)");
    } catch {
      /* a reader holds the WAL: the periodic checkpoint finishes it */
    }
    parentPort.postMessage({ ok: true });
  } catch (err) {
    parentPort.postMessage({ error: { message: err.message, code: err.code || "ECOMPACT" } });
  } finally {
    conn?.close();
  }
}

module.exports = {
  KEYS,
  MAX_DAYS,
  retentionConfig,
  purgeSessionIds,
  countSessionIds,
  purgeableBefore,
  applyRetention,
  reclaimFreePages,
  runScheduledMaintenance,
  spaceInfo,
  runCompact,
  isCompacting,
  lastCompactResult,
  lastMaintenance,
};