AgentHub/server/lib/codex-pool.js

/**
 * @file codex-pool.js
 * @description The Codex account POOL, modelled on claude-pool.js: accounts an
 * admin added from the host's shell (`agenthub accounts add <name> --tool
 * codex`), and the balancer that gives every new Codex pane on no AI Proxy one
 * of them. An empty pool changes nothing: panes run on their CODEX_HOME's own
 * login as before. Two kinds:
 *   - `apikey`: an OpenAI API key, any number of panes at once. The key reaches
 *     the pane only through tmux `-e AGENTHUB_POOL_KEY` (a member's launch
 *     file); the argv only adds a `-c` provider whose `env_key` NAMES it.
 *   - `chatgpt`: a ChatGPT login (`auth.json`). Codex reads it from no env
 *     var, so it is written into the MEMBER's own CODEX_HOME (S3) AS the member
 *     (write-as, 0600) just before the pane starts — never into a host pane's
 *     CODEX_HOME (that is the host's own login), never over a login AgentHub
 *     did not place there, and never a different account than the one a live
 *     pane of that member already runs on (one auth.json per home). Refresh
 *     tokens rotate, so `syncBack` copies a member's auth.json back into the
 *     pool whenever its `last_refresh` is newer (newest wins) — only from a
 *     home AgentHub placed that account in, never dated in the future, and
 *     only when its id_token names the same account (`mayCopyBack`); the 2026-10-08
 *     spike showed an old refresh token still refreshes, so panes sharing an
 *     account do not log each other out.
 *
 * Records live in `<CLAUDE_HOME>/agenthub-codex-accounts/<name>.json` (dir
 * 0700, file 0600); no route returns a key or a token.
 *
 * `pickAccount` ranks like the Claude pool: an account whose last reading is
 * ≥ `BLOCK_PERCENT` goes last, then fewest live Codex panes, then least $ in
 * the last 5 h, then 7 d. The reading is the newest `rate_limits` in the
 * rollouts of the sessions that account's runs started (no network), each
 * from the CODEX_HOME of the account its workspace runs as.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const fs = require("node:fs");
const path = require("node:path");
const { getClaudeHome } = require("./claude-home");
const { getCodexHome } = require("./codex-home");

const NAME_RE = /^[a-z0-9][a-z0-9._-]{0,63}$/i;
const KEY_RE = /^sk-[A-Za-z0-9_-]{20,500}$/;
const QUOTA_TTL_MS = 60 * 1000;
const BLOCK_PERCENT = 90;
const PROVIDER = "agenthub_pool";
const KEY_ENV = "AGENTHUB_POOL_KEY";
const SESSION_IN_NAME = /-([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})\.jsonl$/i;

const err = (code, message) => Object.assign(new Error(message), { code });
const runAs = () => require("./run-as");

const poolDir = () => path.join(getClaudeHome(), "agenthub-codex-accounts");
const fileOf = (name) => path.join(poolDir(), `${name}.json`);

function validName(name) {
  const n = String(name || "").trim();
  if (!NAME_RE.test(n))
    throw err("EBADNAME", "account name: 1–64 letters, digits, '.', '_' or '-'");
  return n;
}

const isChatgptAuth = (auth) =>
  ["id_token", "access_token", "refresh_token", "account_id"].every(
    (k) => typeof auth?.tokens?.[k] === "string" && auth.tokens[k]
  );

function readAccount(name) {
  try {
    const a = JSON.parse(fs.readFileSync(fileOf(name), "utf8"));
    if (a.kind === "apikey" && typeof a.key === "string") return a;
    if (a.kind === "chatgpt" && isChatgptAuth(a.auth)) return a;
  } catch {
    /* absent or not ours */
  }
  return null;
}

function allAccounts() {
  let names = [];
  try {
    names = fs.readdirSync(poolDir()).filter((f) => f.endsWith(".json"));
  } catch {
    return [];
  }
  return names
    .map((f) => f.slice(0, -5))
    .filter((n) => NAME_RE.test(n))
    .map(readAccount)
    .filter(Boolean)
    .sort((a, b) => a.name.localeCompare(b.name));
}

/** Write a record atomically, 0600. `create` refuses an existing name. */
function save(record, { create = false } = {}) {
  fs.mkdirSync(poolDir(), { recursive: true, mode: 0o700 });
  fs.chmodSync(poolDir(), 0o700);
  if (create) {
    fs.writeFileSync(fileOf(record.name), JSON.stringify(record, null, 2), {
      mode: 0o600,
      flag: "wx",
    });
    return;
  }
  // ponytail: no cross-process lock — the server's read-compare-write is one
  // synchronous run on its event loop, and the rename is atomic; a `remove`
  // from the CLI racing a copy-back can at worst bring the account back.
  if (!fs.existsSync(fileOf(record.name))) return; // removed meanwhile
  const tmp = `${fileOf(record.name)}.${process.pid}.tmp`;
  fs.writeFileSync(tmp, JSON.stringify(record, null, 2), { mode: 0o600 });
  fs.renameSync(tmp, fileOf(record.name));
}

/**
 * Add an account: `{apiKey}`, or `{auth}` = the parsed `auth.json` of a
 * `codex login` (a ChatGPT login, or one holding only an API key).
 */
function addAccount(name, { apiKey, auth } = {}) {
  const n = validName(name);
  let record;
  if (isChatgptAuth(auth)) {
    record = { name: n, kind: "chatgpt", accountId: auth.tokens.account_id, auth };
  } else {
    const key = String((auth && auth.OPENAI_API_KEY) || apiKey || "").trim();
    if (!KEY_RE.test(key)) {
      throw err(
        "EBADTOKEN",
        "not a Codex login — give an OpenAI API key (sk-…) or the auth.json `codex login` wrote"
      );
    }
    record = { name: n, kind: "apikey", key };
  }
  if (readAccount(n)) throw err("EEXISTS", `account ${n} is already in the pool`);
  record.addedAt = new Date().toISOString();
  save(record, { create: true });
  quotaCache = null;
  return { name: n, kind: record.kind, addedAt: record.addedAt };
}

function removeAccount(name) {
  const n = validName(name);
  if (!readAccount(n)) throw err("ENOACCOUNT", `no account ${n} in the pool`);
  fs.rmSync(fileOf(n));
  quotaCache = null;
}

// ── What AgentHub recorded per account ─────────────────────────────────────

function livePanes(account) {
  const { db } = require("../db");
  return db
    .prepare(
      `SELECT COUNT(*) AS n FROM dashboard_runs
        WHERE tool = 'codex' AND account = ? AND status IN ('running', 'spawning')`
    )
    .get(account).n;
}

/**
 * $ of the Codex sessions this account's runs started since `hours` ago.
 * ponytail: Codex usage has no hourly table, so a session counts whole once
 * its run started in the window; add hourly rows if that skews a ranking.
 */
function spendSince(account, hours) {
  const { db, stmts } = require("../db");
  const { calculateCost } = require("../routes/pricing");
  const since = new Date(Date.now() - hours * 3600e3).toISOString();
  const rows = db
    .prepare(
      `SELECT model, speed, inference_geo, service_tier,
              SUM(input_tokens) AS input_tokens, SUM(output_tokens) AS output_tokens,
              SUM(cache_read_tokens) AS cache_read_tokens,
              SUM(cache_write_tokens) AS cache_write_tokens
         FROM token_usage
        WHERE session_id IN (SELECT session_id FROM dashboard_runs
                              WHERE tool = 'codex' AND account = ?
                                AND session_id IS NOT NULL AND started_at >= ?)
        GROUP BY model, speed, inference_geo, service_tier`
    )
    .all(account, since);
  return rows.length ? calculateCost(rows, stmts.listPricing.all()).total_cost : 0;
}

// ── Quota, from the account's own rollouts ─────────────────────────────────

let quotaCache = null; // { at, byName: Map<name, {windows, asOf}> }

/**
 * name → the newest `rate_limits` its sessions' rollouts carry. A run's
 * rollout is read only from the CODEX_HOME of the account its workspace runs
 * as (a member's own, else the host's), never dated in the future — a member
 * can speak only for their own session's reading (an accepted ceiling: no
 * network check), never for another account's run.
 */
function quotas() {
  if (quotaCache && Date.now() - quotaCache.at < QUOTA_TTL_MS) return quotaCache.byName;
  const byName = new Map();
  const { db } = require("../db");
  const runs = db
    .prepare(
      `SELECT DISTINCT account, session_id, lane_id FROM dashboard_runs
        WHERE tool = 'codex' AND account IS NOT NULL AND session_id IS NOT NULL`
    )
    .all();
  if (runs.length) {
    const { allRolloutFiles } = require("./codex-ingest");
    const { lastRateLimits, codexWindow } = require("./quota");
    const homeOf = (laneId) => {
      try {
        const acct = runAs().accountForLane(laneId);
        return acct ? path.join(acct.home, ".codex") : getCodexHome();
      } catch {
        return null; // a gone member: nobody's reading
      }
    };
    // ponytail: one walk of each CODEX_HOME in use per minute; index rollouts
    // by session id if a sessions/ tree ever makes that slow.
    const byHome = new Map(); // home → session id → rollout files
    const filesIn = (home, sid) => {
      if (!byHome.has(home)) {
        const bySession = new Map();
        for (const f of allRolloutFiles(home)) {
          if (f.mtime > Date.now() + FUTURE_SLACK_MS) continue;
          const id = SESSION_IN_NAME.exec(path.basename(f.path))?.[1];
          if (id) bySession.set(id, [...(bySession.get(id) || []), f]);
        }
        byHome.set(home, bySession);
      }
      return byHome.get(home).get(sid) || [];
    };
    const filesOf = new Map();
    for (const r of runs) {
      const home = homeOf(r.lane_id);
      if (!home) continue;
      filesOf.set(r.account, [...(filesOf.get(r.account) || []), ...filesIn(home, r.session_id)]);
    }
    for (const [name, files] of filesOf) {
      for (const f of files.sort((a, b) => b.mtime - a.mtime)) {
        const found = lastRateLimits(f);
        if (!found) continue;
        const windows = [found.rl.primary, found.rl.secondary]
          .map((w) => codexWindow(w, found.mtime))
          .filter(Boolean);
        byName.set(name, { windows, asOf: new Date(found.mtime).toISOString() });
        break;
      }
    }
  }
  quotaCache = { at: Date.now(), byName };
  return byName;
}

function quotaPercent(name) {
  const q = quotas().get(name);
  return q && q.windows.length ? Math.max(...q.windows.map((w) => w.usedPercent)) : null;
}

/** Every account with what the balancer sees — never a key or a token. */
function poolStatus() {
  return allAccounts().map((a) => {
    const percent = quotaPercent(a.name);
    return {
      name: a.name,
      kind: a.kind,
      addedAt: a.addedAt,
      livePanes: livePanes(a.name),
      spend5h: spendSince(a.name, 5),
      spend7d: spendSince(a.name, 24 * 7),
      quota: quotas().get(a.name) || null,
      quotaPercent: percent,
      full: percent != null && percent >= BLOCK_PERCENT,
    };
  });
}

// ── A member's auth.json ───────────────────────────────────────────────────

const authFile = (home) => path.join(home, ".codex", "auth.json");

/** The auth.json in `home`'s CODEX_HOME (read guarded: the member's own file), or null. */
function homeAuth(home) {
  try {
    return JSON.parse(runAs().readFileSafe(authFile(home), "utf8", { max: 1 << 20 }));
  } catch {
    return null;
  }
}

const refreshedAt = (auth) => Date.parse(auth?.last_refresh) || 0;
const FUTURE_SLACK_MS = 5 * 60 * 1000;

// Who an id_token says it is: Codex's carries the ChatGPT account under this
// claim (and the email); a real refresh keeps both.
const AUTH_CLAIM = "https://api.openai.com/auth";
function identity(auth) {
  try {
    const c = JSON.parse(
      Buffer.from(String(auth.tokens.id_token).split(".")[1], "base64url").toString("utf8")
    );
    return { email: c.email ?? null, account: c[AUTH_CLAIM]?.chatgpt_account_id ?? null };
  } catch {
    return null;
  }
}

/**
 * Whether `disk` (a member's auth.json) may replace pool account `a`: AgentHub
 * placed `a` in that home itself (`placedIn`), its `last_refresh` is no later
 * than now, and its id_token names the same account and email as ours. A
 * member edits their own auth.json at will; without these, one member could
 * write a forged record over a login every other member's panes get.
 * ponytail: the id_token is not signature-checked (no network here), so a
 * member P was placed with can still poison P with a self-consistent forgery
 * until a real refresh from another home wins it back; a server-side refresh
 * is the upgrade.
 */
function mayCopyBack(a, disk, login) {
  if (!placedHere(a, disk, login)) return false;
  const t = refreshedAt(disk);
  return t > refreshedAt(a.auth) && t <= Date.now() + FUTURE_SLACK_MS;
}

/**
 * Whether `disk` (a member's auth.json) is pool account `a` that AgentHub
 * placed in that home: `placedIn`, and the same email and account in the
 * id_token. `account_id` alone says nothing — it is the ChatGPT WORKSPACE,
 * shared by every seat of a Team, so a member's own seat would match it.
 */
function placedHere(a, disk, login) {
  if (a.kind !== "chatgpt" || !(a.placedIn || []).includes(login)) return false;
  const ours = identity(a.auth);
  const theirs = identity(disk);
  return Boolean(ours && theirs && ours.email === theirs.email && ours.account === theirs.account);
}

/** Copy `home`'s auth.json back into the pool account AgentHub placed there, when it is a newer login of it. */
function syncHome(home) {
  const disk = homeAuth(home);
  if (!isChatgptAuth(disk)) return;
  const login = require("./linux-accounts").loginFor(path.basename(home));
  for (const a of allAccounts()) {
    if (a.kind !== "chatgpt" || a.accountId !== disk.tokens.account_id) continue;
    if (mayCopyBack(a, disk, login)) save({ ...a, auth: disk });
  }
}

/** Every member's CODEX_HOME → the pool, newest login wins (the run-heal beat). */
function syncBack() {
  if (!allAccounts().some((a) => a.kind === "chatgpt")) return;
  for (const home of runAs().memberHomes()) syncHome(home);
}

// ── Picking ────────────────────────────────────────────────────────────────

/** The full record a new Codex pane should run on, or null. `member` = run-as account. */
function pick(member) {
  let usable = allAccounts();
  if (usable.length === 0) return null;
  if (!member) {
    // A host pane's CODEX_HOME holds the host's own login.
    usable = usable.filter((a) => a.kind === "apikey");
  } else {
    const disk = homeAuth(member.home);
    const placed = disk ? usable.find((a) => placedHere(a, disk, member.login)) : null;
    // A login AgentHub did not place there is the member's own: keep it.
    if (!placed && (isChatgptAuth(disk) || disk?.OPENAI_API_KEY)) {
      usable = usable.filter((a) => a.kind !== "chatgpt");
    } else {
      // One auth.json per home: while a pane of theirs runs on a pool
      // login, the next pane stays on that one.
      const live = usable.find((a) => a.kind === "chatgpt" && liveInHome(a.name, member));
      if (live) return live;
    }
  }
  if (usable.length === 0) return null;
  const ranked = usable
    .map((a) => {
      const percent = quotaPercent(a.name);
      return {
        a,
        key: [
          percent != null && percent >= BLOCK_PERCENT ? 1 : 0,
          livePanes(a.name),
          spendSince(a.name, 5),
          spendSince(a.name, 24 * 7),
        ],
      };
    })
    .sort((x, y) => {
      for (let i = 0; i < x.key.length; i++) {
        if (x.key[i] !== y.key[i]) return x.key[i] - y.key[i];
      }
      return x.a.name.localeCompare(y.a.name);
    });
  return ranked[0].a;
}

function liveInHome(name, member) {
  const { db } = require("../db");
  return db
    .prepare(
      `SELECT lane_id FROM dashboard_runs
        WHERE tool = 'codex' AND account = ? AND status IN ('running', 'spawning')
          AND lane_id IS NOT NULL`
    )
    .all(name)
    .some((r) => {
      try {
        return runAs().accountForLane(r.lane_id)?.login === member.login;
      } catch {
        return false; // a gone member's lane is not this (enabled) member's
      }
    });
}

/** `{name, kind}` of the account a new Codex pane would get, or null. */
function pickAccount(member = null) {
  const a = pick(member);
  return a ? { name: a.name, kind: a.kind } : null;
}

/**
 * What a new Codex pane needs from the pool, or null (empty pool, or nothing
 * this pane may use): `{name, argv, env, place}`. `env` goes through tmux `-e`
 * (a member's launch file); `place()` runs just before the pane starts.
 */
function forSpawn(member = null) {
  const a = pick(member);
  if (!a) return null;
  if (a.kind === "apikey") {
    const provider = `{name="OpenAI (AgentHub pool)",base_url="https://api.openai.com/v1",env_key="${KEY_ENV}",wire_api="responses"}`;
    return {
      name: a.name,
      argv: ["-c", `model_providers.${PROVIDER}=${provider}`, "-c", `model_provider="${PROVIDER}"`],
      env: { [KEY_ENV]: a.key },
      place: () => {},
    };
  }
  return {
    name: a.name,
    argv: [],
    env: {},
    place: () => {
      syncHome(member.home); // a live pane may hold a newer login than ours
      const record = readAccount(a.name) || a;
      runAs().mkdirAs(member, path.dirname(authFile(member.home)));
      runAs().writeAs(
        member,
        authFile(member.home),
        "private",
        JSON.stringify(record.auth, null, 2)
      );
      // Only a home AgentHub placed this login in may copy a refresh back.
      if (!(record.placedIn || []).includes(member.login)) {
        save({ ...record, placedIn: [...(record.placedIn || []), member.login] });
      }
    },
  };
}

module.exports = {
  BLOCK_PERCENT,
  addAccount,
  removeAccount,
  poolStatus,
  pickAccount,
  forSpawn,
  syncBack,
  poolDir,
  __read: readAccount,
  resetQuota: () => {
    quotaCache = null;
  },
};