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 — summed from
 *      `token_usage_hourly` for the sessions its runs (`dashboard_runs.account`)
 *      started.
 * 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.
 * @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 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;
}

function readAccount(name) {
  try {
    const a = JSON.parse(fs.readFileSync(fileOf(name), "utf8"));
    return typeof a.token === "string" ? a : null;
  } catch {
    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));
}

const expiresAt = (a) => new Date(Date.parse(a.addedAt) + TOKEN_LIFE_MS).toISOString();
const expired = (a) => 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) };
}

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

// ── 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(a.token);
        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 ─────────────────────────────────────

const USAGE_COLS = [
  "input_tokens",
  "output_tokens",
  "cache_read_tokens",
  "cache_write_tokens",
  "cache_write_1h_tokens",
  "web_search_requests",
  "web_fetch_requests",
  "code_execution_requests",
  "long_input_tokens",
  "long_output_tokens",
  "long_cache_read_tokens",
  "long_cache_write_tokens",
  "long_cache_write_1h_tokens",
];

/**
 * $ spent on `account` since `sinceHour` (`YYYY-MM-DDTHH`, UTC): the hourly
 * usage of every session one of its runs started, from that run's start hour.
 * ponytail: a session later resumed on ANOTHER account keeps counting here
 * too; split by run end time if that ever skews a ranking.
 */
function spendSince(account, sinceHour) {
  const { db, stmts } = require("../db");
  const { calculateCost } = require("../routes/pricing");
  const rows = db
    .prepare(
      `WITH s AS (
         SELECT session_id, MIN(substr(started_at, 1, 13)) AS since
           FROM dashboard_runs
          WHERE tool = 'claude' AND account = ? AND session_id IS NOT NULL
          GROUP BY session_id)
       SELECT h.model, h.speed, h.inference_geo, h.service_tier,
              ${USAGE_COLS.map((c) => `SUM(h.${c}) AS ${c}`).join(", ")}
         FROM s JOIN token_usage_hourly h
           ON h.session_id = s.session_id AND h.hour >= s.since AND h.hour >= ?
        GROUP BY h.model, h.speed, h.inference_geo, h.service_tier`
    )
    .all(account, sinceHour);
  return rows.length ? calculateCost(rows, stmts.listPricing.all()).total_cost : 0;
}

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

const hourAgo = (hours) => new Date(Date.now() - hours * 3600e3).toISOString().slice(0, 13);

/** 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,
      addedAt: a.addedAt,
      expiresAt: expiresAt(a),
      expired: expired(a),
      livePanes: livePanes(a.name),
      spend5h: spendSince(a.name, hourAgo(5)),
      spend7d: spendSince(a.name, hourAgo(24 * 7)),
      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,
    };
  });
}

/**
 * The account a new Claude pane should run on — `{name, token}` — or null when
 * the pool has no usable account. Synchronous (a spawn cannot wait on the
 * network): it ranks on the cached quota and starts a background re-read.
 */
function pickAccount() {
  const usable = allAccounts().filter((a) => !expired(a));
  if (usable.length === 0) return null;
  refreshQuotas().catch(() => {});
  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, hourAgo(5)),
          spendSince(a.name, hourAgo(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);
    });
  const { name, token } = ranked[0].a;
  return { name, token };
}

module.exports = {
  BLOCK_PERCENT,
  addAccount,
  removeAccount,
  poolStatus,
  pickAccount,
  refreshQuotas,
  poolDir,
  __quotaCache: quotaCache,
};