/** * @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ĩ */ 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} */ 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, };