Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
10 kB · 265 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266import { Buffer } from "node:buffer";import type { AuthStorage, OAuthCredential } from "../../core/auth-storage.js";import { credentialIdentity } from "../../core/auth-storage.js";import type { ContextTreeNode } from "../../core/context-tree.js";import { formatContextTree } from "./components/context-tree-format.js";import { theme } from "./theme/theme.js";
const CODEX_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";const BAR_WIDTH = 24;
type FetchLike = (input: string | URL, init?: RequestInit) => Promise<Response>;
export interface UsageLimit { label: string; usedPercent?: number; windowLabel?: string; resetsAt?: number; status?: "ok" | "warning" | "exhausted" | "unknown";}
export interface UsageReport { provider: string; fetchedAt: number; accountId?: string; email?: string; planType?: string; limits: UsageLimit[]; error?: string;}
interface CodexUsageWindow { used_percent?: unknown; limit_window_seconds?: unknown; reset_after_seconds?: unknown; reset_at?: unknown;}
function asRecord(value: unknown): Record<string, unknown> | undefined { return value && typeof value === "object" && !Array.isArray(value) ? (value as Record<string, unknown>) : undefined;}
function asNumber(value: unknown): number | undefined { if (typeof value === "number" && Number.isFinite(value)) return value; if (typeof value === "string" && value.trim()) { const parsed = Number(value); return Number.isFinite(parsed) ? parsed : undefined; } return undefined;}
function decodeJwtPayload(token: string): Record<string, unknown> | undefined { const encoded = token.split(".")[1]; if (!encoded) return undefined; try { const normalized = encoded.replace(/-/g, "+").replace(/_/g, "/"); return asRecord(JSON.parse(Buffer.from(normalized, "base64").toString("utf8"))); } catch { return undefined; }}
function tokenAccountId(token: string): string | undefined { const payload = decodeJwtPayload(token); const auth = asRecord(payload?.["https://api.openai.com/auth"]); return typeof auth?.chatgpt_account_id === "string" ? auth.chatgpt_account_id : undefined;}
function tokenEmail(token: string): string | undefined { const payload = decodeJwtPayload(token); const profile = asRecord(payload?.["https://api.openai.com/profile"]); const email = profile?.email; return typeof email === "string" && email.trim() ? email.trim().toLowerCase() : undefined;}
function accountLabel(report: UsageReport, index: number): string { if (report.email) return report.email; if (report.accountId) return report.accountId; return `account ${index + 1}`;}
function formatDuration(ms: number): string { const seconds = Math.max(0, Math.round(ms / 1000)); if (seconds < 60) return `${seconds}s`; const minutes = Math.round(seconds / 60); if (minutes < 60) return `${minutes}m`; const hours = Math.round(minutes / 60); if (hours < 48) return `${hours}h`; return `${Math.round(hours / 24)}d`;}
function formatWindowLabel(seconds: number | undefined): string | undefined { if (seconds === undefined || seconds <= 0) return undefined; if (seconds >= 86_400) return `${Math.round(seconds / 86_400)}d window`; return `${Math.max(1, Math.round(seconds / 3_600))}h window`;}
function parseWindow(window: unknown, nowMs: number): UsageLimit | undefined { const value = asRecord(window) as CodexUsageWindow | undefined; if (!value) return undefined; const usedPercent = asNumber(value.used_percent); const resetAt = asNumber(value.reset_at); const resetAfter = asNumber(value.reset_after_seconds); const resetMs = resetAt === undefined ? undefined : resetAt > 1_000_000_000_000 ? resetAt : resetAt * 1000; return { label: "Usage", usedPercent, windowLabel: formatWindowLabel(asNumber(value.limit_window_seconds)), resetsAt: resetMs ?? (resetAfter === undefined ? undefined : nowMs + resetAfter * 1000), status: usedPercent === undefined ? "unknown" : usedPercent >= 100 ? "exhausted" : usedPercent >= 90 ? "warning" : "ok", };}
function parseUsagePayload(payload: unknown, nowMs: number): { planType?: string; limits: UsageLimit[] } { const root = asRecord(payload); const rateLimit = asRecord(root?.rate_limit); const limits: UsageLimit[] = []; const primary = parseWindow(rateLimit?.primary_window, nowMs); const secondary = parseWindow(rateLimit?.secondary_window, nowMs); if (primary) limits.push({ ...primary, label: "Primary" }); if (secondary) limits.push({ ...secondary, label: "Secondary" }); const additional = Array.isArray(root?.additional_rate_limits) ? root.additional_rate_limits : []; for (const entry of additional) { const item = asRecord(entry); const extraRateLimit = asRecord(item?.rate_limit); const name = typeof item?.limit_name === "string" ? item.limit_name : typeof item?.metered_feature === "string" ? item.metered_feature : "Additional"; const extraPrimary = parseWindow(extraRateLimit?.primary_window, nowMs); const extraSecondary = parseWindow(extraRateLimit?.secondary_window, nowMs); if (extraPrimary) limits.push({ ...extraPrimary, label: `${name} · Primary` }); if (extraSecondary) limits.push({ ...extraSecondary, label: `${name} · Secondary` }); } return { planType: typeof root?.plan_type === "string" ? root.plan_type : undefined, limits, };}
async function fetchCodexUsage( credential: OAuthCredential, options: { fetchImpl: FetchLike; signal?: AbortSignal },): Promise<UsageReport> { const nowMs = Date.now(); const accountId = typeof credential.accountId === "string" ? credential.accountId : tokenAccountId(credential.access); const email = typeof credential.email === "string" ? credential.email : tokenEmail(credential.access); const label: UsageReport = { provider: "openai-codex", fetchedAt: nowMs, accountId, email, limits: [] }; try { const headers: Record<string, string> = { Authorization: `Bearer ${credential.access}`, "User-Agent": "prime-agent", }; if (accountId) headers["ChatGPT-Account-Id"] = accountId; const response = await options.fetchImpl(CODEX_USAGE_URL, { headers, signal: options.signal }); if (!response.ok) return { ...label, error: `HTTP ${response.status}` }; const parsed = parseUsagePayload(await response.json(), nowMs); return { ...label, ...parsed }; } catch (error) { if (options.signal?.aborted) return { ...label, error: "cancelled" }; return { ...label, error: error instanceof Error ? error.message : String(error) }; }}
/** Fetch live Codex quota once for every stored OAuth account. */export async function fetchCodexUsageReports( authStorage: AuthStorage, options: { fetchImpl?: FetchLike; signal?: AbortSignal } = {},): Promise<UsageReport[]> { const credentials = authStorage .getCredentialList("openai-codex") .filter((credential): credential is OAuthCredential => credential.type === "oauth"); if (credentials.length === 0) return []; const fetchImpl = options.fetchImpl ?? fetch; return Promise.all( credentials.map(async (credential) => { const identity = credentialIdentity(credential); let resolved: Awaited<ReturnType<AuthStorage["getApiKeyForCredential"]>>; try { resolved = await authStorage.getApiKeyForCredential("openai-codex", identity); } catch { resolved = undefined; } if (!resolved || resolved.credential.type !== "oauth") { const accountId = typeof credential.accountId === "string" ? credential.accountId : tokenAccountId(credential.access); const email = typeof credential.email === "string" ? credential.email : tokenEmail(credential.access); return { provider: "openai-codex", fetchedAt: Date.now(), accountId, email, limits: [], error: "OAuth token unavailable; run /login", }; } return fetchCodexUsage(resolved.credential, { fetchImpl, signal: options.signal }); }), );}
function renderBar(usedPercent: number | undefined): string { if (usedPercent === undefined) return `[${"·".repeat(BAR_WIDTH)}] unknown`; const fraction = Math.max(0, Math.min(1, usedPercent / 100)); const filled = Math.round(fraction * BAR_WIDTH); return `[${"█".repeat(filled)}${"░".repeat(BAR_WIDTH - filled)}] ${Math.round(usedPercent)}%`;}
function renderLimit(limit: UsageLimit, nowMs: number): string[] { const used = limit.usedPercent === undefined ? "unknown" : `${limit.usedPercent.toFixed(1)}% used`; const left = limit.usedPercent === undefined ? "" : ` · ${(Math.max(0, 100 - limit.usedPercent)).toFixed(1)}% left`; const reset = limit.resetsAt && limit.resetsAt > nowMs ? ` · resets in ${formatDuration(limit.resetsAt - nowMs)}` : ""; return [ ` ${limit.label}${limit.windowLabel ? ` (${limit.windowLabel})` : ""}: ${used}${left}${reset}`, ` ${renderBar(limit.usedPercent)}`, ];}
export function renderUsageReportText( reports: readonly UsageReport[], contextTree: ContextTreeNode | undefined, width: number, options: { loading?: boolean; fetchedAt?: number } = {},): string { const nowMs = Date.now(); const lines = [theme.bold(theme.fg("accent", "Usage"))]; if (options.fetchedAt) lines[0] += theme.fg("dim", ` (${formatDuration(nowMs - options.fetchedAt)} ago)`); if (options.loading) lines.push(theme.fg("dim", "Refreshing provider quotas in the background…")); if (reports.length === 0) { lines.push(theme.fg("dim", "No Codex quota snapshot available yet.")); } else { lines.push(""); lines.push(theme.bold(theme.fg("accent", "ChatGPT Codex"))); reports.forEach((report, index) => { lines.push(""); lines.push( ` ${theme.fg("accent", accountLabel(report, index))}${report.planType ? ` · ${report.planType}` : ""}`, ); if (report.error) { lines.push(` ${theme.fg("warning", `Unable to fetch quota: ${report.error}`)}`); } else if (report.limits.length === 0) { lines.push(theme.fg("dim", " No limits reported")); } else { for (const limit of report.limits) lines.push(...renderLimit(limit, nowMs)); } }); } if (contextTree) { lines.push("", theme.bold(theme.fg("accent", "This session"))); lines.push(formatContextTree(contextTree, Math.max(60, width))); } else { lines.push("", theme.fg("dim", "Session usage will appear when a context snapshot is available.")); } return lines.join("\n");}