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.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130import { jsonBody, makeLimit, unknownReport } from "../driver.js";import { defineUsageFeed } from "../registry.js";import type { UsageLimit, UsageReportBody, UsageState, UsageStep } from "../types.js";
/** Endpoint and window shape of the ChatGPT backend usage call. */export const CODEX_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";
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 windowLabel(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`;}
/** `reset_at` above 1e12 is already milliseconds; below it is seconds. */function parseWindow(window: unknown, nowMs: number): UsageLimit | undefined { const value = asRecord(window); 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 makeLimit("Usage", usedPercent, { windowLabel: windowLabel(asNumber(value.limit_window_seconds)), resetsAt: resetMs ?? (resetAfter === undefined ? undefined : nowMs + resetAfter * 1000), });}
export function parseCodexUsage(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); if (primary) limits.push({ ...primary, label: "Primary" }); const secondary = parseWindow(rateLimit?.secondary_window, nowMs); 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 extra = 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(extra?.primary_window, nowMs); if (extraPrimary) limits.push({ ...extraPrimary, label: `${name} · Primary` }); const extraSecondary = parseWindow(extra?.secondary_window, nowMs); if (extraSecondary) limits.push({ ...extraSecondary, label: `${name} · Secondary` }); }
return { planType: typeof root?.plan_type === "string" ? root.plan_type : undefined, limits };}
export const codexUsageFeed = defineUsageFeed({ provider: "openai-codex", providerName: "ChatGPT Codex", credentialKinds: ["oauth"], advance(state: UsageState, responses): UsageStep { const credential = state.credential; if (credential.kind !== "oauth") return { report: unknownReport("openai-codex needs an OAuth account") }; const access = credential.access; const accountId = credential.accountId ?? jwtClaim(access, "https://api.openai.com/auth", "chatgpt_account_id"); const email = credential.email ?? jwtClaimEmail(access);
if (responses.length === 0) { const headers: Record<string, string> = { Authorization: `Bearer ${access}`, "User-Agent": "prime-agent", }; if (accountId) headers["ChatGPT-Account-Id"] = accountId; return { requests: [{ kind: "http", method: "GET", url: CODEX_USAGE_URL, headers }] }; }
const response = responses[0]; if (!response || response.status !== 200) return { report: unknownReport(`HTTP ${response?.status ?? "no response"}`) }; const payload = jsonBody(response); if (payload === undefined) return { report: unknownReport("usage body is not JSON") };
const body: UsageReportBody = { ...parseCodexUsage(payload, state.context.nowMs), accountId, email, error: null, }; return { report: body }; },});
function jwtPayload(token: string): Record<string, unknown> | undefined { const parts = token.split("."); if (parts.length !== 3) return undefined; try { const normalized = parts[1]!.replace(/-/g, "+").replace(/_/g, "/"); const padded = normalized.padEnd(normalized.length + ((4 - (normalized.length % 4)) % 4), "="); const decoded = typeof Buffer === "undefined" ? atob(padded) : Buffer.from(padded, "base64").toString("utf8"); return asRecord(JSON.parse(decoded)); } catch { return undefined; }}
function jwtClaim(token: string, namespace: string, key: string): string | undefined { const claim = asRecord(jwtPayload(token)?.[namespace])?.[key]; return typeof claim === "string" ? claim : undefined;}
function jwtClaimEmail(token: string): string | undefined { const email = asRecord(jwtPayload(token)?.["https://api.openai.com/profile"])?.email; if (typeof email !== "string") return undefined; const trimmed = email.trim().toLowerCase(); return trimmed === "" ? undefined : trimmed;}