AgentHub/server/lib/user-quota.js

/**
 * @file user-quota.js
 * @description Spend quotas per user (server mode only; local mode has one
 * user and never reads any of this). An admin sets a DEFAULT quota (row
 * `user_id = 0` of `user_quotas`) and may override it per user — a limit of
 * their own, or none at all (`limit_usd` NULL). A quota is `limit_usd` per
 * `period`: `day` (from the local midnight of `tz_offset`, the admin's
 * browser's getTimezoneOffset() when they saved it) or `5h` (a sliding
 * window: the current UTC hour and the four before it — spend booked in hour
 * H counts until H + 5 h). Spend is the users page's own booking
 * (`hourlyTokenRows` over the sessions a user OWNS through a workspace,
 * priced by `calculateCost`), never a second formula.
 *
 * A user over their quota cannot chat on until it resets:
 *  - `assertSpawnAllowed` refuses every new pane (`EQUOTA`, 429) — start,
 *    resume and a prompt from the workspace home all pass `spawnRun`;
 *  - every pane carries `AGENTHUB_QUOTA_FILE` (`flagEnv`), and `sweep` writes
 *    that file `{blocked, until, message}` — the host's under the data dir,
 *    a member's (S3) in their handoff dir through `write-as` — so the hook
 *    handler refuses a `UserPromptSubmit` of Claude Code or Codex locally,
 *    without asking the dashboard (scripts/hook-handler.js; fail-open);
 *  - a member's omp on the shared login loses its gateway key while blocked
 *    (`isRunBlocked`, omp-share.js `clientsFor`);
 *  - a pane no hook can stop (Antigravity, Cursor, a Codex without trusted
 *    hooks, an omp on its own login, any other catalog tool) is ended.
 * Spend lands after a turn, so a user may end one turn past the limit.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const fs = require("node:fs");
const path = require("node:path");

const PERIODS = ["day", "5h"];
const DEFAULT_ID = 0;
const HOUR_MS = 3_600_000;
const WINDOW_HOURS = 5;

const db = () => require("../db").db;
const fail = (code, message) => Object.assign(new Error(message), { code });
const isServerMode = () => require("./auth").isServerMode();
const hourOf = (ms) => new Date(ms).toISOString().slice(0, 13);
const hourMs = (h) => Date.parse(`${h}:00:00.000Z`);
const round = (n) => Math.round(n * 10000) / 10000;

// ── settings ──────────────────────────────────────────────────────────────

function row(userId) {
  return db()
    .prepare("SELECT limit_usd, period, tz_offset FROM user_quotas WHERE user_id = ?")
    .get(userId);
}

const asConfig = (r, source) =>
  r && r.limit_usd != null
    ? { source, limit_usd: r.limit_usd, period: r.period, tz_offset: r.tz_offset }
    : null;

/** The default quota, or null when there is none. */
function getDefault() {
  return asConfig(row(DEFAULT_ID), "default");
}

/**
 * A user's own setting: `default` (no row — the default applies),
 * `unlimited`, or `custom` with its limit.
 */
function getOverride(userId) {
  const r = row(userId);
  if (!r) return { mode: "default" };
  if (r.limit_usd == null) return { mode: "unlimited" };
  return { mode: "custom", limit_usd: r.limit_usd, period: r.period, tz_offset: r.tz_offset };
}

/** The quota that binds `userId` now, or null (no limit). */
function effective(userId) {
  const r = row(userId);
  if (r) return asConfig(r, "user");
  return getDefault();
}

/** `{limit_usd, period, tz_offset}` checked; throws EBADQUOTA. */
function validate(cfg) {
  const limit = Number(cfg?.limit_usd);
  if (!Number.isFinite(limit) || limit <= 0 || limit > 1_000_000) {
    throw fail("EBADQUOTA", "limit_usd must be a positive number of dollars");
  }
  if (!PERIODS.includes(cfg.period)) {
    throw fail("EBADQUOTA", `period must be one of: ${PERIODS.join(", ")}`);
  }
  const tz = cfg.tz_offset == null ? 0 : cfg.tz_offset;
  if (!Number.isInteger(tz) || Math.abs(tz) > 14 * 60) {
    throw fail("EBADQUOTA", "tz_offset must be minutes within ±840");
  }
  return { limit_usd: limit, period: cfg.period, tz_offset: tz };
}

function upsert(userId, limit, period, tz) {
  db()
    .prepare(
      `INSERT INTO user_quotas (user_id, limit_usd, period, tz_offset, updated_at)
       VALUES (?, ?, ?, ?, strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
       ON CONFLICT(user_id) DO UPDATE SET limit_usd = excluded.limit_usd,
         period = excluded.period, tz_offset = excluded.tz_offset,
         updated_at = excluded.updated_at`
    )
    .run(userId, limit, period, tz);
}

/** Set the default quota (`null` removes it: no limit for anyone without their own). */
function setDefault(cfg) {
  if (cfg == null) {
    db().prepare("DELETE FROM user_quotas WHERE user_id = ?").run(DEFAULT_ID);
    return null;
  }
  const v = validate(cfg);
  upsert(DEFAULT_ID, v.limit_usd, v.period, v.tz_offset);
  return getDefault();
}

/** A user's setting checked: `{mode}` or `{mode: "custom", limit_usd, period, tz_offset}`. */
function validateUserQuota(q) {
  const mode = q?.mode;
  if (mode === "default" || mode === "unlimited") return { mode };
  if (mode === "custom") return { mode, ...validate(q) };
  throw fail("EBADQUOTA", "quota.mode must be one of: default, unlimited, custom");
}

/** A user's setting: `{mode: "default" | "unlimited" | "custom", …}`. */
function setUserQuota(userId, q) {
  if (!Number.isInteger(userId) || userId <= 0) throw fail("EBADQUOTA", "no such user");
  const v = validateUserQuota(q);
  if (v.mode === "default") {
    db().prepare("DELETE FROM user_quotas WHERE user_id = ?").run(userId);
  } else if (v.mode === "unlimited") {
    upsert(userId, null, "day", 0);
  } else {
    upsert(userId, v.limit_usd, v.period, v.tz_offset);
  }
  return getOverride(userId);
}

/** `$5/day`, `$2.50/5h`, `no limit` — for the audit log. */
function describe(q) {
  if (!q || q.mode === "unlimited") return "no limit";
  if (q.mode === "default") return "the default";
  return `$${q.limit_usd}/${q.period}`;
}

// ── spend ─────────────────────────────────────────────────────────────────

/**
 * The window a quota counts over at `now`: `since` (a UTC hour,
 * `YYYY-MM-DDTHH`, as hourly token rows are keyed) and, for a day, when the
 * next one starts. Buckets are whole UTC hours, so a half-hour zone's day
 * starts up to 30 min early.
 */
function windowOf(cfg, now) {
  if (cfg.period === "5h") {
    return { since: hourOf(now - (WINDOW_HOURS - 1) * HOUR_MS), nextReset: null };
  }
  const off = cfg.tz_offset || 0;
  const local = new Date(now - off * 60_000);
  local.setUTCHours(0, 0, 0, 0);
  const start = local.getTime() + off * 60_000;
  return { since: hourOf(start), nextReset: start + 24 * HOUR_MS };
}

/**
 * When a 5 h quota opens again: the oldest hours leave the window first, so
 * it is the end of the first hour whose leaving takes the spend under the limit.
 */
function slidingReset(byHour, spend, limit) {
  let left = spend;
  for (const [h, cost] of [...byHour].sort((a, b) => (a[0] < b[0] ? -1 : 1))) {
    left -= cost;
    if (left < limit) return hourMs(h) + WINDOW_HOURS * HOUR_MS;
  }
  return null;
}

/**
 * Every user's quota and spend now: Map(user id → status). Only users with a
 * quota are listed. One token read for all of them.
 * @returns {Map<number, {limit_usd, period, tz_offset, source, spend, pct,
 *   exceeded: boolean, resetsAt: string|null, since: string}>}
 */
function statusAll({ now = Date.now(), userIds = null } = {}) {
  const out = new Map();
  if (!isServerMode()) return out;
  const ids =
    userIds ||
    db()
      .prepare("SELECT id FROM users WHERE disabled = 0")
      .all()
      .map((r) => r.id);
  const quotas = [];
  for (const id of ids) {
    const cfg = effective(id);
    if (cfg) quotas.push({ id, cfg, win: windowOf(cfg, now) });
  }
  if (!quotas.length) return out;
  const { calculateCost, hourlyTokenRows } = require("../routes/pricing");
  const { OWNED_SESSIONS } = require("../routes/users");
  const rules = require("../db").stmts.listPricing.all();
  const oldest = quotas.reduce((a, q) => (q.win.since < a ? q.win.since : a), quotas[0].win.since);
  const owner = new Map();
  for (const r of db().prepare(OWNED_SESSIONS).all()) {
    if (!owner.has(r.session_id)) owner.set(r.session_id, r.uid);
  }
  // A session that spent in the window was touched in it.
  const rows = hourlyTokenRows(
    `WHERE COALESCE(NULLIF(s.updated_at, ''), s.started_at) >= ? AND s.id IN (SELECT session_id FROM (${OWNED_SESSIONS}))`,
    [oldest]
  );
  const byUser = new Map();
  for (const r of rows) {
    const uid = owner.get(r.session_id);
    if (uid == null) continue;
    if (!byUser.has(uid)) byUser.set(uid, []);
    byUser.get(uid).push(r);
  }
  for (const { id, cfg, win } of quotas) {
    const mine = (byUser.get(id) || []).filter((r) => r.date >= win.since);
    const spend = mine.length ? calculateCost(mine, rules).total_cost : 0;
    const exceeded = spend >= cfg.limit_usd;
    let resetsAt = win.nextReset;
    if (cfg.period === "5h") {
      const hours = new Map();
      for (const r of mine) {
        if (!hours.has(r.date)) hours.set(r.date, []);
        hours.get(r.date).push(r);
      }
      const byHour = new Map(
        [...hours].map(([h, list]) => [h, calculateCost(list, rules).total_cost])
      );
      resetsAt = exceeded ? slidingReset(byHour, spend, cfg.limit_usd) : null;
    }
    out.set(id, {
      ...cfg,
      spend: round(spend),
      pct: round((spend / cfg.limit_usd) * 100),
      exceeded,
      resetsAt: resetsAt ? new Date(resetsAt).toISOString() : null,
      since: `${win.since}:00:00.000Z`,
    });
  }
  return out;
}

/** One user's status, or null without a quota. */
function statusFor(userId, opts = {}) {
  if (userId == null) return null;
  return statusAll({ ...opts, userIds: [userId] }).get(userId) || null;
}

/** "HH:MM" (and the date when it is not today) in the quota's own zone. */
function localTime(iso, tzOffset, now = Date.now()) {
  if (!iso) return null;
  const off = (tzOffset || 0) * 60_000;
  const at = new Date(Date.parse(iso) - off).toISOString();
  const today = new Date(now - off).toISOString().slice(0, 10);
  const utc = -(tzOffset || 0) / 60;
  const zone = `UTC${utc >= 0 ? "+" : ""}${utc}`;
  return `${at.slice(0, 10) === today ? "" : `${at.slice(0, 10)} `}${at.slice(11, 16)} ${zone}`;
}

/** The text a refused prompt or start shows. */
function message(s, now = Date.now()) {
  const per = s.period === "5h" ? "5 hours" : "day";
  const when = localTime(s.resetsAt, s.tz_offset, now);
  return (
    `AgentHub: your spend quota is used up — $${s.spend.toFixed(2)} of $${s.limit_usd.toFixed(2)} per ${per}.` +
    (when ? ` It opens again at ${when}.` : "") +
    " Ask an admin to raise it if you need more."
  );
}

/** Throws EQUOTA when the workspace owner is over their quota. No-op in local mode. */
function assertSpawnAllowed(ownerId) {
  if (!isServerMode() || ownerId == null) return;
  const s = statusFor(ownerId);
  if (s?.exceeded) throw Object.assign(fail("EQUOTA", message(s)), { quota: s });
}

// ── the flag file the hook handler reads ─────────────────────────────────

const FLAG = "quota.json";
const hostDir = () => path.join(require("./claude-home").getDataDir(), "user-quota");
const hostFlag = (userId) => path.join(hostDir(), `u${userId}.json`);

/**
 * The env a pane of `ownerId` gets: where its hooks look for the flag —
 * the member's handoff dir (S3: the pane runs as them) or the host's data dir.
 */
function flagEnv(ownerId, member) {
  if (!isServerMode() || ownerId == null) return {};
  const file = member ? path.join(require("./run-as").handoffDir(member), FLAG) : hostFlag(ownerId);
  return { AGENTHUB_QUOTA_FILE: file };
}

function writeHostFlag(userId, text) {
  fs.mkdirSync(hostDir(), { recursive: true, mode: 0o700 });
  const file = hostFlag(userId);
  const tmp = `${file}.${process.pid}.tmp`;
  fs.writeFileSync(tmp, text, { mode: 0o600 });
  fs.renameSync(tmp, file);
}

/** Write a user's flag everywhere their panes may read it. Best effort per place. */
function writeFlags(user, s, now) {
  const body = s?.exceeded
    ? { blocked: true, until: s.resetsAt, message: message(s, now) }
    : { blocked: false };
  const text = `${JSON.stringify(body)}\n`;
  try {
    writeHostFlag(user.id, text);
  } catch (err) {
    console.warn(`user quota: flag for ${user.username} not written:`, err.message);
  }
  let acct = null;
  try {
    acct = require("./run-as").accountForUser(user.username);
  } catch {
    acct = null;
  }
  if (!acct) return;
  try {
    require("./run-as").handoffWrite(acct, FLAG, text);
  } catch (err) {
    console.warn(`user quota: ${user.username}'s flag not written:`, err.message);
  }
}

// ── the sweep ────────────────────────────────────────────────────────────

/** Users over their quota now (by the last sweep). */
let blocked = new Set();
/** user id → the flag text last written, so an unchanged state writes nothing. */
const written = new Map();

function laneOwner(laneId) {
  return db().prepare("SELECT owner_id FROM lanes WHERE id = ?").get(laneId)?.owner_id ?? null;
}

/** Is run `runId`'s workspace owner over quota (by the last sweep)? omp-share asks. */
function isRunBlocked(runId) {
  if (!blocked.size) return false;
  const laneId = require("./pty-run").laneIdFromRunId(runId);
  return laneId != null && blocked.has(laneOwner(laneId));
}

/**
 * Can a hook refuse this pane's next prompt? Claude Code always (the plugin's
 * hooks, a member's managed-settings hooks); Codex only on the host's own
 * CODEX_HOME with our UserPromptSubmit hook installed and trusted (members
 * get no Codex hooks); a member's omp on the shared login through its
 * gateway key. Everything else must be ended.
 */
function hookCanStop(run, member) {
  if (run.tool === "claude" || run.tool === "shell") return true;
  if (run.tool === "codex" && !member) {
    try {
      const ev = require("./codex-hooks").status().events?.UserPromptSubmit;
      return !!(ev?.installed && ev?.trusted);
    } catch {
      return false;
    }
  }
  if (run.tool === "omp" && member) {
    try {
      const share = require("./omp-share");
      return share.available() && share.enabled();
    } catch {
      return false;
    }
  }
  return false;
}

/** End the live panes of `userIds` that no hook can stop. Returns the run ids ended. */
function endUnstoppable(userIds) {
  if (!userIds.size) return [];
  const runs = require("./pty-run");
  const runAs = require("./run-as");
  const ended = [];
  for (const run of runs.listRuns()) {
    if (run.status !== "running" || !userIds.has(laneOwner(run.laneId))) continue;
    if (hookCanStop(run, runAs.accountForLane(run.laneId))) continue;
    try {
      if (runs.killRun(run.id)) ended.push(run.id);
    } catch (err) {
      console.warn(`user quota: ${run.id} not ended:`, err.message);
    }
  }
  return ended;
}

/**
 * Recompute every user's state, rewrite the flags that changed, end the
 * panes no hook can stop and re-list the omp gateway's keys. Called every
 * minute by the server and right after an admin changes a quota.
 */
function sweep({ now = Date.now() } = {}) {
  if (!isServerMode()) return { blocked: [] };
  const users = db().prepare("SELECT id, username FROM users").all();
  const status = statusAll({ now });
  const next = new Set();
  for (const u of users) {
    const s = status.get(u.id) || null;
    if (s?.exceeded) next.add(u.id);
    const text = s?.exceeded ? `${s.resetsAt}|${s.limit_usd}` : "open";
    // A user never blocked has no flag to clear: nothing is written for them.
    if (written.get(u.id) === text || (!written.has(u.id) && !s?.exceeded && !hadFlag(u.id))) {
      continue;
    }
    writeFlags(u, s, now);
    written.set(u.id, text);
  }
  const changed = next.size !== blocked.size || [...next].some((id) => !blocked.has(id));
  blocked = next;
  const ended = endUnstoppable(next);
  if (changed) {
    try {
      require("./omp-share").syncClients();
    } catch {
      /* sharing off or not available */
    }
  }
  return { blocked: [...next], ended };
}

/** A flag left by an earlier server (it may still say blocked). */
function hadFlag(userId) {
  try {
    return fs.existsSync(hostFlag(userId));
  } catch {
    return false;
  }
}

/** Test seam: forget what the sweep remembered. */
function __reset() {
  blocked = new Set();
  written.clear();
}

module.exports = {
  PERIODS,
  getDefault,
  getOverride,
  effective,
  setDefault,
  setUserQuota,
  validateUserQuota,
  describe,
  windowOf,
  slidingReset,
  statusAll,
  statusFor,
  message,
  assertSpawnAllowed,
  flagEnv,
  isRunBlocked,
  hookCanStop,
  sweep,
  __reset,
};