/** * @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ĩ */ 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="))); }); });