AgentHub/server/__tests__/user-quota.test.js

/**
 * @file user-quota.test.js
 * @description Spend quotas per user (server/lib/user-quota.js): the default
 * and a user's own setting, spend over a day and a sliding 5 h window, the
 * admin routes (password again), `/api/me/quota`, the EQUOTA refusal of a new
 * pane, the sweep's flag files and the panes it ends, the omp gateway keys it
 * leaves out, and the hook handler refusing a prompt from the flag alone.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const { describe, it, before, after } = require("node:test");
const assert = require("node:assert/strict");
const path = require("node:path");
const fs = require("node:fs");
const os = require("node:os");
const { spawnSync } = require("node:child_process");

const TMP = fs.mkdtempSync(path.join(os.tmpdir(), "user-quota-"));
process.env.DASHBOARD_DB_PATH = path.join(TMP, "dashboard.db");
process.env.DASHBOARD_DATA_DIR = path.join(TMP, "data");
process.env.CLAUDE_HOME = path.join(TMP, ".claude");
process.env.CODEX_HOME = path.join(TMP, ".codex");
process.env.DASHBOARD_LIVENESS_PROBE = "0";
process.env.DASHBOARD_QUOTA_SWEEP_MS = "0";

const { createApp, startServer } = require("../index");
const { db } = require("../db");
const auth = require("../lib/auth");
const quota = require("../lib/user-quota");
const runs = require("../lib/pty-run");

let server;
let BASE;
const BOSS_PASS = "admin-pass-1";
const ids = {};
let laneId;

async function call(method, p, { cookie, body } = {}) {
  const res = await fetch(BASE + p, {
    method,
    headers: { "Content-Type": "application/json", ...(cookie ? { cookie } : {}) },
    body: body ? JSON.stringify(body) : undefined,
  });
  return { status: res.status, body: await res.json().catch(() => null) };
}

async function login(username, password) {
  const res = await fetch(BASE + "/api/auth/login", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ username, password }),
  });
  return res.headers.get("set-cookie").split(";")[0];
}

const HOUR = 3_600_000;
const hourOf = (ms) => new Date(ms).toISOString().slice(0, 13);

/** `tokens` output tokens of Sonnet 4.5 ($15 / M) booked to `sid` in the hour of `at`. */
function spend(sid, tokens, at = Date.now()) {
  db.prepare(
    `INSERT INTO token_usage (session_id, model, output_tokens)
     VALUES (?, 'claude-sonnet-4-5-20250929', ?)
     ON CONFLICT DO UPDATE SET output_tokens = output_tokens + excluded.output_tokens`
  ).run(sid, tokens);
  db.prepare(
    `INSERT INTO token_usage_hourly (session_id, hour, model, output_tokens)
     VALUES (?, ?, 'claude-sonnet-4-5-20250929', ?)`
  ).run(sid, hourOf(at), tokens);
}

before(async () => {
  process.env.AGENTHUB_MODE = "server";
  server = await startServer(createApp(), 0);
  BASE = `http://127.0.0.1:${server.address().port}`;
  ids.boss = (await auth.createUser("boss", BOSS_PASS, "admin")).id;
  ids.mia = (await auth.createUser("mia", "member-pass-1", "developer")).id;
  ids.lou = (await auth.createUser("lou", "member-pass-2", "developer")).id;
  laneId = Number(
    db
      .prepare(
        "INSERT INTO lanes (cwd, title, tool, owner_id) VALUES ('/w/mia', 'Mia', 'claude', ?)"
      )
      .run(ids.mia).lastInsertRowid
  );
  db.prepare(
    "INSERT INTO sessions (id, name, status, cwd, lane_id, updated_at) VALUES ('sm', 'mia work', 'active', '/w/mia', ?, ?)"
  ).run(laneId, new Date().toISOString());
  spend("sm", 1_000_000); // $15 now
});

after(() => {
  delete process.env.AGENTHUB_MODE;
  server?.close();
  fs.rmSync(TMP, { recursive: true, force: true });
});

describe("settings", () => {
  it("a user without a row of their own gets the default; unlimited and custom override it", () => {
    assert.equal(quota.effective(ids.mia), null);
    quota.setDefault({ limit_usd: 10, period: "day", tz_offset: -420 });
    assert.deepEqual(quota.effective(ids.mia), {
      source: "default",
      limit_usd: 10,
      period: "day",
      tz_offset: -420,
    });
    quota.setUserQuota(ids.mia, { mode: "unlimited" });
    assert.equal(quota.effective(ids.mia), null);
    assert.deepEqual(quota.getOverride(ids.mia), { mode: "unlimited" });
    quota.setUserQuota(ids.mia, { mode: "custom", limit_usd: 50, period: "5h" });
    assert.equal(quota.effective(ids.mia).source, "user");
    assert.equal(quota.effective(ids.mia).limit_usd, 50);
    quota.setUserQuota(ids.mia, { mode: "default" });
    assert.equal(quota.effective(ids.mia).source, "default");
    quota.setDefault(null);
    assert.equal(quota.effective(ids.mia), null);
  });

  it("refuses a bad limit, period or mode", () => {
    for (const bad of [
      { mode: "custom", limit_usd: 0, period: "day" },
      { mode: "custom", limit_usd: 5, period: "week" },
      { mode: "custom", limit_usd: 5, period: "day", tz_offset: 9999 },
      { mode: "nope" },
    ]) {
      assert.throws(() => quota.setUserQuota(ids.mia, bad), { code: "EBADQUOTA" });
    }
  });
});

describe("spend windows", () => {
  it("a day starts at the zone's local midnight and resets at the next one", () => {
    const now = Date.parse("2026-10-28T03:30:00Z"); // 10:30 at UTC+7
    const w = quota.windowOf({ period: "day", tz_offset: -420 }, now);
    assert.equal(w.since, "2026-10-27T17");
    assert.equal(new Date(w.nextReset).toISOString(), "2026-10-28T17:00:00.000Z");
  });

  it("5 h is the current hour and the four before it; it opens as the oldest hours leave", () => {
    const now = Date.parse("2026-10-28T10:20:00Z");
    assert.equal(quota.windowOf({ period: "5h" }, now).since, "2026-10-28T06");
    const byHour = new Map([
      ["2026-10-28T06", 4],
      ["2026-10-28T08", 3],
      ["2026-10-28T10", 2],
    ]);
    // 9 spent, limit 6: 06's $4 leaving brings it to 5 < 6 → at 06 + 5 h.
    assert.equal(
      new Date(quota.slidingReset(byHour, 9, 6)).toISOString(),
      "2026-10-28T11:00:00.000Z"
    );
    // limit 3: everything up to 08 must go → 08 + 5 h.
    assert.equal(
      new Date(quota.slidingReset(byHour, 9, 3)).toISOString(),
      "2026-10-28T13:00:00.000Z"
    );
  });

  it("counts the spend of the sessions a user owns, inside the window only", () => {
    db.prepare(
      "INSERT INTO sessions (id, name, status, cwd, lane_id, updated_at) VALUES ('old', 'old', 'completed', '/w/mia', ?, ?)"
    ).run(laneId, new Date().toISOString());
    spend("old", 1_000_000, Date.now() - 8 * HOUR); // outside 5 h
    quota.setUserQuota(ids.mia, { mode: "custom", limit_usd: 100, period: "5h" });
    const s = quota.statusFor(ids.mia);
    assert.equal(s.spend, 15);
    assert.equal(s.exceeded, false);
    assert.equal(s.resetsAt, null);
    assert.equal(quota.statusFor(ids.lou), null);
    quota.setUserQuota(ids.mia, { mode: "default" });
  });
});

describe("routes", () => {
  it("an admin sets the default with their password; the list shows each user's quota", async () => {
    const boss = await login("boss", BOSS_PASS);
    const noPass = await call("PUT", "/api/users/quota-default", {
      cookie: boss,
      body: { quota: { limit_usd: 10, period: "day" } },
    });
    assert.equal(noPass.status, 403);
    const bad = await call("PUT", "/api/users/quota-default", {
      cookie: boss,
      body: { quota: { limit_usd: -1, period: "day" }, currentPassword: BOSS_PASS },
    });
    assert.equal(bad.status, 400);
    assert.equal(bad.body.error.code, "EBADQUOTA");
    const ok = await call("PUT", "/api/users/quota-default", {
      cookie: boss,
      body: { quota: { limit_usd: 10, period: "day" }, currentPassword: BOSS_PASS },
    });
    assert.equal(ok.status, 200, JSON.stringify(ok.body));
    assert.equal(ok.body.quotaDefault.limit_usd, 10);

    const list = await call("GET", "/api/users", { cookie: boss });
    assert.equal(list.body.quotaDefault.limit_usd, 10);
    const mia = list.body.users.find((u) => u.id === ids.mia);
    assert.deepEqual(mia.quotaSetting, { mode: "default" });
    assert.equal(mia.quota.spend, 15);
    assert.equal(mia.quota.exceeded, true);
    assert.ok(mia.quota.resetsAt);
  });

  it("a developer reads their own quota and may not change any", async () => {
    const me = await login("mia", "member-pass-1");
    const r = await call("GET", "/api/me/quota", { cookie: me });
    assert.equal(r.status, 200);
    assert.equal(r.body.quota.exceeded, true);
    const put = await call("PUT", "/api/users/quota-default", {
      cookie: me,
      body: { quota: null, currentPassword: "member-pass-1" },
    });
    assert.equal(put.status, 403);
  });

  it("PATCH /api/users/:id {quota} overrides one user's, audited", async () => {
    const boss = await login("boss", BOSS_PASS);
    const r = await call("PATCH", `/api/users/${ids.mia}`, {
      cookie: boss,
      body: { quota: { mode: "custom", limit_usd: 40, period: "day" }, currentPassword: BOSS_PASS },
    });
    assert.equal(r.status, 200, JSON.stringify(r.body));
    assert.equal(quota.statusFor(ids.mia).exceeded, false);
    const audit = db
      .prepare("SELECT path FROM audit_log WHERE owner_id = ? ORDER BY id DESC LIMIT 1")
      .get(ids.mia);
    assert.match(audit.path, /spend quota \$40\/day/);
    const bad = await call("PATCH", `/api/users/${ids.mia}`, {
      cookie: boss,
      body: { quota: { mode: "lots" }, currentPassword: BOSS_PASS },
    });
    assert.equal(bad.status, 400);
    quota.setUserQuota(ids.mia, { mode: "default" }); // back to the $10 default: over
  });
});

describe("enforcement", () => {
  it("refuses a new pane of an owner over quota with EQUOTA and a reason", () => {
    assert.throws(
      () => quota.assertSpawnAllowed(ids.mia),
      (err) => {
        assert.equal(err.code, "EQUOTA");
        assert.match(err.message, /\$15\.00 of \$10\.00 per day/);
        return true;
      }
    );
    quota.assertSpawnAllowed(ids.lou); // under the default: no throw
    quota.assertSpawnAllowed(null); // a workspace with no owner
  });

  it("gives a pane its owner's flag file in server mode only", () => {
    const env = quota.flagEnv(ids.mia, null);
    assert.match(env.AGENTHUB_QUOTA_FILE, /user-quota[/\\]u\d+\.json$/);
    delete process.env.AGENTHUB_MODE;
    try {
      assert.deepEqual(quota.flagEnv(ids.mia, null), {});
      quota.assertSpawnAllowed(ids.mia); // local mode never refuses
      assert.equal(quota.statusAll().size, 0);
    } finally {
      process.env.AGENTHUB_MODE = "server";
    }
  });

  it("sweep: writes the flag, ends the panes no hook can stop, keeps Claude's, clears on raise", () => {
    const realList = runs.listRuns;
    const realKill = runs.killRun;
    const killed = [];
    runs.listRuns = () => [
      { id: `agenthub-lane-${laneId}`, laneId, status: "running", tool: "claude" },
      { id: `agenthub-lane-${laneId}-1`, laneId, status: "running", tool: "cursor" },
    ];
    runs.killRun = (id) => killed.push(id) > 0;
    try {
      quota.__reset();
      const r = quota.sweep();
      assert.deepEqual(r.blocked, [ids.mia]);
      assert.deepEqual(killed, [`agenthub-lane-${laneId}-1`]);
      const file = quota.flagEnv(ids.mia, null).AGENTHUB_QUOTA_FILE;
      const flag = JSON.parse(fs.readFileSync(file, "utf8"));
      assert.equal(flag.blocked, true);
      assert.match(flag.message, /quota is used up/);
      assert.ok(Date.parse(flag.until) > Date.now());
      // Lou never went over: no file is written for them.
      assert.equal(fs.existsSync(quota.flagEnv(ids.lou, null).AGENTHUB_QUOTA_FILE), false);
      assert.equal(quota.isRunBlocked(`agenthub-lane-${laneId}-3`), true);

      quota.setUserQuota(ids.mia, { mode: "unlimited" });
      killed.length = 0;
      quota.sweep();
      assert.deepEqual(killed, []);
      assert.equal(JSON.parse(fs.readFileSync(file, "utf8")).blocked, false);
      assert.equal(quota.isRunBlocked(`agenthub-lane-${laneId}`), false);
    } finally {
      runs.listRuns = realList;
      runs.killRun = realKill;
      quota.setUserQuota(ids.mia, { mode: "default" });
      quota.__reset();
    }
  });

  it("the omp gateway loses a blocked owner's key, never the queue's", () => {
    const { clientsFor } = require("../lib/omp-share");
    const t = { standby: ["s1"], used: { a: { hash: "ha" }, b: { hash: "hb" } } };
    const out = clientsFor(t, (id) => id === "a");
    assert.equal(out.a, undefined);
    assert.deepEqual(out.b, { tokenSha256: "hb" });
    assert.equal(Object.keys(out).filter((k) => k.startsWith("standby-")).length, 1);
  });
});

describe("hook handler", () => {
  const HANDLER = path.resolve(__dirname, "..", "..", "scripts", "hook-handler.js");
  const flag = path.join(TMP, "flag.json");
  const run = (args, file) =>
    spawnSync(process.execPath, [HANDLER, ...args], {
      input: "{}",
      env: { ...process.env, AGENTHUB_QUOTA_FILE: file, CLAUDE_DASHBOARD_PORT: "1" },
      encoding: "utf8",
    });

  it("refuses a Claude Code or Codex prompt while the flag says blocked", () => {
    fs.writeFileSync(
      flag,
      JSON.stringify({ blocked: true, until: "2999-01-01T00:00:00Z", message: "quota used up" })
    );
    for (const args of [["UserPromptSubmit"], ["UserPromptSubmit", "codex"]]) {
      const r = run(args, flag);
      assert.equal(r.status, 2);
      assert.match(r.stderr, /quota used up/);
    }
  });

  it("lets everything else through: other events, Antigravity, an expired or missing flag", () => {
    assert.equal(run(["PreToolUse"], flag).status, 0);
    assert.equal(run(["UserPromptSubmit", "antigravity"], flag).status, 0);
    assert.equal(run(["UserPromptSubmit"], path.join(TMP, "none.json")).status, 0);
    fs.writeFileSync(flag, JSON.stringify({ blocked: true, until: "2000-01-01T00:00:00Z" }));
    assert.equal(run(["UserPromptSubmit"], flag).status, 0);
    fs.writeFileSync(flag, "{not json");
    assert.equal(run(["UserPromptSubmit"], flag).status, 0);
  });
});