diff --git a/index.ts b/index.ts index 1838950..30dd587 100644 --- a/index.ts +++ b/index.ts @@ -693,23 +693,50 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { accounts: menuAccounts, input: process.stdin, output: process.stdout, - handlers: { - onCheckQuotas: async () => { - await Promise.all( - accounts.map(async (acc, index) => { - if (acc.enabled === false) return; - const live = accountManager.getAccountByIndex(index); - if (!live) return; - const auth = accountManager.toAuthDetails(live); - if (auth.access && auth.expires > Date.now()) { - await codexStatus.fetchFromBackend(live, auth.access); - } - }), - ); - const snapshots = await codexStatus.getAllSnapshots(); - const report = renderQuotaReport(menuAccounts, snapshots, Date.now()); - process.stdout.write(report.join("\n") + "\n"); - }, + handlers: { + onCheckQuotas: async () => { + await Promise.all( + accounts.map(async (acc, index) => { + if (acc.enabled === false) return; + const live = accountManager.getAccountByIndex(index); + if (!live) return; + try { + let accessToken: string | null = null; + const auth = accountManager.toAuthDetails(live); + if (auth.access && auth.expires > Date.now()) { + accessToken = auth.access; + } else { + const refreshed = await accountManager.refreshAccountWithFallback(live); + if (refreshed.type === "success") { + if (refreshed.headers) { + await codexStatus.updateFromHeaders( + live, + Object.fromEntries(refreshed.headers.entries()), + ); + } + const refreshedAuth = { + type: "oauth" as const, + access: refreshed.access, + refresh: refreshed.refresh, + expires: refreshed.expires, + }; + accountManager.updateFromAuth(live, refreshedAuth); + await accountManager.saveToDisk(); + accessToken = refreshed.access; + } + } + if (accessToken) { + await codexStatus.fetchFromBackend(live, accessToken); + } + } catch { + // Keep rendering other accounts even if one fails quota refresh/fetch. + } + }), + ); + const snapshots = await codexStatus.getAllSnapshots(); + const report = renderQuotaReport(accounts, snapshots, Date.now()); + process.stdout.write(report.join("\n") + "\n"); + }, onConfigureModels: async () => { process.stdout.write( "Edit your opencode.jsonc (or opencode.json) to configure models.\n", @@ -769,18 +796,18 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { authorize: async (_inputs?: Record) => { let replaceExisting = false; - const existingStorage = await loadAccounts(); - if (existingStorage?.accounts?.length && process.stdin.isTTY && process.stdout.isTTY) { - const menuResult = await runInteractiveAuthMenu({ allowExit: true }); - if (menuResult === "exit") { - return { + const existingStorage = await loadAccounts(); + if (existingStorage?.accounts?.length && process.stdin.isTTY && process.stdout.isTTY) { + const menuResult = await runInteractiveAuthMenu({ allowExit: true }); + if (menuResult === "exit") { + return { url: "about:blank", method: "code" as const, instructions: "Login cancelled.", - callback: async () => ({ type: "failed" as const }), - }; + callback: async () => ({ type: "failed" as const }), + }; + } } - } const { pkce, state, url } = await createAuthorizationFlow(); let serverInfo = null; diff --git a/lib/ui/auth-menu.ts b/lib/ui/auth-menu.ts index daf8136..242322a 100644 --- a/lib/ui/auth-menu.ts +++ b/lib/ui/auth-menu.ts @@ -46,6 +46,12 @@ function colorize(text: string, color: string, useColor: boolean): string { return useColor ? `${color}${text}${ANSI.reset}` : text; } +function formatAccountDisplayName(account: AccountInfo): string { + const base = account.email || `Account ${account.index + 1}`; + const plan = typeof account.plan === "string" ? account.plan.trim() : ""; + return plan ? `${base} (${plan})` : base; +} + function getStatusBadge(status: AccountStatus | undefined, useColor: boolean): string { switch (status) { case "rate-limited": @@ -76,7 +82,7 @@ export function formatStatusBadges( } function buildAccountLabel(account: AccountInfo, useColor: boolean): string { - const baseLabel = account.email || `Account ${account.index + 1}`; + const baseLabel = formatAccountDisplayName(account); const badges = formatStatusBadges(account, useColor); return badges ? `${baseLabel} ${badges}` : baseLabel; } @@ -190,7 +196,7 @@ export async function showAccountDetails( ): Promise { const useColor = options.useColor ?? shouldUseColor(); const output = options.output ?? process.stdout; - const label = account.email || `Account ${account.index + 1}`; + const label = formatAccountDisplayName(account); const badges = formatStatusBadges(account, useColor); const bold = useColor ? ANSI.bold : ""; diff --git a/lib/ui/codex-quota-report.ts b/lib/ui/codex-quota-report.ts index 193fa9d..6c75a72 100644 --- a/lib/ui/codex-quota-report.ts +++ b/lib/ui/codex-quota-report.ts @@ -1,106 +1,151 @@ import { createHash } from "node:crypto"; import type { CodexRateLimitSnapshot } from "../codex-status.js"; +import { ANSI, shouldUseColor } from "./tty/ansi.js"; type AccountLike = { - accountId?: string; - email?: string; - plan?: string; - refreshToken?: string; + accountId?: string; + email?: string; + plan?: string; + refreshToken?: string; }; -const LINE = "=".repeat(52); const BAR_WIDTH = 20; +function colorize(text: string, color: string, useColor: boolean): string { + return useColor ? `${color}${text}${ANSI.reset}` : text; +} + function getSnapshotKey(account: AccountLike): string | null { - if (account.accountId && account.email && account.plan) { - return `${account.accountId}|${account.email.toLowerCase()}|${account.plan}`; - } - if (account.refreshToken) { - return createHash("sha256").update(account.refreshToken).digest("hex"); - } - return null; + if (account.accountId && account.email && account.plan) { + return `${account.accountId}|${account.email.toLowerCase()}|${account.plan}`; + } + if (account.refreshToken) { + return createHash("sha256").update(account.refreshToken).digest("hex"); + } + return null; } function findSnapshot( - account: AccountLike, - snapshots: CodexRateLimitSnapshot[], + account: AccountLike, + snapshots: CodexRateLimitSnapshot[], ): CodexRateLimitSnapshot | undefined { - const key = getSnapshotKey(account); - if (key) { - const direct = snapshots.find((snapshot) => (snapshot as any).key === key); - if (direct) return direct; - } + const key = getSnapshotKey(account); + if (key) { + const direct = snapshots.find((snapshot) => (snapshot as any).key === key); + if (direct) return direct; + } - return snapshots.find( - (snapshot) => - snapshot.accountId === account.accountId && - snapshot.email?.toLowerCase() === account.email?.toLowerCase() && - snapshot.plan === account.plan, - ); + return snapshots.find( + (snapshot) => + snapshot.accountId === account.accountId && + snapshot.email?.toLowerCase() === account.email?.toLowerCase() && + snapshot.plan === account.plan, + ); } function formatReset(resetAt: number, now: number): string { - if (!resetAt || resetAt <= now) return ""; - const date = new Date(resetAt); - const timeStr = `${String(date.getHours()).padStart(2, "0")}:${String(date.getMinutes()).padStart(2, "0")}`; - if (resetAt - now <= 24 * 60 * 60 * 1000) { - return ` (resets ${timeStr})`; - } - const months = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]; - return ` (resets ${timeStr} ${date.getDate()} ${months[date.getMonth()]})`; + if (!resetAt || resetAt <= now) return ""; + const date = new Date(resetAt); + const timeStr = `${String(date.getHours()).padStart(2, "0")}:${String(date.getMinutes()).padStart(2, "0")}`; + if (resetAt - now <= 24 * 60 * 60 * 1000) { + return ` (resets ${timeStr})`; + } + const months = [ + "Jan", + "Feb", + "Mar", + "Apr", + "May", + "Jun", + "Jul", + "Aug", + "Sep", + "Oct", + "Nov", + "Dec", + ]; + return ` (resets ${timeStr} ${date.getDate()} ${months[date.getMonth()]})`; } -function renderBar(usedPercent: number | undefined | null): { bar: string; percent: string } { - if (usedPercent === undefined || usedPercent === null || Number.isNaN(usedPercent)) { - const empty = "░".repeat(BAR_WIDTH); - return { bar: empty, percent: "???" }; - } - const left = Math.max(0, Math.min(100, Math.round(100 - usedPercent))); - const filled = Math.round((left / 100) * BAR_WIDTH); - const bar = "█".repeat(filled) + "░".repeat(BAR_WIDTH - filled); - const percent = `${String(left).padStart(3, " ")}%`; - return { bar, percent }; +function renderBar( + usedPercent: number | undefined | null, + useColor: boolean, +): { bar: string; percent: string } { + if ( + usedPercent === undefined || + usedPercent === null || + Number.isNaN(usedPercent) + ) { + const empty = "░".repeat(BAR_WIDTH); + return { bar: empty, percent: "???" }; + } + const left = Math.max(0, Math.min(100, Math.round(100 - usedPercent))); + const filled = Math.round((left / 100) * BAR_WIDTH); + const filledBar = "█".repeat(filled); + const emptyBar = "░".repeat(BAR_WIDTH - filled); + const fillColor = + left >= 70 ? ANSI.green : left >= 35 ? ANSI.yellow : ANSI.red; + const bar = `${colorize(filledBar, fillColor, useColor)}${colorize(emptyBar, ANSI.dim, useColor)}`; + const percent = `${String(left).padStart(3, " ")}%`; + return { bar, percent }; +} + +function formatAccountName(account: AccountLike): string { + const label = account.email || account.accountId || "Unknown account"; + const plan = typeof account.plan === "string" ? account.plan.trim() : ""; + return plan ? `${label} (${plan})` : label; } export function renderQuotaReport( - accounts: AccountLike[], - snapshots: CodexRateLimitSnapshot[], - now = Date.now(), + accounts: AccountLike[], + snapshots: CodexRateLimitSnapshot[], + now = Date.now(), + useColor = shouldUseColor(), ): string[] { - const lines: string[] = ["Checking quotas for all accounts..."]; - for (const account of accounts) { - const email = account.email || account.accountId || "Unknown account"; - const snapshot = findSnapshot(account, snapshots); - lines.push(LINE); - lines.push(` ${email}`); - lines.push(LINE); - lines.push(" +- Codex CLI Quota"); - - const primary = renderBar(snapshot?.primary?.usedPercent); - const primaryReset = snapshot?.primary?.resetAt - ? formatReset(snapshot.primary.resetAt, now) - : primary.percent === "???" - ? " ???" - : ""; - lines.push(` | |- GPT-5 ${primary.bar} ${primary.percent}${primaryReset}`); + const lines: string[] = []; + lines.push(`${colorize("┌", ANSI.dim, useColor)} Quota Report`); + for (const account of accounts) { + const name = formatAccountName(account); + const snapshot = findSnapshot(account, snapshots); + lines.push(`${colorize("│", ANSI.cyan, useColor)}`); + lines.push( + `${colorize("├", ANSI.cyan, useColor)} ${colorize(name, ANSI.bold, useColor)}`, + ); + lines.push( + `${colorize("│", ANSI.cyan, useColor)} ${colorize("Codex CLI Quota", ANSI.dim, useColor)}`, + ); - const secondary = renderBar(snapshot?.secondary?.usedPercent); - const secondaryReset = snapshot?.secondary?.resetAt - ? formatReset(snapshot.secondary.resetAt, now) - : secondary.percent === "???" - ? " ???" - : ""; - lines.push(` | |- Weekly ${secondary.bar} ${secondary.percent}${secondaryReset}`); + const primary = renderBar(snapshot?.primary?.usedPercent, useColor); + const primaryReset = snapshot?.primary?.resetAt + ? formatReset(snapshot.primary.resetAt, now) + : primary.percent === "???" + ? " ???" + : ""; + lines.push( + `${colorize("│", ANSI.cyan, useColor)} ${colorize("●", ANSI.green, useColor)} 5h ${primary.bar} ${primary.percent}${primaryReset}`, + ); - const creditInfo = snapshot?.credits; - const creditStr = creditInfo - ? creditInfo.unlimited - ? "unlimited" - : `${creditInfo.balance} credits` - : "0 credits"; - lines.push(` |---- Credits ${creditStr}`); - } + const secondary = renderBar(snapshot?.secondary?.usedPercent, useColor); + const secondaryReset = snapshot?.secondary?.resetAt + ? formatReset(snapshot.secondary.resetAt, now) + : secondary.percent === "???" + ? " ???" + : ""; + lines.push( + `${colorize("│", ANSI.cyan, useColor)} ${colorize("●", ANSI.green, useColor)} Weekly ${secondary.bar} ${secondary.percent}${secondaryReset}`, + ); - return lines; + const creditInfo = snapshot?.credits; + const creditStr = creditInfo + ? creditInfo.unlimited + ? "unlimited" + : `${creditInfo.balance} credits` + : "0 credits"; + lines.push( + `${colorize("│", ANSI.cyan, useColor)} ${colorize("●", ANSI.green, useColor)} Credits ${creditStr}`, + ); + } + lines.push(`${colorize("└", ANSI.cyan, useColor)}`); + return lines; } diff --git a/test/auth-menu.test.ts b/test/auth-menu.test.ts index 8d7ce17..c123241 100644 --- a/test/auth-menu.test.ts +++ b/test/auth-menu.test.ts @@ -76,6 +76,7 @@ describe("auth menu helpers", () => { const accountItem = items.find((item) => item.label.includes(account.email)); expect(accountItem).toBeTruthy(); + expect(accountItem!.label).toContain(`(${account.plan})`); expect(accountItem!.label).toContain("[enabled]"); expect(accountItem!.label).toContain("[last active]"); expect(accountItem!.hint).toContain("used"); diff --git a/test/codex-quota-report.test.ts b/test/codex-quota-report.test.ts index a7368fa..3240d8d 100644 --- a/test/codex-quota-report.test.ts +++ b/test/codex-quota-report.test.ts @@ -4,28 +4,47 @@ import { describe, it, expect } from "vitest"; import { renderQuotaReport } from "../lib/ui/codex-quota-report.js"; +const ANSI_REGEX = /\u001b\[[0-9;]*m/g; + describe("codex quota report", () => { - it("renders account blocks with aligned percent", () => { - const accountsFixture = JSON.parse( - readFileSync(new URL("./fixtures/openai-codex-accounts.json", import.meta.url), "utf-8"), - ) as { accounts: Array<{ accountId: string; email: string; plan: string; refreshToken: string }> }; - const snapshotEntries = JSON.parse( - readFileSync(new URL("./fixtures/codex-status-snapshots.json", import.meta.url), "utf-8"), - ) as Array<[string, any]>; - const snapshots = snapshotEntries.map((entry) => entry[1]); - const account = accountsFixture.accounts[0]!; - const now = snapshots[0]?.updatedAt ?? Date.now(); + it("renders account blocks with aligned percent", () => { + const accountsFixture = JSON.parse( + readFileSync( + new URL("./fixtures/openai-codex-accounts.json", import.meta.url), + "utf-8", + ), + ) as { + accounts: Array<{ + accountId: string; + email: string; + plan: string; + refreshToken: string; + }>; + }; + const snapshotEntries = JSON.parse( + readFileSync( + new URL("./fixtures/codex-status-snapshots.json", import.meta.url), + "utf-8", + ), + ) as Array<[string, any]>; + const snapshots = snapshotEntries.map((entry) => entry[1]); + const account = accountsFixture.accounts[0]!; + const now = snapshots[0]?.updatedAt ?? Date.now(); - const output = renderQuotaReport([account], snapshots, now).join("\n"); - expect(output).toContain("Checking quotas for all accounts..."); - expect(output).toContain(account.email); - expect(output).toContain("Codex CLI Quota"); - expect(output).toContain("GPT-5"); - expect(output).toContain("Weekly"); + const output = renderQuotaReport([account], snapshots, now).join("\n"); + const plainOutput = output.replace(ANSI_REGEX, ""); + expect(plainOutput).toContain("Quota Report"); + expect(plainOutput).toContain(account.email); + expect(plainOutput).toContain(`(${account.plan})`); + expect(plainOutput).toContain("Codex CLI Quota"); + expect(plainOutput).toContain("5h"); + expect(plainOutput).toContain("Weekly"); + expect(plainOutput).toContain("┌"); + expect(plainOutput).toContain("│"); - // Percents aligned to 3 characters. - expect(output).toMatch(/\s\d{1,3}%/); - // Bars use block characters. - expect(output).toMatch(/[█░]{20}/); - }); + // Percents aligned to 3 characters. + expect(plainOutput).toMatch(/\s\d{1,3}%/); + // Bars use block characters. + expect(plainOutput).toMatch(/[█░]{20}/); + }); });