AgentHub/server/lib/claude-pool.js

/**
 * @file claude-pool.js
 * @description The Claude account POOL: subscription accounts an admin added
 * (`agenthub accounts add`, a `claude setup-token` token each), and the
 * balancer that gives every new Claude pane one of them through
 * `CLAUDE_CODE_OAUTH_TOKEN` (pty-run.js) — so sessions spread over the
 * accounts without touching the machine's own login. An empty pool changes
 * nothing: panes run on the machine's login as before.
 *
 * Tokens live in `<CLAUDE_HOME>/agenthub-accounts/<name>.json` (dir 0700, file
 * 0600); no route returns one. A setup-token is valid for one year from when it
 * was added; past that the account is skipped.
 *
 * `pickAccount` ranks the accounts by what AgentHub itself recorded:
 *   1. an account whose last quota reading is ≥ `BLOCK_PERCENT` in either
 *      window goes last (still used when every account is that full);
 *   2. fewest live panes (spreads concurrent sessions);
 *   3. least $ spent in the last 5 hours, then in the last 7 days — the
 *      sessions its runs (`dashboard_runs.account`) held, over the hours they
 *      held them (pool-usage.js).
 * The quota is read with the account's own token in the background, at most
 * every `QUOTA_TTL_MS`; a token the usage endpoint refuses simply has no
 * reading, and ranking falls back to spend.
 *
 * Two kinds of account:
 *   - `token`: a `claude setup-token` token, given to the pane as
 *     `CLAUDE_CODE_OAUTH_TOKEN` (any pane, host or member).
 *   - `login`: a real Claude login (`claudeAiOauth` + `oauthAccount`), moved
 *     in from a saved switch-user profile (cc-profiles.js) or pasted from a
 *     `.credentials.json`. Claude Code reads it from no env var, so — like a
 *     Codex ChatGPT login — it is written as `~/.claude/.credentials.json`
 *     into a MEMBER's own home (S3), as the member (write-as, 0600), just
 *     before the pane starts: never into a host pane's home (the host's own
 *     login — unless that login itself was moved in, `hostHolds`, and the
 *     host's home is then one more placement), never over a login AgentHub
 *     did not place, never a second
 *     account into a home where a live pane of that member runs on one.
 *     Its refresh token is SINGLE USE, so the server is the one that
 *     refreshes it (`refreshLogins`, `REFRESH_AHEAD_MS` before expiry) and
 *     pushes the new pair into every home it placed the login in — Claude
 *     Code re-reads its credentials before refreshing and takes tokens
 *     "another process landed". A home whose Claude refreshed first is copied
 *     back, but only after Anthropic confirms the token is the same account
 *     (`fetchOwner`): a member could otherwise plant a login every other
 *     member's panes would get.
 * @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 NAME_RE = /^[a-z0-9][a-z0-9._-]{0,63}$/i;
const TOKEN_RE = /^sk-ant-[A-Za-z0-9_-]{20,500}$/;
const TOKEN_LIFE_MS = 365 * 864e5;
const QUOTA_TTL_MS = 5 * 60 * 1000;
const QUOTA_FRESH_MS = 30 * 60 * 1000;
const BLOCK_PERCENT = 90;
const REFRESH_AHEAD_MS = 30 * 60 * 1000;
const FUTURE_SLACK_MS = 5 * 60 * 1000;

const err = (code, message) => Object.assign(new Error(message), { code });

const poolDir = () => path.join(getClaudeHome(), "agenthub-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 isOauth = (o) =>
  Boolean(o) &&
  typeof o.accessToken === "string" &&
  o.accessToken.length > 0 &&
  typeof o.refreshToken === "string" &&
  o.refreshToken.length > 0;

function readAccount(name) {
  try {
    const a = JSON.parse(fs.readFileSync(fileOf(name), "utf8"));
    if (a.kind === "login") return isOauth(a.oauth) ? a : null;
    return typeof a.token === "string" ? { ...a, kind: "token" } : null;
  } catch {
    return null;
  }
}

/** Rewrite one record (0600, through a temp file). */
function save(record) {
  const tmp = `${fileOf(record.name)}.tmp`;
  fs.writeFileSync(tmp, JSON.stringify(record, null, 2), { mode: 0o600 });
  fs.renameSync(tmp, fileOf(record.name));
}

/** The secret the usage endpoint takes for `a`. */
const bearerOf = (a) => (a.kind === "login" ? a.oauth.accessToken : a.token);

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));
}

// A setup-token lives a year; a login lives as long as its refresh token works.
const expiresAt = (a) =>
  a.kind === "login" ? null : new Date(Date.parse(a.addedAt) + TOKEN_LIFE_MS).toISOString();
const expired = (a) =>
  a.kind === "login" ? a.refreshError === "ESTALE" : Date.parse(expiresAt(a)) <= Date.now();

/** Add an account; refuses a name already in the pool. */
function addAccount(name, token) {
  const n = validName(name);
  const t = String(token || "").trim();
  if (!TOKEN_RE.test(t)) {
    throw err(
      "EBADTOKEN",
      "not a Claude token — paste what `claude setup-token` printed (sk-ant-…)"
    );
  }
  if (readAccount(n)) throw err("EEXISTS", `account ${n} is already in the pool`);
  fs.mkdirSync(poolDir(), { recursive: true, mode: 0o700 });
  fs.chmodSync(poolDir(), 0o700);
  const record = { name: n, token: t, addedAt: new Date().toISOString() };
  fs.writeFileSync(fileOf(n), JSON.stringify(record, null, 2), { mode: 0o600, flag: "wx" });
  quotaCache.delete(n);
  return { name: n, addedAt: record.addedAt, expiresAt: expiresAt(record) };
}

/**
 * Add a real login (`claudeAiOauth`, plus the `oauthAccount`/`userID` Claude
 * Code keeps in `.claude.json`); refuses a name, or an account, already in the
 * pool. The caller hands over its copy: two holders of one refresh token log
 * each other out, so a source that keeps using it must not be offered — with
 * one exception, `hostHolds`: the login this machine's own Claude Code runs
 * on, whose home then becomes one more placement (`HOST`) the server
 * refreshes and pushes into like any member's.
 */
function addLoginAccount(name, { oauth, account = null, userID = null, hostHolds = false } = {}) {
  const n = validName(name);
  if (!isOauth(oauth)) {
    throw err(
      "EBADTOKEN",
      "not a Claude login — it needs claudeAiOauth with an accessToken and a refreshToken"
    );
  }
  if (readAccount(n)) throw err("EEXISTS", `account ${n} is already in the pool`);
  const uuid = account?.accountUuid;
  const twin =
    uuid && allAccounts().find((a) => a.kind === "login" && a.account?.accountUuid === uuid);
  if (twin) throw err("EEXISTS", `that Claude account is already in the pool as ${twin.name}`);
  fs.mkdirSync(poolDir(), { recursive: true, mode: 0o700 });
  fs.chmodSync(poolDir(), 0o700);
  const record = {
    name: n,
    kind: "login",
    oauth,
    account,
    userID,
    addedAt: new Date().toISOString(),
    placed: hostHolds ? { [HOST]: oauth.refreshToken } : {},
  };
  fs.writeFileSync(fileOf(n), JSON.stringify(record, null, 2), { mode: 0o600, flag: "wx" });
  quotaCache.delete(n);
  return { name: n, kind: "login", addedAt: record.addedAt, email: account?.emailAddress ?? null };
}

function removeAccount(name) {
  const n = validName(name);
  const a = readAccount(n);
  if (!a) throw err("ENOACCOUNT", `no account ${n} in the pool`);
  // Take the login back out of every home it is still ours in.
  if (a.kind === "login") {
    for (const [member, acct] of placedHomes(a)) {
      // The host keeps its own login: removing it would log this machine out.
      if (acct.host || homeCreds(acct)?.refreshToken !== a.placed[member]) continue;
      try {
        runAs().writeAs(acct, credsPath(acct), "remove");
      } catch {
        /* best effort: the member can still /logout */
      }
    }
  }
  fs.rmSync(fileOf(n));
  quotaCache.delete(n);
}

// ── A member's .credentials.json (kind `login`) ────────────────────────────

const runAs = () => require("./run-as");
/** The placement key of this machine's own Claude Code (`hostHolds`). */
const HOST = "@host";
const HOST_ACCT = { name: HOST, host: true };
const credsPath = (acct) =>
  acct.host
    ? path.join(getClaudeHome(), ".credentials.json")
    : path.join(acct.home, ".claude", ".credentials.json");

/** The whole `.credentials.json` of a placement (a member's read guarded), or null. */
function readCredsFile(acct) {
  try {
    const text = acct.host
      ? fs.readFileSync(credsPath(acct), "utf8")
      : runAs().readFileSafe(credsPath(acct), "utf8", { max: 1 << 20 });
    const c = JSON.parse(text);
    return c && typeof c === "object" && !Array.isArray(c) ? c : null;
  } catch {
    return null;
  }
}

/** The `claudeAiOauth` of a placement, or null. */
function homeCreds(acct) {
  const c = readCredsFile(acct);
  return c && typeof c.claudeAiOauth === "object" ? c.claudeAiOauth : null;
}

/** `[key, account]` for every home `a` is placed in that still runs as its member (or the host). */
function placedHomes(a) {
  const out = [];
  for (const member of Object.keys(a.placed || {})) {
    if (member === HOST) {
      out.push([HOST, HOST_ACCT]);
      continue;
    }
    try {
      const acct = runAs().accountForUser(member);
      if (acct) out.push([member, acct]);
    } catch {
      /* a disabled member's home is left as it is */
    }
  }
  return out;
}

/** Whether `member`'s credentials are still the pair AgentHub last wrote there for `a`. */
const ownsHome = (a, member, disk) =>
  a.kind === "login" && Boolean(disk) && disk.refreshToken === a.placed?.[member];

/**
 * Write `a`'s login into a placement — a member's home as the member, the
 * host's directly (0600, through a temp file) — keeping every other key of
 * the file (`mcpOAuth`…), and remember the pair.
 */
function writeLogin(a, acct) {
  const text = `${JSON.stringify({ ...(readCredsFile(acct) || {}), claudeAiOauth: a.oauth }, null, 2)}\n`;
  if (acct.host) {
    const tmp = `${credsPath(acct)}.agenthub-tmp`;
    fs.writeFileSync(tmp, text, { mode: 0o600 });
    fs.renameSync(tmp, credsPath(acct));
  } else {
    runAs().mkdirAs(acct, path.dirname(credsPath(acct)));
    runAs().writeAs(acct, credsPath(acct), "private", text);
  }
  return { ...a, placed: { ...(a.placed || {}), [acct.name]: a.oauth.refreshToken } };
}

const refreshing = new Set();

/**
 * Keep every login alive, run on a beat: copy back a home's newer pair (only
 * once Anthropic says it is the same account), refresh what expires within
 * `REFRESH_AHEAD_MS` — here, the one holder that may spend the refresh
 * token — and push the newest pair into every home it is placed in. A home
 * whose credentials are no longer ours (the member logged in themselves, or
 * out) is dropped from the placement and never written again.
 */
async function refreshLogins() {
  const usage = require("./cc-usage");
  for (const stale of allAccounts().filter((x) => x.kind === "login")) {
    if (refreshing.has(stale.name)) continue;
    refreshing.add(stale.name);
    try {
      let a = stale;
      const homes = placedHomes(a);
      // 1. A home whose own Claude refreshed first holds the only live pair.
      for (const [member, acct] of homes) {
        const disk = homeCreds(acct);
        if (!disk || ownsHome(a, member, disk)) continue;
        const newer =
          isOauth(disk) &&
          Number(disk.expiresAt) > Number(a.oauth.expiresAt) &&
          Number(disk.expiresAt) <= Date.now() + 24 * 3600e3 + FUTURE_SLACK_MS;
        const owner = newer ? await usage.fetchOwner(disk.accessToken) : null;
        const placed = { ...a.placed };
        if (owner && owner.uuid === a.account?.accountUuid) {
          a = { ...a, oauth: { ...a.oauth, ...disk } };
          placed[member] = disk.refreshToken;
        } else {
          delete placed[member]; // their own login now: never ours to write over
        }
        a = { ...a, placed };
        save(a);
      }
      // 2. The server refreshes ahead of every home.
      if (!(Number(a.oauth.expiresAt) > Date.now() + REFRESH_AHEAD_MS)) {
        try {
          a = { ...a, oauth: await usage.exchangeRefresh(a.oauth), refreshError: null };
        } catch (e) {
          a = { ...a, refreshError: e.code || "EUPSTREAM" };
        }
        save(a);
        quotaCache.delete(a.name);
      }
      // 3. Every home still ours gets the newest pair.
      for (const [member, acct] of placedHomes(a)) {
        const disk = homeCreds(acct);
        if (!ownsHome(a, member, disk) || disk.refreshToken === a.oauth.refreshToken) continue;
        try {
          a = writeLogin(a, acct);
          save(a);
        } catch (e) {
          console.warn(`[claude-pool] ${a.name} → ${member}: ${e.message}`);
        }
      }
    } catch (e) {
      console.warn(`[claude-pool] ${stale.name}: ${e.message}`);
    } finally {
      refreshing.delete(stale.name);
    }
  }
}

// ── Quota, read in the background ──────────────────────────────────────────

/** name → { at, fiveHour?, sevenDay?, error? } */
const quotaCache = new Map();
const inFlight = new Set();

/** Read every account's quota whose reading is older than QUOTA_TTL_MS. */
async function refreshQuotas({ force = false } = {}) {
  const { fetchUsage } = require("./cc-usage");
  await Promise.all(
    allAccounts().map(async (a) => {
      const cached = quotaCache.get(a.name);
      if (inFlight.has(a.name) || (!force && cached && Date.now() - cached.at < QUOTA_TTL_MS)) {
        return;
      }
      inFlight.add(a.name);
      try {
        const u = await fetchUsage(bearerOf(a));
        quotaCache.set(a.name, { at: Date.now(), ...u });
      } catch (e) {
        // Keep the last good windows: a hiccup must not unblock a full account.
        quotaCache.set(a.name, { ...cached, at: Date.now(), error: e.code || "EUPSTREAM" });
      } finally {
        inFlight.delete(a.name);
      }
    })
  );
}

/** The highest known utilization (%) of a fresh reading, else null. */
function quotaPercent(name) {
  const q = quotaCache.get(name);
  if (!q || Date.now() - q.at > QUOTA_FRESH_MS) return null;
  const now = Date.now();
  // A window whose reset time has passed is empty again.
  const used = [q.fiveHour, q.sevenDay]
    .filter((w) => w && w.utilization != null)
    .map((w) => (w.resetsAt && Date.parse(w.resetsAt) <= now ? 0 : w.utilization));
  return used.length ? Math.max(...used) : null;
}

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

/** $ and tokens of the last 5 h / 7 d (pool-usage.js). */
const usageOf = (account) => require("./pool-usage").accountUsage("claude", account);

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

/** Every account with what the balancer sees — never a token. */
function poolStatus() {
  return allAccounts().map((a) => {
    const q = quotaCache.get(a.name);
    const percent = quotaPercent(a.name);
    return {
      name: a.name,
      kind: a.kind,
      email: a.kind === "login" ? (a.account?.emailAddress ?? null) : null,
      refreshError: a.kind === "login" ? (a.refreshError ?? null) : null,
      addedAt: a.addedAt,
      expiresAt: expiresAt(a),
      expired: expired(a),
      livePanes: livePanes(a.name),
      ...usageOf(a.name),
      quota: q
        ? {
            fiveHour: q.fiveHour || null,
            sevenDay: q.sevenDay || null,
            checkedAt: new Date(q.at).toISOString(),
            error: q.error || null,
          }
        : null,
      quotaPercent: percent,
      full: percent != null && percent >= BLOCK_PERCENT,
    };
  });
}

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

/** The full record a new Claude pane should run on, or null. `member` = run-as account. */
function pick(member) {
  let usable = allAccounts().filter((a) => !expired(a));
  if (usable.length === 0) return null;
  if (!member) {
    // A host pane's home holds the host's own login — a pool account only when
    // it was moved in from there (`hostHolds`) and is still the pair on disk:
    // the pane then runs on it as it is, and its usage is that account's.
    const disk = homeCreds(HOST_ACCT);
    usable = usable.filter((a) => a.kind === "token" || ownsHome(a, HOST, disk));
  } else {
    const disk = homeCreds(member);
    const placed = disk && usable.find((a) => ownsHome(a, member.name, disk));
    if (!placed && disk) {
      // A login AgentHub did not place there is the member's own: keep it.
      usable = usable.filter((a) => a.kind === "token");
    } else {
      // One .credentials.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 === "login" && liveInHome(a.name, member));
      if (live) return live;
    }
  }
  if (usable.length === 0) return null;
  refreshQuotas().catch(() => {});
  const ranked = usable
    .map((a) => {
      const percent = quotaPercent(a.name);
      const u = usageOf(a.name);
      return {
        a,
        key: [
          percent != null && percent >= BLOCK_PERCENT ? 1 : 0,
          livePanes(a.name),
          u.spend5h,
          u.spend7d,
        ],
      };
    })
    .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;
}

/**
 * The account a new Claude pane should run on — `{name, kind, token?}` — or
 * null when the pool has nothing this pane may use. Synchronous (a spawn
 * cannot wait on the network): it ranks on the cached quota and starts a
 * background re-read. A `login` is offered for a member, and for a host pane
 * only when it is the host's own login (`hostHolds`), which stays as it is.
 */
function pickAccount(member = null) {
  const a = pick(member);
  if (!a) return null;
  return a.kind === "login"
    ? { name: a.name, kind: "login" }
    : { name: a.name, kind: "token", token: a.token };
}

/**
 * What a new Claude pane needs from the pool, or null: `{name, env, place}`.
 * `env` goes through tmux `-e` (a member's launch file); `place()` runs just
 * before the pane starts and writes a login into the member's home.
 */
function forSpawn(member = null) {
  const a = pick(member);
  if (!a) return null;
  if (a.kind === "token") {
    return { name: a.name, env: { CLAUDE_CODE_OAUTH_TOKEN: a.token }, place: () => {} };
  }
  // The host's own login is already in place; the server keeps it fresh.
  if (!member) return { name: a.name, env: {}, place: () => {} };
  return {
    name: a.name,
    env: {},
    place: () => {
      const record = readAccount(a.name) || a;
      save(writeLogin(record, member));
      runAs().markClaudeOnboarded(member, {
        ...(record.account ? { oauthAccount: record.account } : {}),
        ...(record.userID ? { userID: record.userID } : {}),
      });
    },
  };
}

module.exports = {
  BLOCK_PERCENT,
  REFRESH_AHEAD_MS,
  addAccount,
  addLoginAccount,
  removeAccount,
  poolStatus,
  pickAccount,
  forSpawn,
  refreshQuotas,
  refreshLogins,
  __read: readAccount,
  poolDir,
  __quotaCache: quotaCache,
};