AgentHub/server/lib/quota.js

/**
 * @file quota.js
 * @description How much of each native AI tool's plan quota is spent, for the
 * status bar's quota chips (`GET /api/pricing/quota`).
 *
 *  - **Claude**: the logged-in account's `GET /api/oauth/usage` (`cc-usage.js`),
 *    with the access token the CLI holds RIGHT NOW. It is never refreshed here —
 *    a refresh rotates the CLI's credentials — so an expired token answers
 *    `ELIVEEXPIRED` until Claude Code next refreshes it itself. Every OTHER
 *    saved auth profile (`cc-profiles.js`) is asked with its own stored token
 *    while that token is valid and never refreshed (a refresh rotates the
 *    saved credentials); past expiry it shows its last figure, `asOf`. Cached
 *    `CLAUDE_TTL_MS` per account, so a poll from every open tab is one
 *    upstream call each. The live entry is `active: true`.
 *  - **Codex**: no network. Every `token_count` event Codex writes into its
 *    rollout journal carries the account's `rate_limits`; the newest one wins.
 *
 * Each tool also carries its signed-in `account` (an email) and `expiresAt`,
 * when its stored access token runs out — Claude's `claudeAiOauth.expiresAt`,
 * Codex's `auth.json` access-token `exp` (claims read, never verified).
 *
 * Only percentages, window lengths, reset/expiry times and the account's
 * email leave this module — never a token.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

const fs = require("node:fs");
const path = require("node:path");
const { getCodexHome } = require("./codex-home");
const { readActiveAuth, readProfile, listProfiles } = require("./cc-profiles");
const { fetchUsage } = require("./cc-usage");
const { allRolloutFiles } = require("./codex-ingest");

const CLAUDE_TTL_MS = 5 * 60_000;
const FORCE_MIN_MS = 30_000; // a Refresh click still cannot hammer the endpoint
const EXPIRY_MARGIN_MS = 60_000;
const TAIL_BYTES = 512 * 1024;
const CODEX_FILES = 5;

const cache = new Map(); // "live" | "profile:<name>" → { at, value }
const lastGood = new Map(); // profile name → { windows, at }

/** One upstream call per key per `CLAUDE_TTL_MS` (`FORCE_MIN_MS` on a click). */
async function cached(key, force, fn) {
  const hit = cache.get(key);
  const age = hit ? Date.now() - hit.at : Infinity;
  if (age < (force ? FORCE_MIN_MS : CLAUDE_TTL_MS)) return hit.value;
  const value = await fn();
  cache.set(key, { at: Date.now(), value });
  return value;
}

function windowLabel(minutes) {
  if (minutes === 300) return "5h";
  if (minutes === 10080) return "wk";
  if (minutes >= 40000 && minutes <= 45000) return "mo";
  return minutes % 1440 === 0 ? `${minutes / 1440}d` : `${Math.round(minutes / 60)}h`;
}

const iso = (ms) => (Number(ms) ? new Date(Number(ms)).toISOString() : null);
const fresh = (ms) => Number(ms) > Date.now() + EXPIRY_MARGIN_MS;

function pct(w) {
  return { usedPercent: w.utilization, resetsAt: w.resetsAt };
}

function claudeWindows(u) {
  const windows = [];
  if (u.fiveHour) windows.push({ label: "5h", ...pct(u.fiveHour) });
  if (u.sevenDay) windows.push({ label: "wk", ...pct(u.sevenDay) });
  return windows;
}

/** The account the CLI is logged into, with the token it holds right now. */
async function liveClaude() {
  let who = {};
  try {
    const { oauth, account } = readActiveAuth();
    if (!oauth || !oauth.accessToken) return null; // API-key user or not logged in
    who = {
      account: account?.emailAddress || account?.displayName || null,
      expiresAt: iso(oauth.expiresAt),
    };
    if (!fresh(oauth.expiresAt)) return { ...who, windows: [], error: "ELIVEEXPIRED" };
    return { ...who, windows: claudeWindows(await fetchUsage(oauth.accessToken)) };
  } catch (e) {
    return { ...who, windows: [], error: e.code || "EUPSTREAM" };
  }
}

/**
 * A saved profile that is not the live account, asked with its OWN stored
 * token only while that token is still valid. Never refreshed here: a refresh
 * rotates the profile's credentials on disk, so it happens only on the
 * Profiles modal's Check quota click (`cc-usage.profileUsage`). Past expiry
 * the last figure this process read is shown with its `asOf`.
 */
async function profileClaude(item) {
  const who = { account: item.email || item.displayName || null, expiresAt: iso(item.expiresAt) };
  let creds = null;
  try {
    creds = readProfile(item.name)?.credentials || null;
  } catch {
    /* a corrupt profile file reads as signed out */
  }
  if (creds?.accessToken && fresh(creds.expiresAt)) {
    try {
      const windows = claudeWindows(await fetchUsage(creds.accessToken));
      lastGood.set(item.name, { windows, at: Date.now() });
      return { ...who, windows };
    } catch (e) {
      return { ...who, windows: [], error: e.code || "EUPSTREAM" };
    }
  }
  const prev = lastGood.get(item.name);
  if (!prev) return { ...who, windows: [], error: "EPROFILEEXPIRED" };
  return { ...who, windows: prev.windows.map(expireReset), asOf: iso(prev.at) };
}

/** A window whose reset has passed is empty again. */
function expireReset(w) {
  return w.resetsAt && Date.parse(w.resetsAt) < Date.now() ? { ...w, usedPercent: 0 } : w;
}

/** The live account first, then every other saved profile (by label). */
async function claudeQuotas(force) {
  let items = [];
  try {
    items = listProfiles().items;
  } catch {
    /* no profiles dir */
  }
  const [live, ...others] = await Promise.all([
    cached("live", force, liveClaude),
    ...items
      .filter((it) => !it.active)
      .map(async (it) => ({
        id: "claude",
        active: false,
        label: it.label,
        ...(await cached(`profile:${it.name}`, force, () => profileClaude(it))),
      })),
  ]);
  const mine = items.find((it) => it.active);
  return [live && { id: "claude", active: true, label: mine?.label || null, ...live }, ...others];
}

/** The newest `rate_limits` in the last `TAIL_BYTES` of a rollout, or null. */
function lastRateLimits(file) {
  let text;
  try {
    const fd = fs.openSync(file.path, "r");
    try {
      const len = Math.min(file.size, TAIL_BYTES);
      const buf = Buffer.alloc(len);
      fs.readSync(fd, buf, 0, len, file.size - len);
      text = buf.toString("utf8");
    } finally {
      fs.closeSync(fd);
    }
  } catch {
    return null;
  }
  const lines = text.split("\n");
  for (let i = lines.length - 1; i >= 0; i--) {
    if (!lines[i].includes('"rate_limits"')) continue;
    try {
      const rl = JSON.parse(lines[i])?.payload?.rate_limits;
      if (rl && (rl.primary || rl.secondary)) return { rl, mtime: file.mtime };
    } catch {
      /* a cut first line or a partial write */
    }
  }
  return null;
}

function codexWindow(w, mtime) {
  if (!w || typeof w.used_percent !== "number") return null;
  const resetsMs = Number(w.resets_at)
    ? Number(w.resets_at) * 1000
    : Number(w.resets_in_seconds)
      ? mtime + Number(w.resets_in_seconds) * 1000
      : null;
  // A window that has reset since Codex last wrote is empty again.
  const used = resetsMs && resetsMs < Date.now() ? 0 : w.used_percent;
  return {
    label: windowLabel(Number(w.window_minutes)),
    usedPercent: used,
    resetsAt: resetsMs ? new Date(resetsMs).toISOString() : null,
  };
}

/** A JWT's claims, unverified — only for showing who is signed in. */
function jwtClaims(token) {
  return JSON.parse(Buffer.from(String(token).split(".")[1], "base64url").toString("utf8"));
}

/** Codex's ChatGPT login from `auth.json`: email + access-token expiry. */
function codexAccount() {
  try {
    const { tokens } = JSON.parse(fs.readFileSync(path.join(getCodexHome(), "auth.json"), "utf8"));
    const { exp } = jwtClaims(tokens.access_token);
    return {
      account: jwtClaims(tokens.id_token).email || null,
      expiresAt: exp ? new Date(exp * 1000).toISOString() : null,
    };
  } catch {
    return {}; // API-key login, no file, or a token that is not a JWT
  }
}

function codexQuota() {
  const files = allRolloutFiles()
    .sort((a, b) => b.mtime - a.mtime)
    .slice(0, CODEX_FILES);
  for (const f of files) {
    const found = lastRateLimits(f);
    if (!found) continue;
    const windows = [found.rl.primary, found.rl.secondary]
      .map((w) => codexWindow(w, found.mtime))
      .filter(Boolean);
    return {
      id: "codex",
      active: true,
      ...codexAccount(),
      windows,
      asOf: new Date(found.mtime).toISOString(),
    };
  }
  return null;
}

async function quotas({ force = false } = {}) {
  const tools = [...(await claudeQuotas(force)), codexQuota()].filter(Boolean);
  return { tools, checkedAt: new Date().toISOString() };
}

function __resetCache() {
  cache.clear();
  lastGood.clear();
}

module.exports = { quotas, codexWindow, codexAccount, windowLabel, __resetCache };