AgentHub/server/__tests__/claude-pool.test.js

/**
 * @file claude-pool.test.js
 * @description The Claude account pool (server/lib/claude-pool.js): token
 * validation and file modes, the balancer's order (full quota last, then fewest
 * live panes, then least recent spend, expired accounts never), and spawnRun
 * handing the pick to a Claude pane through tmux `-e CLAUDE_CODE_OAUTH_TOKEN`
 * and recording it — but not for a Codex pane, and not with an empty pool.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const { describe, it, before, after, beforeEach } = 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 TMP = fs.mkdtempSync(path.join(os.tmpdir(), "claude-pool-test-"));
process.env.DASHBOARD_DB_PATH = path.join(TMP, "dashboard.db");
process.env.CLAUDE_HOME = path.join(TMP, "claude");

const { db } = require("../db");
const ccUsage = require("../lib/cc-usage");
const pool = require("../lib/claude-pool");
const tmux = require("../lib/tmux");
const pty = require("../lib/pty-run");

const TOKEN = (c) => `sk-ant-oat01-${c.repeat(40)}`;
const realFetchUsage = ccUsage.fetchUsage;

function run(id, account, { status = "completed", session = null, started, ended = null } = {}) {
  db.prepare(
    `INSERT INTO dashboard_runs (id, session_id, mode, cwd, status, started_at, ended_at, account)
     VALUES (?, ?, 'interactive', '/tmp', ?, ?, ?, ?)`
  ).run(id, session, status, started || new Date().toISOString(), ended, account);
}

const MODEL = "claude-sonnet-4-5-20250929";
const hoursAgo = (h) => new Date(Date.now() - h * 3600e3).toISOString();

/**
 * `tokens` output tokens of `session` spent `atHoursAgo` hours ago — the
 * hourly row plus the session total, as the hook writer keeps them both.
 * `{hourly: false}` = a total with no hourly row (a sub-agent's bucket).
 */
function spend(session, tokens, { atHoursAgo = 0, hourly = true, model = MODEL } = {}) {
  db.prepare(
    "INSERT OR IGNORE INTO sessions (id, name, status, cwd, started_at) VALUES (?, ?, 'completed', '/tmp', ?)"
  ).run(session, session, hoursAgo(atHoursAgo));
  db.prepare(
    `INSERT INTO token_usage (session_id, model, output_tokens) VALUES (?, ?, ?)
     ON CONFLICT DO UPDATE SET output_tokens = output_tokens + excluded.output_tokens`
  ).run(session, model, tokens);
  if (!hourly) return;
  db.prepare(
    `INSERT INTO token_usage_hourly (session_id, hour, model, output_tokens) VALUES (?, ?, ?, ?)`
  ).run(session, hoursAgo(atHoursAgo).slice(0, 13), model, tokens);
}

before(() => {
  // No network: every quota read answers "no reading" unless a test says otherwise.
  ccUsage.fetchUsage = async () => {
    throw Object.assign(new Error("offline"), { code: "EUNAUTHORIZED" });
  };
});

after(() => {
  ccUsage.fetchUsage = realFetchUsage;
  tmux.__setExecImpl(null);
  fs.rmSync(TMP, { recursive: true, force: true });
});

describe("claude pool", () => {
  beforeEach(() => {
    fs.rmSync(pool.poolDir(), { recursive: true, force: true });
    db.prepare("DELETE FROM dashboard_runs").run();
    db.prepare("DELETE FROM token_usage_hourly").run();
    db.prepare("DELETE FROM token_usage").run();
    db.prepare("DELETE FROM sessions").run();
    pool.__quotaCache.clear();
  });

  it("is empty until an admin adds an account, and stores it 0600", () => {
    assert.equal(pool.pickAccount(), null);
    assert.throws(() => pool.addAccount("a", "not-a-token"), { code: "EBADTOKEN" });
    assert.throws(() => pool.addAccount("../x", TOKEN("a")), { code: "EBADNAME" });
    const added = pool.addAccount("a", TOKEN("a"));
    assert.ok(Date.parse(added.expiresAt) > Date.now() + 360 * 864e5);
    assert.throws(() => pool.addAccount("a", TOKEN("b")), { code: "EEXISTS" });
    const mode = fs.statSync(path.join(pool.poolDir(), "a.json")).mode & 0o777;
    assert.equal(mode, 0o600);
    assert.equal(fs.statSync(pool.poolDir()).mode & 0o777, 0o700);
    assert.ok(
      !JSON.stringify(pool.poolStatus()).includes(TOKEN("a")),
      "status never shows a token"
    );
    pool.removeAccount("a");
    assert.equal(pool.pickAccount(), null);
  });

  it("prefers fewer live panes, then less recent spend, and skips expired tokens", () => {
    for (const n of ["a", "b", "c"]) pool.addAccount(n, TOKEN(n));
    assert.equal(pool.pickAccount().name, "a");

    run("r1", "a", { status: "running" });
    assert.equal(pool.pickAccount().name, "b");

    run("r2", "b", { session: "s-b" });
    spend("s-b", 100000);
    const picked = pool.pickAccount();
    assert.equal(picked.name, "c");
    assert.equal(picked.token, TOKEN("c"));

    const status = Object.fromEntries(pool.poolStatus().map((a) => [a.name, a]));
    assert.equal(status.a.livePanes, 1);
    assert.ok(status.b.spend5h > 0 && status.b.spend7d >= status.b.spend5h);
    assert.equal(status.c.spend5h, 0);

    // c's token is a year old: never handed out.
    const file = path.join(pool.poolDir(), "c.json");
    const rec = JSON.parse(fs.readFileSync(file, "utf8"));
    rec.addedAt = new Date(Date.now() - 366 * 864e5).toISOString();
    fs.writeFileSync(file, JSON.stringify(rec));
    assert.equal(pool.pickAccount().name, "b");
  });

  it("books a session to an account only for the hours its runs held it, sub-agents included", () => {
    for (const n of ["a", "b"]) pool.addAccount(n, TOKEN(n));
    // s1 ran on a 30 h ago, then was resumed on b 2 h ago.
    run("r-a", "a", { session: "s1", started: hoursAgo(30), ended: hoursAgo(29) });
    run("r-b", "b", { session: "s1", started: hoursAgo(2), status: "running" });
    spend("s1", 1000, { atHoursAgo: 30 });
    spend("s1", 3000, { atHoursAgo: 2 });
    // s2 (on a, now): its main agent and a Haiku sub-agent with no hourly row.
    run("r-a2", "a", { session: "s2" });
    spend("s2", 500);
    spend("s2", 200, { hourly: false, model: "claude-haiku-4-5-20251001" });
    const status = Object.fromEntries(pool.poolStatus().map((x) => [x.name, x]));
    assert.equal(status.a.tokens5h, 700, "s2, sub-agent included; s1 is b's now");
    assert.equal(status.a.tokens7d, 1700);
    assert.equal(status.b.tokens5h, 3000, "not the hours a held s1");
    assert.equal(status.b.tokens7d, 3000);
    assert.ok(status.a.spend7d > status.a.spend5h && status.a.spend5h > 0);
    assert.ok(status.b.spend5h > 0);
  });

  it("puts an account whose quota is nearly used up last", async () => {
    for (const n of ["a", "b"]) pool.addAccount(n, TOKEN(n));
    ccUsage.fetchUsage = async (token) => ({
      fiveHour: { utilization: token === TOKEN("a") ? 95 : 10, resetsAt: null },
      sevenDay: { utilization: 20, resetsAt: null },
    });
    try {
      await pool.refreshQuotas({ force: true });
      run("r1", "b", { status: "running" }); // b busier, but a is full
      assert.equal(pool.pickAccount().name, "b");
      assert.equal(pool.poolStatus().find((x) => x.name === "a").full, true);
    } finally {
      ccUsage.fetchUsage = async () => {
        throw Object.assign(new Error("offline"), { code: "EUNAUTHORIZED" });
      };
    }
  });

  it("gives a Claude pane the picked account's token and records it", () => {
    pool.addAccount("a", TOKEN("a"));
    let argv = null;
    tmux.__setExecImpl((args) => {
      if (args[0] === "has-session") throw Object.assign(new Error("gone"), { status: 1 });
      if (args[0] === "new-session") argv = args;
      return "";
    });
    pty.spawnRun({ laneId: 7, cwd: "/tmp", tool: "claude" });
    assert.ok(argv.includes(`CLAUDE_CODE_OAUTH_TOKEN=${TOKEN("a")}`));
    assert.ok(argv.indexOf(`CLAUDE_CODE_OAUTH_TOKEN=${TOKEN("a")}`) < argv.indexOf("--"));
    const row = db.prepare("SELECT account FROM dashboard_runs WHERE id = 'agenthub-lane-7'").get();
    assert.equal(row.account, "a");

    argv = null;
    pty.spawnRun({ laneId: 8, cwd: "/tmp", tool: "codex" });
    assert.ok(!argv.some((x) => String(x).startsWith("CLAUDE_CODE_OAUTH_TOKEN=")));

    pool.removeAccount("a");
    argv = null;
    pty.spawnRun({ laneId: 9, cwd: "/tmp", tool: "claude" });
    assert.ok(!argv.some((x) => String(x).startsWith("CLAUDE_CODE_OAUTH_TOKEN=")));
  });
});