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

/**
 * @file codex-pool.test.js
 * @description The Codex account pool (server/lib/codex-pool.js), fake
 * accounts only: key/auth.json validation and file modes, the balancer
 * (full quota last — read from the account's own rollouts —, then fewest live
 * Codex panes, then spend), an API-key account reaching a pane only through
 * tmux `-e` with a `-c` provider naming the variable, a ChatGPT account's
 * `auth.json` written AS the member into their CODEX_HOME (never a host pane's,
 * never over the member's own login, sticky while a pane of theirs runs on
 * one), the newest `auth.json` copied back into the pool, and an empty pool
 * leaving every spawn as it was.
 * @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.realpathSync(fs.mkdtempSync(path.join(os.tmpdir(), "codex-pool-test-")));
process.env.DASHBOARD_DB_PATH = path.join(TMP, "dashboard.db");
process.env.CLAUDE_HOME = path.join(TMP, "claude");
process.env.CODEX_HOME = path.join(TMP, "codex");

const { db } = require("../db");
const pool = require("../lib/codex-pool");
const tmux = require("../lib/tmux");
const runAs = require("../lib/run-as");
const pty = require("../lib/pty-run");
const { fakeSudo } = require("./fixtures/fake-sudo");

const KEY = (c) => `sk-proj-${c.repeat(40)}`;
const jwt = (claims) =>
  `e30.${Buffer.from(JSON.stringify(claims)).toString("base64url")}.sig-${Math.random()}`;
const AUTH = (accountId, lastRefresh = "2026-10-01T00:00:00Z", rt = `rt-${accountId}`) => ({
  auth_mode: "chatgpt",
  OPENAI_API_KEY: null,
  tokens: {
    id_token: jwt({ email: `${accountId}@example.com` }),
    access_token: jwt({ exp: 2e9 }),
    refresh_token: rt,
    account_id: accountId,
  },
  last_refresh: lastRefresh,
});

const MIA = { name: "mia", login: "ah-mia", home: path.join(TMP, "users", "mia"), uid: 5001 };
const MIA_AUTH = path.join(MIA.home, ".codex", "auth.json");
const realAccountForLane = runAs.accountForLane;
const realAccountForPath = runAs.accountForPath;
const realMemberHomes = runAs.memberHomes;
let sudo;
let tmuxArgv; // the last new-session argv

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

function rollout(sessionId, usedPercent, home = process.env.CODEX_HOME) {
  const dir = path.join(home, "sessions", "2026", "10", "09");
  fs.mkdirSync(dir, { recursive: true });
  const line = {
    type: "event_msg",
    payload: {
      type: "token_count",
      rate_limits: {
        primary: { used_percent: usedPercent, window_minutes: 300, resets_at: 4e9 },
        secondary: { used_percent: 5, window_minutes: 10080, resets_at: 4e9 },
      },
    },
  };
  const file = path.join(dir, `rollout-2026-10-09T00-00-00-${sessionId}.jsonl`);
  fs.writeFileSync(
    file,
    `${JSON.stringify({ type: "session_meta", payload: { id: sessionId } })}\n${JSON.stringify(line)}\n`
  );
  return file;
}

/** Record that AgentHub placed pool account `name` in mia's home (what place() does). */
function placedWithMia(name) {
  const file = path.join(pool.poolDir(), `${name}.json`);
  fs.writeFileSync(file, JSON.stringify({ ...pool.__read(name), placedIn: ["ah-mia"] }), {
    mode: 0o600,
  });
}

const envOf = (argv, name) =>
  argv
    .slice(0, argv.indexOf("--"))
    .find((x) => String(x).startsWith(`${name}=`))
    ?.slice(name.length + 1);

before(() => {
  tmux.__setExecImpl((args) => {
    if (args[0] === "has-session") throw Object.assign(new Error("gone"), { status: 1 });
    if (args[0] === "new-session") tmuxArgv = args;
    return "";
  });
});

after(() => {
  tmux.__setExecImpl(null);
  runAs.__setExecImpl(null);
  runAs.accountForLane = realAccountForLane;
  runAs.accountForPath = realAccountForPath;
  runAs.memberHomes = realMemberHomes;
  fs.rmSync(TMP, { recursive: true, force: true });
});

beforeEach(() => {
  fs.rmSync(pool.poolDir(), { recursive: true, force: true });
  fs.rmSync(path.join(TMP, "users"), { recursive: true, force: true });
  fs.rmSync(path.join(process.env.CODEX_HOME, "sessions"), { recursive: true, force: true });
  db.prepare("DELETE FROM dashboard_runs").run();
  pool.resetQuota();
  tmuxArgv = null;
  sudo = fakeSudo();
  runAs.__setExecImpl(sudo.exec);
  runAs.accountForLane = (laneId) => (laneId >= 50 ? MIA : null);
  // ...and her folder is hers (a pane's cwd must be its workspace's account's).
  runAs.accountForPath = (p) => (String(p).startsWith(MIA.home) ? MIA : realAccountForPath(p));
  runAs.memberHomes = () => [MIA.home];
  fs.mkdirSync(path.join(MIA.home, ".agenthub", "handoff"), { recursive: true });
});

describe("codex pool: accounts", () => {
  it("is empty until an admin adds one, and stores it 0600 in a 0700 dir", () => {
    assert.equal(pool.pickAccount(null), null);
    assert.throws(() => pool.addAccount("a", { apiKey: "nope" }), { code: "EBADTOKEN" });
    assert.throws(() => pool.addAccount("../x", { apiKey: KEY("a") }), { code: "EBADNAME" });
    assert.throws(() => pool.addAccount("c", { auth: { tokens: {} } }), { code: "EBADTOKEN" });
    assert.equal(pool.addAccount("a", { apiKey: KEY("a") }).kind, "apikey");
    assert.equal(pool.addAccount("c", { auth: AUTH("acct-c") }).kind, "chatgpt");
    // An auth.json holding only an API key is an API-key account.
    assert.equal(pool.addAccount("k", { auth: { OPENAI_API_KEY: KEY("k") } }).kind, "apikey");
    assert.throws(() => pool.addAccount("a", { apiKey: KEY("b") }), { code: "EEXISTS" });
    assert.equal(fs.statSync(path.join(pool.poolDir(), "a.json")).mode & 0o777, 0o600);
    assert.equal(fs.statSync(pool.poolDir()).mode & 0o777, 0o700);
    const shown = JSON.stringify(pool.poolStatus());
    assert.ok(
      !shown.includes(KEY("a")) && !shown.includes("rt-acct-c"),
      "status never shows a secret"
    );
    pool.removeAccount("a");
    assert.throws(() => pool.removeAccount("a"), { code: "ENOACCOUNT" });
  });

  it("ranks by Codex live panes (a Claude run of the same name does not count), then quota", () => {
    for (const n of ["a", "b"]) pool.addAccount(n, { apiKey: KEY(n) });
    run("c1", "a", { status: "running", tool: "claude" });
    assert.equal(pool.pickAccount(null).name, "a");
    run("x1", "a", { status: "running" });
    assert.equal(pool.pickAccount(null).name, "b");

    // b's own rollout says its 5-hour window is 95% used: b goes last.
    run("x2", "b", { session: "11111111-2222-3333-4444-555555555555" });
    rollout("11111111-2222-3333-4444-555555555555", 95);
    pool.resetQuota();
    assert.equal(pool.pickAccount(null).name, "a");
    const b = pool.poolStatus().find((x) => x.name === "b");
    assert.equal(b.full, true);
    assert.equal(b.quota.windows[0].usedPercent, 95);
  });

  it("reads a run's quota only from its own account's CODEX_HOME, never dated in the future", () => {
    pool.addAccount("b", { apiKey: KEY("b") });
    const MIA_CODEX = path.join(MIA.home, ".codex");
    // A host run's session id, forged in mia's home: not the host's reading.
    run("h1", "b", { session: "11111111-2222-3333-4444-000000000001" });
    rollout("11111111-2222-3333-4444-000000000001", 95, MIA_CODEX);
    pool.resetQuota();
    assert.equal(pool.poolStatus().find((x) => x.name === "b").quota, null);
    // Mia's own run, but her rollout is dated 2030 to stay the newest forever.
    run("m1", "b", { session: "11111111-2222-3333-4444-000000000002", lane: 50 });
    const future = rollout("11111111-2222-3333-4444-000000000002", 0, MIA_CODEX);
    fs.utimesSync(future, new Date("2030-01-01"), new Date("2030-01-01"));
    pool.resetQuota();
    assert.equal(pool.poolStatus().find((x) => x.name === "b").quota, null);
    // Her own current reading of her own run still counts.
    fs.utimesSync(future, new Date(), new Date());
    pool.resetQuota();
    assert.equal(pool.poolStatus().find((x) => x.name === "b").quotaPercent, 5);
  });

  it("never gives a host pane a ChatGPT account (it would replace the host's login)", () => {
    pool.addAccount("c", { auth: AUTH("acct-c") });
    assert.equal(pool.pickAccount(null), null);
    assert.equal(pool.pickAccount(MIA).name, "c");
  });
});

describe("codex pool: spawn", () => {
  it("an empty pool leaves a Codex spawn exactly as before", () => {
    pty.spawnRun({ laneId: 7, cwd: "/tmp", tool: "codex" });
    const plain = tmuxArgv;
    assert.ok(!plain.some((x) => /AGENTHUB_POOL_KEY|agenthub_pool/.test(String(x))));
    const row = db.prepare("SELECT account FROM dashboard_runs WHERE id = 'agenthub-lane-7'").get();
    assert.equal(row.account, null);
  });

  it("an API-key account reaches a host pane through tmux -e only", () => {
    pool.addAccount("a", { apiKey: KEY("a") });
    pty.spawnRun({ laneId: 7, cwd: "/tmp", tool: "codex", initialPrompt: "hi" });
    assert.equal(envOf(tmuxArgv, "AGENTHUB_POOL_KEY"), KEY("a"));
    const after = tmuxArgv.slice(tmuxArgv.indexOf("--"));
    assert.ok(!after.some((x) => String(x).includes(KEY("a"))), "never in the pane's argv");
    assert.ok(after.includes('model_provider="agenthub_pool"'));
    assert.ok(after.some((x) => String(x).includes('env_key="AGENTHUB_POOL_KEY"')));
    assert.ok(after.indexOf('model_provider="agenthub_pool"') < after.indexOf("hi"));
    const row = db.prepare("SELECT account FROM dashboard_runs WHERE id = 'agenthub-lane-7'").get();
    assert.equal(row.account, "a");

    // Claude panes and Codex panes on an AI Proxy keep their own way in.
    tmuxArgv = null;
    pty.spawnRun({ laneId: 8, cwd: "/tmp", tool: "claude" });
    assert.ok(!tmuxArgv.some((x) => String(x).includes("AGENTHUB_POOL_KEY")));
  });

  it("writes a ChatGPT account's auth.json AS the member, 0600, nothing secret over sudo", () => {
    pool.addAccount("c", { auth: AUTH("acct-c") });
    pty.spawnRun({ laneId: 50, cwd: path.join(MIA.home, "p"), tool: "codex" });
    const written = JSON.parse(fs.readFileSync(MIA_AUTH, "utf8"));
    assert.equal(written.tokens.refresh_token, "rt-acct-c");
    assert.equal(fs.statSync(MIA_AUTH).mode & 0o777, 0o600);
    assert.ok(
      sudo.calls.some((a) => a.includes(MIA_AUTH)),
      "written through write-as"
    );
    assert.ok(!JSON.stringify(sudo.calls).includes("rt-acct-c"));
    assert.ok(!tmuxArgv.some((x) => String(x).includes("rt-acct-c")));
    const row = db
      .prepare("SELECT account FROM dashboard_runs WHERE id = 'agenthub-lane-50'")
      .get();
    assert.equal(row.account, "c");
  });

  it("never replaces a member's own login, and stays on the account their live pane uses", () => {
    pool.addAccount("c", { auth: AUTH("acct-c") });
    pool.addAccount("d", { auth: AUTH("acct-d") });
    fs.mkdirSync(path.dirname(MIA_AUTH), { recursive: true });
    fs.writeFileSync(MIA_AUTH, JSON.stringify(AUTH("mias-own")));
    assert.equal(pool.pickAccount(MIA), null);
    pty.spawnRun({ laneId: 50, cwd: path.join(MIA.home, "p"), tool: "codex" });
    assert.equal(JSON.parse(fs.readFileSync(MIA_AUTH, "utf8")).tokens.account_id, "mias-own");

    // A pool login in place, its pane live on d: the next pane is d too, though
    // d already has more live panes than c.
    fs.writeFileSync(MIA_AUTH, JSON.stringify(AUTH("acct-d")));
    placedWithMia("d");
    run("agenthub-lane-51", "d", { status: "running", lane: 51 });
    assert.equal(pool.pickAccount(MIA).name, "d");
    db.prepare("DELETE FROM dashboard_runs").run();
    assert.equal(pool.pickAccount(MIA).name, "c");
  });

  it("a member's own seat in a pool account's ChatGPT Team stays theirs", () => {
    // Codex's account_id is the WORKSPACE: every seat of a Team shares it.
    pool.addAccount("c", { auth: AUTH("team-x") });
    placedWithMia("c"); // even in a home AgentHub once placed c in
    const own = AUTH("team-x", "2026-10-02T00:00:00Z", "rt-mias-seat");
    own.tokens.id_token = jwt({ email: "mia@example.com" });
    fs.mkdirSync(path.dirname(MIA_AUTH), { recursive: true });
    fs.writeFileSync(MIA_AUTH, JSON.stringify(own));
    assert.equal(pool.pickAccount(MIA), null);
    pty.spawnRun({ laneId: 50, cwd: path.join(MIA.home, "p"), tool: "codex" });
    assert.equal(JSON.parse(fs.readFileSync(MIA_AUTH, "utf8")).tokens.refresh_token, "rt-mias-seat");
  });

  it("copies the newest auth.json back into the pool", () => {
    pool.addAccount("c", { auth: AUTH("acct-c", "2026-10-01T00:00:00Z") });
    // A refresh comes back only from a home AgentHub placed the account in.
    fs.mkdirSync(path.dirname(MIA_AUTH), { recursive: true });
    fs.writeFileSync(MIA_AUTH, JSON.stringify(AUTH("acct-c", "2026-10-05T00:00:00Z", "rt-new")));
    pool.syncBack();
    assert.equal(pool.__read("c").auth.tokens.refresh_token, "rt-acct-c", "never placed there");
    fs.rmSync(MIA_AUTH);
    pty.spawnRun({ laneId: 50, cwd: path.join(MIA.home, "p"), tool: "codex" });
    assert.deepEqual(pool.__read("c").placedIn, ["ah-mia"]);
    fs.writeFileSync(MIA_AUTH, JSON.stringify(AUTH("acct-c", "2026-09-01T00:00:00Z", "rt-old")));
    pool.syncBack();
    assert.equal(pool.__read("c").auth.tokens.refresh_token, "rt-acct-c", "older: kept ours");
    fs.writeFileSync(MIA_AUTH, JSON.stringify(AUTH("acct-c", "2026-10-05T00:00:00Z", "rt-new")));
    pool.syncBack();
    assert.equal(pool.__read("c").auth.tokens.refresh_token, "rt-new");
    assert.equal(fs.statSync(path.join(pool.poolDir(), "c.json")).mode & 0o777, 0o600);
    // A new pane starts from the newest copy.
    pty.spawnRun({ laneId: 50, cwd: path.join(MIA.home, "p"), tool: "codex" });
    assert.equal(JSON.parse(fs.readFileSync(MIA_AUTH, "utf8")).tokens.refresh_token, "rt-new");
  });

  it("never takes back a forged auth.json: a future date, or another account's id_token", () => {
    pool.addAccount("c", { auth: AUTH("acct-c", "2026-10-01T00:00:00Z") });
    pty.spawnRun({ laneId: 50, cwd: path.join(MIA.home, "p"), tool: "codex" });
    const forged = AUTH("acct-c", "9999-12-31T00:00:00Z", "rt-forged");
    fs.writeFileSync(MIA_AUTH, JSON.stringify(forged));
    pool.syncBack();
    assert.equal(pool.__read("c").auth.tokens.refresh_token, "rt-acct-c", "future last_refresh");
    // mia's own ChatGPT tokens relabelled with the pool account's id.
    const own = AUTH("mias-own", "2026-10-05T00:00:00Z", "rt-mia");
    own.tokens.account_id = "acct-c";
    fs.writeFileSync(MIA_AUTH, JSON.stringify(own));
    pool.syncBack();
    assert.equal(pool.__read("c").auth.tokens.refresh_token, "rt-acct-c", "another identity");
    // A garbage id_token proves nothing either.
    own.tokens.id_token = "x";
    fs.writeFileSync(MIA_AUTH, JSON.stringify(own));
    pool.syncBack();
    assert.equal(pool.__read("c").auth.tokens.refresh_token, "rt-acct-c", "no identity");
  });
});