codex-proxy / src /auth /usage-refresher.ts
icebear
feat: usage stats page with time-series token tracking (#147)
df56b50 unverified
raw
history blame
7.22 kB
/**
* Usage Refresher β€” background quota refresh for all accounts.
*
* - Periodically fetches official quota from Codex backend
* - Caches quota on AccountEntry for fast dashboard reads
* - Marks quota-exhausted accounts as rate_limited
* - Evaluates warning thresholds and updates warning store
* - Non-fatal: all errors log warnings but never crash the server
*/
import { CodexApi } from "../proxy/codex-api.js";
import { CodexApiError } from "../proxy/codex-types.js";
import { getConfig } from "../config.js";
import { jitter } from "../utils/jitter.js";
import { toQuota } from "./quota-utils.js";
import {
evaluateThresholds,
updateWarnings,
type QuotaWarning,
} from "./quota-warnings.js";
import type { AccountPool } from "./account-pool.js";
import type { CookieJar } from "../proxy/cookie-jar.js";
import type { ProxyPool } from "../proxy/proxy-pool.js";
import type { UsageStatsStore } from "./usage-stats.js";
/** Check if a CodexApiError indicates the account is banned/suspended (non-CF 403). */
function isBanError(err: unknown): boolean {
if (!(err instanceof CodexApiError)) return false;
if (err.status !== 403) return false;
// Cloudflare challenge pages contain cf_chl or are HTML β€” not a ban
const body = err.body.toLowerCase();
if (body.includes("cf_chl") || body.includes("<!doctype") || body.includes("<html")) return false;
return true;
}
/** Check if a CodexApiError is a 401 token invalidation. */
function isTokenInvalidError(err: unknown): boolean {
if (!(err instanceof CodexApiError)) return false;
return err.status === 401;
}
const INITIAL_DELAY_MS = 3_000; // 3s after startup
let _refreshTimer: ReturnType<typeof setTimeout> | null = null;
let _accountPool: AccountPool | null = null;
let _cookieJar: CookieJar | null = null;
let _proxyPool: ProxyPool | null = null;
let _usageStats: UsageStatsStore | null = null;
async function fetchQuotaForAllAccounts(
pool: AccountPool,
cookieJar: CookieJar,
proxyPool: ProxyPool | null,
): Promise<void> {
if (!pool.isAuthenticated()) return;
const entries = pool.getAllEntries().filter((e) => e.status === "active" || e.status === "rate_limited" || e.status === "banned");
if (entries.length === 0) return;
const config = getConfig();
const thresholds = config.quota.warning_thresholds;
console.log(`[QuotaRefresh] Refreshing quota for ${entries.length} active/rate-limited/banned account(s)`);
const results = await Promise.allSettled(
entries.map(async (entry) => {
const proxyUrl = proxyPool?.resolveProxyUrl(entry.id);
const api = new CodexApi(entry.token, entry.accountId, cookieJar, entry.id, proxyUrl);
const usage = await api.getUsage();
const quota = toQuota(usage);
// Auto-recover banned accounts that respond successfully
if (entry.status === "banned") {
pool.markStatus(entry.id, "active");
console.log(`[QuotaRefresh] Account ${entry.id} (${entry.email ?? "?"}) unbanned β€” quota fetch succeeded`);
}
// Cache quota on the account
pool.updateCachedQuota(entry.id, quota);
// Sync rate limit window
const resetAt = usage.rate_limit.primary_window?.reset_at ?? null;
const windowSec = usage.rate_limit.primary_window?.limit_window_seconds ?? null;
pool.syncRateLimitWindow(entry.id, resetAt, windowSec);
// Mark exhausted if limit reached (primary or secondary)
if (config.quota.skip_exhausted) {
const primaryExhausted = quota.rate_limit.limit_reached;
const secondaryExhausted = quota.secondary_rate_limit?.limit_reached ?? false;
if (primaryExhausted || secondaryExhausted) {
const exhaustResetAt = primaryExhausted
? quota.rate_limit.reset_at
: quota.secondary_rate_limit?.reset_at ?? null;
pool.markQuotaExhausted(entry.id, exhaustResetAt);
console.log(`[QuotaRefresh] Account ${entry.id} (${entry.email ?? "?"}) quota exhausted β€” marked rate_limited`);
}
}
// Evaluate warning thresholds
const warnings: QuotaWarning[] = [];
const pw = evaluateThresholds(
entry.id,
entry.email,
quota.rate_limit.used_percent,
quota.rate_limit.reset_at,
"primary",
thresholds.primary,
);
if (pw) warnings.push(pw);
const sw = evaluateThresholds(
entry.id,
entry.email,
quota.secondary_rate_limit?.used_percent ?? null,
quota.secondary_rate_limit?.reset_at ?? null,
"secondary",
thresholds.secondary,
);
if (sw) warnings.push(sw);
updateWarnings(entry.id, warnings);
}),
);
let succeeded = 0;
for (let i = 0; i < results.length; i++) {
const r = results[i];
if (r.status === "fulfilled") {
succeeded++;
} else {
const entry = entries[i];
const msg = r.reason instanceof Error ? r.reason.message : String(r.reason);
// Detect banned accounts (non-CF 403)
if (isBanError(r.reason)) {
pool.markStatus(entry.id, "banned");
console.warn(`[QuotaRefresh] Account ${entry.id} (${entry.email ?? "?"}) banned β€” 403 from upstream`);
} else if (isTokenInvalidError(r.reason)) {
pool.markStatus(entry.id, "expired");
console.warn(`[QuotaRefresh] Account ${entry.id} (${entry.email ?? "?"}) token invalidated β€” 401 from upstream`);
} else {
console.warn(`[QuotaRefresh] Account ${entry.id} quota fetch failed: ${msg}`);
}
}
}
console.log(`[QuotaRefresh] Done: ${succeeded}/${entries.length} succeeded`);
// Record usage snapshot for time-series history
if (_usageStats) {
try {
_usageStats.recordSnapshot(pool);
} catch (err) {
console.warn("[QuotaRefresh] Failed to record usage snapshot:", err instanceof Error ? err.message : err);
}
}
}
function scheduleNext(
pool: AccountPool,
cookieJar: CookieJar,
): void {
const config = getConfig();
const intervalMs = jitter(config.quota.refresh_interval_minutes * 60 * 1000, 0.15);
_refreshTimer = setTimeout(async () => {
try {
await fetchQuotaForAllAccounts(pool, cookieJar, _proxyPool);
} finally {
scheduleNext(pool, cookieJar);
}
}, intervalMs);
}
/**
* Start the background quota refresh loop.
*/
export function startQuotaRefresh(
accountPool: AccountPool,
cookieJar: CookieJar,
proxyPool?: ProxyPool,
usageStats?: UsageStatsStore,
): void {
_accountPool = accountPool;
_cookieJar = cookieJar;
_proxyPool = proxyPool ?? null;
_usageStats = usageStats ?? null;
_refreshTimer = setTimeout(async () => {
try {
await fetchQuotaForAllAccounts(accountPool, cookieJar, _proxyPool);
} finally {
scheduleNext(accountPool, cookieJar);
}
}, INITIAL_DELAY_MS);
const config = getConfig();
console.log(`[QuotaRefresh] Scheduled initial quota refresh in 3s (interval: ${config.quota.refresh_interval_minutes}min)`);
}
/**
* Stop the background refresh timer.
*/
export function stopQuotaRefresh(): void {
if (_refreshTimer) {
clearTimeout(_refreshTimer);
_refreshTimer = null;
console.log("[QuotaRefresh] Stopped quota refresh");
}
}