diff --git a/app/electron/main.test.ts b/app/electron/main.test.ts index f0e03a87..07b6daa1 100644 --- a/app/electron/main.test.ts +++ b/app/electron/main.test.ts @@ -31,6 +31,7 @@ function fakeSpawn(result: unknown = { current: { cost: 12.34 } }) { // must spawn. cliStatus is the one channel that resolves without spawning. const CHANNELS = [ 'codeburn:getOverview', + 'codeburn:getQuota', 'codeburn:getPlans', 'codeburn:getActReport', 'codeburn:getModels', @@ -100,14 +101,15 @@ function flattenMenuItems(items: any[]): any[] { } describe('createBridgeHandlers (channel → argv for all channels)', () => { + const deps = (extra = {}) => ({ spawnCli: vi.fn(), spawnCliAction: vi.fn(), resolveCodeburnPath: () => null, getQuota: vi.fn(async () => []), ...extra }) it('exposes exactly the bridge channels', () => { - const handlers = createBridgeHandlers({ spawnCli: vi.fn(), spawnCliAction: vi.fn(), resolveCodeburnPath: () => null }) + const handlers = createBridgeHandlers(deps()) expect(Object.keys(handlers).sort()).toEqual([...CHANNELS].sort()) }) it.each(ARGV_CASES)('$channel with $args spawns the expected argv', async ({ channel, args, argv }) => { const { spawnCli, spawnCliAction, calls } = fakeSpawn() - const handlers = createBridgeHandlers({ spawnCli, spawnCliAction, resolveCodeburnPath: () => '/bin/codeburn' }) + const handlers = createBridgeHandlers(deps({ spawnCli, spawnCliAction, resolveCodeburnPath: () => '/bin/codeburn' })) const res = await handlers[channel]!(...args) expect(calls[0]).toEqual(argv) expect(res).toMatchObject({ ok: true }) @@ -115,7 +117,7 @@ describe('createBridgeHandlers (channel → argv for all channels)', () => { it('codeburn:cliStatus resolves from resolveCodeburnPath without spawning', async () => { const spawnCli = vi.fn() - const handlers = createBridgeHandlers({ spawnCli, spawnCliAction: vi.fn(), resolveCodeburnPath: () => '/opt/homebrew/bin/codeburn' }) + const handlers = createBridgeHandlers(deps({ spawnCli, resolveCodeburnPath: () => '/opt/homebrew/bin/codeburn' })) const res = await handlers['codeburn:cliStatus']!() expect(spawnCli).not.toHaveBeenCalled() expect(res).toEqual({ ok: true, value: { found: true, path: '/opt/homebrew/bin/codeburn' } }) @@ -123,9 +125,21 @@ describe('createBridgeHandlers (channel → argv for all channels)', () => { }) describe('createBridgeHandlers (IPC wiring)', () => { + const withQuota = (value: T) => ({ ...value, getQuota: vi.fn(async () => []) }) + it('returns normalized quota through its own IPC channel and sanitizes unexpected failures', async () => { + const base = { spawnCli: vi.fn(), spawnCliAction: vi.fn(), resolveCodeburnPath: () => null } + const value = [{ provider: 'claude' as const, connection: 'connected' as const, primary: null, details: [], planLabel: 'Pro', footerLines: [] }] + const ok = createBridgeHandlers({ ...base, getQuota: vi.fn(async () => value) }) + expect(await ok['codeburn:getQuota']!()).toEqual({ ok: true, value }) + + const failed = createBridgeHandlers({ ...base, getQuota: vi.fn(async () => { throw new Error('Bearer secret sk-ant-leak') }) }) + const result = await failed['codeburn:getQuota']!() + expect(result).toMatchObject({ ok: false, error: { kind: 'nonzero' } }) + expect(JSON.stringify(result)).not.toMatch(/secret|sk-ant-leak/) + }) it('getOverview spawns menubar-json for the period, omitting --provider for "all"', async () => { const { spawnCli, spawnCliAction, calls } = fakeSpawn() - const handlers = createBridgeHandlers({ spawnCli, spawnCliAction, resolveCodeburnPath: () => '/bin/codeburn' }) + const handlers = createBridgeHandlers(withQuota({ spawnCli, spawnCliAction, resolveCodeburnPath: () => '/bin/codeburn' })) const res = await handlers['codeburn:getOverview']!('30days', 'all') expect(calls[0]).toEqual(['status', '--format', 'menubar-json', '--period', '30days']) expect(res).toEqual({ ok: true, value: { current: { cost: 12.34 } } }) @@ -133,7 +147,7 @@ describe('createBridgeHandlers (IPC wiring)', () => { it('adds --provider and --by-task when requested', async () => { const { spawnCli, spawnCliAction, calls } = fakeSpawn([]) - const handlers = createBridgeHandlers({ spawnCli, spawnCliAction, resolveCodeburnPath: () => null }) + const handlers = createBridgeHandlers(withQuota({ spawnCli, spawnCliAction, resolveCodeburnPath: () => null })) await handlers['codeburn:getModels']!('week', 'claude', true) expect(calls[0]).toEqual(['models', '--format', 'json', '--period', 'week', '--provider', 'claude', '--by-task']) }) @@ -142,17 +156,17 @@ describe('createBridgeHandlers (IPC wiring)', () => { const spawnCli = vi.fn(async () => { throw new CliError('nonzero', 'boom') }) - const handlers = createBridgeHandlers({ spawnCli, spawnCliAction: vi.fn(), resolveCodeburnPath: () => '/bin/codeburn' }) + const handlers = createBridgeHandlers(withQuota({ spawnCli, spawnCliAction: vi.fn(), resolveCodeburnPath: () => '/bin/codeburn' })) const res = await handlers['codeburn:getYield']!('today') expect(res).toEqual({ ok: false, error: { kind: 'nonzero', message: 'boom' } }) }) it('cliStatus reports the resolved binary path', async () => { - const handlers = createBridgeHandlers({ + const handlers = createBridgeHandlers(withQuota({ spawnCli: vi.fn(), spawnCliAction: vi.fn(), resolveCodeburnPath: () => '/opt/homebrew/bin/codeburn', - }) + })) const res = await handlers['codeburn:cliStatus']!() expect(res).toEqual({ ok: true, value: { found: true, path: '/opt/homebrew/bin/codeburn' } }) }) diff --git a/app/electron/main.ts b/app/electron/main.ts index ba0fedee..6b9ebbb5 100644 --- a/app/electron/main.ts +++ b/app/electron/main.ts @@ -2,6 +2,7 @@ import { app, BrowserWindow, dialog, ipcMain, Menu, nativeTheme, shell, type Men import path from 'node:path' import { CliError, resolveCodeburnPath, spawnCli, spawnCliAction, type ActionResult } from './cli' +import { getQuota, sanitizeError } from './quota' // Result envelope: handlers never throw across IPC so the structured error // `kind` survives contextBridge serialization. preload.ts unwraps it. @@ -26,6 +27,7 @@ type Deps = { spawnCli: (args: string[], opts?: { timeoutMs?: number }) => Promise spawnCliAction: (args: string[], opts?: { timeoutMs?: number }) => Promise resolveCodeburnPath: () => string | null + getQuota: typeof getQuota } type Handler = (...args: any[]) => Promise @@ -35,7 +37,7 @@ type Handler = (...args: any[]) => Promise * shell) and returns a result envelope. Pure + injectable so the wiring is * unit-testable without launching Electron. */ -export function createBridgeHandlers(deps: Deps = { spawnCli, spawnCliAction, resolveCodeburnPath }): Record { +export function createBridgeHandlers(deps: Deps = { spawnCli, spawnCliAction, resolveCodeburnPath, getQuota }): Record { const run = (build: (...args: any[]) => string[]): Handler => async (...args: any[]) => { try { return { ok: true, value: await deps.spawnCli(build(...args)) } @@ -52,6 +54,10 @@ export function createBridgeHandlers(deps: Deps = { spawnCli, spawnCliAction, re } return { + 'codeburn:getQuota': async () => { + try { return { ok: true, value: await deps.getQuota() } } + catch (error) { return { ok: false, error: { kind: 'nonzero', message: sanitizeError(error) } } } + }, 'codeburn:getOverview': run((period: string, provider: string, range?: DateRange) => [ 'status', '--format', 'menubar-json', '--period', period, ...providerArgs(provider), ...rangeArgs(range), ]), diff --git a/app/electron/preload.ts b/app/electron/preload.ts index 16e0b6be..b8af3628 100644 --- a/app/electron/preload.ts +++ b/app/electron/preload.ts @@ -18,6 +18,7 @@ async function invoke(channel: string, ...args: unknown[]): Promise { // Shape matches CodeburnBridge (app/renderer/lib/types.ts); typing is enforced // renderer-side where `window.codeburn` is declared as CodeburnBridge. const bridge = { + getQuota: () => invoke('codeburn:getQuota'), getOverview: (period: string, provider: string, range?: DateRange) => invoke('codeburn:getOverview', period, provider, range), getPlans: (period: string) => invoke('codeburn:getPlans', period), getActReport: () => invoke('codeburn:getActReport'), diff --git a/app/electron/quota/claude.test.ts b/app/electron/quota/claude.test.ts new file mode 100644 index 00000000..356ba002 --- /dev/null +++ b/app/electron/quota/claude.test.ts @@ -0,0 +1,81 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { decodeClaudeUsage, fetchClaudeQuota } from './claude' + +const credential = JSON.stringify({ + claudeAiOauth: { + accessToken: 'sk-ant-test-secret', + refreshToken: 'unused', + expiresAt: Date.now() + 86_400_000, + rateLimitTier: 'max_20x', + }, +}) + +afterEach(() => vi.restoreAllMocks()) + +describe('Claude quota', () => { + it('decodes ordered five-hour, weekly, model, and scoped windows with credential tier', () => { + const quota = decodeClaudeUsage({ + five_hour: { utilization: 25, resets_at: '2026-07-12T12:00:00Z' }, + seven_day: { utilization: 50, resets_at: '2026-07-19T12:00:00.123Z' }, + seven_day_opus: { utilization: 75, resets_at: '2026-07-19T12:00:00Z' }, + seven_day_sonnet: { utilization: 90, resets_at: '2026-07-19T12:00:00Z' }, + limits: [ + { kind: 'weekly_all', percent: 88, scope: { model: { display_name: 'Duplicate' } } }, + { kind: 'weekly_scoped', percent: 10, resets_at: '2026-07-20T00:00:00Z', scope: { model: { display_name: 'Haiku' } } }, + ], + }, { accessToken: 'hidden', rateLimitTier: 'max_20x' }) + + expect(quota.connection).toBe('connected') + expect(quota.planLabel).toBe('Max 20x') + expect(quota.primary?.label).toBe('Weekly') + expect(quota.details.map(row => row.label)).toEqual(['5-hour', 'Weekly', 'Weekly · Opus', 'Weekly · Sonnet', 'Weekly · Haiku']) + expect(quota.details.map(row => row.percent)).toEqual([0.25, 0.5, 0.75, 0.9, 0.1]) + }) + + it('returns disconnected without credentials and never fetches', async () => { + const fetchMock = vi.fn() + const result = await fetchClaudeQuota({ fetch: fetchMock, readFile: vi.fn(async () => null) }) + expect(result.quota.connection).toBe('disconnected') + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('sanitizes newline-corrupted credential JSON and uses exact request headers', async () => { + const broken = credential.replace('sk-ant-test-secret', 'sk-ant-test-\n secret') + const fetchMock = vi.fn(async () => new Response(JSON.stringify({ seven_day: { utilization: 4, resets_at: '2026-07-19T00:00:00Z' } }), { status: 200 })) + const result = await fetchClaudeQuota({ fetch: fetchMock, readFile: vi.fn(async () => broken) }) + expect(result.quota.connection).toBe('connected') + const [url, init] = fetchMock.mock.calls[0]! as unknown as [string, RequestInit] + expect(url).toBe('https://api.anthropic.com/api/oauth/usage') + expect(init.method).toBe('GET') + expect(init.headers).toEqual({ + Authorization: 'Bearer sk-ant-test-secret', Accept: 'application/json', + 'anthropic-beta': 'oauth-2025-04-20', 'User-Agent': 'claude-code/2.1.0', + }) + expect(init.headers).not.toHaveProperty('anthropic-version') + }) + + it('uses the body retry_after for 429 backoff', async () => { + const fetchMock = vi.fn(async () => new Response(JSON.stringify({ retry_after: '42' }), { status: 429 })) + const result = await fetchClaudeQuota({ fetch: fetchMock, readFile: vi.fn(async () => credential) }) + expect(result).toMatchObject({ quota: { connection: 'transientFailure' }, retryAfterSeconds: 60 }) + }) + + it('never calls an Anthropic refresh endpoint when the token is unchanged after 401', async () => { + const fetchMock = vi.fn(async () => new Response('', { status: 401 })) + const result = await fetchClaudeQuota({ fetch: fetchMock, readFile: vi.fn(async () => credential) }) + expect(result.quota.connection).toBe('transientFailure') + expect(fetchMock).toHaveBeenCalledTimes(1) + expect((fetchMock.mock.calls as unknown as Array<[string]>).every(call => String(call[0]) === 'https://api.anthropic.com/api/oauth/usage')).toBe(true) + }) + + it('redacts tokens and NUL from diagnostics without surfacing them', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const fetchMock = vi.fn(async () => { throw new Error('Bearer rawbearer sk-ant-leak sk-other eyJabc.def.ghi\0tail') }) + const result = await fetchClaudeQuota({ fetch: fetchMock, readFile: vi.fn(async () => credential) }) + const logged = warn.mock.calls.flat().join(' ') + expect(result.quota).not.toHaveProperty('error') + expect(logged).not.toMatch(/rawbearer|sk-ant-leak|sk-other|eyJabc|\0/) + expect(logged).toContain('[REDACTED]') + }) +}) diff --git a/app/electron/quota/claude.ts b/app/electron/quota/claude.ts new file mode 100644 index 00000000..be1c333a --- /dev/null +++ b/app/electron/quota/claude.ts @@ -0,0 +1,158 @@ +import { execFile } from 'node:child_process' +import os from 'node:os' +import path from 'node:path' +import { promisify } from 'node:util' + +import { fraction, quotaRequestSignal, readSecureFile, sanitizeError } from './security' +import type { QuotaProvider, QuotaWindow } from './types' + +const execFileAsync = promisify(execFile) +const ENDPOINT = 'https://api.anthropic.com/api/oauth/usage' + +type ClaudeCredential = { accessToken: string; expiresAt?: number; rateLimitTier?: string } +export type ClaudeDeps = { + fetch: typeof fetch + credentialPath: string + readFile: typeof readSecureFile + now: () => number + keychain?: () => Promise +} + +const defaults: ClaudeDeps = { + fetch: globalThis.fetch, + credentialPath: path.join(os.homedir(), '.claude', '.credentials.json'), + readFile: readSecureFile, + now: Date.now, +} + +function empty(connection: QuotaProvider['connection']): QuotaProvider { + return { provider: 'claude', connection, primary: null, details: [], planLabel: null, footerLines: [] } +} + +function parseCredential(raw: string): ClaudeCredential | null { + const clean = raw.replace(/\r/g, '').replace(/\n[ \t]*/g, '') + const oauth = (JSON.parse(clean) as { claudeAiOauth?: Record }).claudeAiOauth + if (!oauth || typeof oauth.accessToken !== 'string' || oauth.accessToken.length === 0) return null + return { + accessToken: oauth.accessToken, + expiresAt: typeof oauth.expiresAt === 'number' ? oauth.expiresAt : undefined, + rateLimitTier: typeof oauth.rateLimitTier === 'string' ? oauth.rateLimitTier : undefined, + } +} + +async function credentialFromFile(deps: ClaudeDeps): Promise { + const raw = await deps.readFile(deps.credentialPath, 64 * 1024) + return raw ? parseCredential(raw) : null +} + +export async function readClaudeKeychain(): Promise { + if (process.platform !== 'darwin') return null + const args = ['find-generic-password', '-s', 'Claude Code-credentials'] + const user = process.env.USER + const attempts = user ? [[...args, '-a', user, '-w'], [...args, '-w']] : [[...args, '-w']] + for (const argv of attempts) { + try { + const { stdout } = await execFileAsync('/usr/bin/security', argv, { timeout: 10_000, maxBuffer: 64 * 1024 }) + if (stdout) return stdout + } catch { /* File fallback absence is a normal disconnected state. */ } + } + return null +} + +function windowOf(label: string, value: unknown): QuotaWindow | null { + if (!value || typeof value !== 'object') return null + const row = value as Record + const percent = fraction(row.utilization) + if (percent === null) return null + const resetsAt = typeof row.resets_at === 'string' && !Number.isNaN(Date.parse(row.resets_at)) + ? new Date(row.resets_at).toISOString() : null + return { label, percent, resetsAt } +} + +function tierLabel(raw: string | undefined): string { + const value = raw?.toLowerCase() ?? '' + if (value.includes('max_20x') || value.includes('max20x') || value.includes('max-20x')) return 'Max 20x' + if (value.includes('max_5x') || value.includes('max5x') || value.includes('max-5x') || value.includes('max')) return 'Max 5x' + if (value.includes('pro')) return 'Pro' + if (value.includes('team')) return 'Team' + if (value.includes('enterprise')) return 'Enterprise' + return 'Subscription' +} + +export function decodeClaudeUsage(body: unknown, credential: ClaudeCredential): QuotaProvider { + const data = body && typeof body === 'object' ? body as Record : {} + const five = windowOf('5-hour', data.five_hour) + const weekly = windowOf('Weekly', data.seven_day) + const opus = windowOf('Weekly · Opus', data.seven_day_opus) + const sonnet = windowOf('Weekly · Sonnet', data.seven_day_sonnet) + const scoped: QuotaWindow[] = [] + if (Array.isArray(data.limits)) { + for (const item of data.limits) { + if (!item || typeof item !== 'object') continue + const row = item as Record + const display = row.scope?.model?.display_name + const percent = fraction(row.percent) + if (row.kind !== 'weekly_scoped' || typeof display !== 'string' || percent === null) continue + const resetsAt = typeof row.resets_at === 'string' && !Number.isNaN(Date.parse(row.resets_at)) + ? new Date(row.resets_at).toISOString() : null + scoped.push({ label: `Weekly · ${display}`, percent, resetsAt }) + } + } + return { + provider: 'claude', connection: 'connected', primary: weekly, + details: [five, weekly, opus, sonnet].filter((row): row is QuotaWindow => row !== null).concat(scoped), + planLabel: tierLabel(credential.rateLimitTier), footerLines: [], + } +} + +async function request(token: string, deps: ClaudeDeps, parent?: AbortSignal): Promise { + return deps.fetch(ENDPOINT, { + method: 'GET', signal: quotaRequestSignal(parent), + headers: { + Authorization: `Bearer ${token}`, + Accept: 'application/json', + 'anthropic-beta': 'oauth-2025-04-20', + 'User-Agent': 'claude-code/2.1.0', + }, + }) +} + +export type ClaudeResult = { quota: QuotaProvider; retryAfterSeconds?: number } + +export async function fetchClaudeQuota(options: Partial & { signal?: AbortSignal; allowKeychain?: boolean } = {}): Promise { + const deps = { ...defaults, ...options } + try { + let credential = await credentialFromFile(deps) + if (!credential && options.allowKeychain && process.platform === 'darwin') { + const raw = await (deps.keychain ?? readClaudeKeychain)() + credential = raw ? parseCredential(raw) : null + } + if (!credential) return { quota: empty('disconnected') } + + let response: Response + if (credential.expiresAt !== undefined && credential.expiresAt - deps.now() <= 5 * 60_000) { + const reread = await credentialFromFile(deps) + if (!reread || reread.accessToken === credential.accessToken) return { quota: empty('transientFailure') } + credential = reread + } + response = await request(credential.accessToken, deps, options.signal) + if (response.status === 401) { + const reread = await credentialFromFile(deps) + if (!reread || reread.accessToken === credential.accessToken) return { quota: empty('transientFailure') } + credential = reread + response = await request(credential.accessToken, deps, options.signal) + } + if (response.status === 429) { + let hint: unknown + try { hint = (await response.json() as Record).retry_after } catch { hint = undefined } + const parsed = typeof hint === 'number' ? hint : typeof hint === 'string' ? Number(hint) : NaN + return { quota: empty('transientFailure'), retryAfterSeconds: Math.max(Number.isFinite(parsed) ? parsed : 300, 60) } + } + if (!response.ok) return { quota: empty(response.status >= 400 && response.status < 500 ? 'terminalFailure' : 'transientFailure') } + return { quota: decodeClaudeUsage(await response.json(), credential) } + } catch (error) { + // Deliberately sanitize before the only diagnostic sink. Tokens are never returned. + console.warn(`Claude quota unavailable: ${sanitizeError(error)}`) + return { quota: empty('transientFailure') } + } +} diff --git a/app/electron/quota/codex.test.ts b/app/electron/quota/codex.test.ts new file mode 100644 index 00000000..81694cd9 --- /dev/null +++ b/app/electron/quota/codex.test.ts @@ -0,0 +1,68 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { decodeCodexUsage, fetchCodexQuota } from './codex' + +const now = Date.parse('2026-07-12T00:00:00Z') +const auth = { + auth_mode: 'chatgpt', OPENAI_API_KEY: 'preserve-me', last_refresh: '2026-07-11T00:00:00Z', + tokens: { access_token: 'eyJaccess.token.sig', refresh_token: 'refresh-secret', id_token: 'old-id', account_id: 'acct_1' }, +} + +afterEach(() => vi.restoreAllMocks()) + +describe('Codex quota', () => { + it('decodes primary/secondary/additional windows, plan and numeric-string credits', () => { + const quota = decodeCodexUsage({ + plan_type: 'pLuS', + rate_limit: { + primary_window: { used_percent: 20, reset_at: 1_800_000_000, limit_window_seconds: 18_000 }, + secondary_window: { used_percent: 80, reset_at: 1_800_100_000, limit_window_seconds: 604_800 }, + }, + additional_rate_limits: [{ + limit_name: 'GPT-5', rate_limit: { + primary_window: { used_percent: 12, reset_at: 1_800_000_000, limit_window_seconds: 3600 }, + secondary_window: { used_percent: 0, reset_at: 1_800_000_000, limit_window_seconds: 86_400 }, + }, + }], + credits: { balance: '3.5' }, + }) + expect(quota.planLabel).toBe('Plus') + expect(quota.primary?.label).toBe('5-hour') + expect(quota.details.map(row => row.label)).toEqual(['5-hour', 'Weekly', 'GPT-5 · Hour']) + expect(quota.footerLines).toEqual(['Credits remaining · $3.50']) + }) + + it('promotes secondary when primary is absent', () => { + const quota = decodeCodexUsage({ rate_limit: { secondary_window: { used_percent: 9, reset_at: 1_800_000_000, limit_window_seconds: 604_800 } } }) + expect(quota.primary?.label).toBe('Weekly') + expect(quota.details).toHaveLength(1) + }) + + it('returns disconnected without credentials', async () => { + const fetchMock = vi.fn() + const result = await fetchCodexQuota({ fetch: fetchMock, readFile: vi.fn(async () => null) }) + expect(result.quota.connection).toBe('disconnected') + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('sends account id and uses Retry-After header for 429', async () => { + const fetchMock = vi.fn(async () => new Response('', { status: 429, headers: { 'Retry-After': '120' } })) + const result = await fetchCodexQuota({ fetch: fetchMock, readFile: vi.fn(async () => JSON.stringify(auth)), now: () => now }) + expect(result.retryAfterSeconds).toBe(120) + const usageInit = (fetchMock.mock.calls[0]! as unknown as [string, RequestInit])[1] + expect(usageInit.headers).toMatchObject({ 'ChatGPT-Account-Id': 'acct_1', 'User-Agent': 'CodeBurn' }) + }) + + it('refreshes after eight days and preserves unrelated auth keys on write-back', async () => { + const stale = { ...auth, last_refresh: '2026-07-01T00:00:00Z' } + const fetchMock = vi.fn(async (url: string) => url.includes('/oauth/token') + ? new Response(JSON.stringify({ access_token: 'new-access', refresh_token: 'new-refresh', id_token: 'new-id' }), { status: 200 }) + : new Response(JSON.stringify({ plan_type: 'pro', rate_limit: {} }), { status: 200 })) + const writeFile = vi.fn(async () => undefined) + await fetchCodexQuota({ fetch: fetchMock as typeof fetch, readFile: vi.fn(async () => JSON.stringify(stale)), writeFile, now: () => now }) + const saved = JSON.parse((writeFile.mock.calls[0]! as unknown as [string, string])[1]) + expect(saved.OPENAI_API_KEY).toBe('preserve-me') + expect(saved.tokens).toMatchObject({ access_token: 'new-access', refresh_token: 'new-refresh', id_token: 'new-id', account_id: 'acct_1' }) + expect((fetchMock.mock.calls[0]! as unknown as [string, RequestInit])[1].method).toBe('POST') + }) +}) diff --git a/app/electron/quota/codex.ts b/app/electron/quota/codex.ts new file mode 100644 index 00000000..f81fff5c --- /dev/null +++ b/app/electron/quota/codex.ts @@ -0,0 +1,177 @@ +import os from 'node:os' +import path from 'node:path' + +import { atomicWriteSecureFile, fraction, quotaRequestSignal, readSecureFile, sanitizeError } from './security' +import type { QuotaProvider, QuotaWindow } from './types' + +const USAGE_ENDPOINT = 'https://chatgpt.com/backend-api/wham/usage' +const TOKEN_ENDPOINT = 'https://auth.openai.com/oauth/token' +const CLIENT_ID = 'app_EMoamEEZ73f0CkXaXp7hrann' +const EIGHT_DAYS = 8 * 24 * 60 * 60_000 + +type AuthDoc = Record & { + auth_mode?: string + tokens?: { access_token?: string; refresh_token?: string; id_token?: string; account_id?: string; [key: string]: unknown } + last_refresh?: string +} + +export type CodexDeps = { + fetch: typeof fetch + authPath: string + readFile: typeof readSecureFile + writeFile: typeof atomicWriteSecureFile + now: () => number +} + +const defaults: CodexDeps = { + fetch: globalThis.fetch, + authPath: path.join(os.homedir(), '.codex', 'auth.json'), + readFile: readSecureFile, + writeFile: atomicWriteSecureFile, + now: Date.now, +} + +function empty(connection: QuotaProvider['connection']): QuotaProvider { + return { provider: 'codex', connection, primary: null, details: [], planLabel: null, footerLines: [] } +} + +async function readAuth(deps: CodexDeps): Promise { + const raw = await deps.readFile(deps.authPath, 64 * 1024) + return raw ? JSON.parse(raw) as AuthDoc : null +} + +function labelForSeconds(value: unknown): string { + const seconds = typeof value === 'number' ? Math.max(0, Math.trunc(value)) : 0 + if (seconds < 3600) return 'Hourly' + if (seconds < 7200) return 'Hour' + if (seconds >= 18_000 && seconds < 19_000) return '5-hour' + if (seconds >= 86_400 && seconds < 87_000) return 'Daily' + if (seconds >= 604_800 && seconds < 605_000) return 'Weekly' + const hours = Math.floor(seconds / 3600) + return hours < 24 ? `${hours}-hour` : `${Math.floor(hours / 24)}-day` +} + +function windowOf(value: unknown, override?: string): QuotaWindow | null { + if (!value || typeof value !== 'object') return null + const row = value as Record + const percent = fraction(row.used_percent) + if (percent === null) return null + const reset = typeof row.reset_at === 'number' && Number.isFinite(row.reset_at) + ? new Date(row.reset_at * 1000).toISOString() : null + return { label: override ?? labelForSeconds(row.limit_window_seconds), percent, resetsAt: reset } +} + +function planLabel(value: unknown): string | null { + if (typeof value !== 'string' || !value.trim()) return null + const raw = value.trim() + const lower = raw.toLowerCase() + const known: Record = { + guest: 'Guest', free: 'Free', go: 'Go', plus: 'Plus', pro: 'Pro', + prolite: 'Pro Lite', pro_lite: 'Pro Lite', 'pro-lite': 'Pro Lite', + free_workspace: 'Free Workspace', team: 'Team', business: 'Business', + education: 'Education', quorum: 'Quorum', k12: 'K-12', enterprise: 'Enterprise', edu: 'Edu', + } + return known[lower] ?? lower.replace(/(^|[_-])\w/g, match => match.replace(/[_-]/, ' ').toUpperCase()) +} + +export function decodeCodexUsage(body: unknown): QuotaProvider { + const data = body && typeof body === 'object' ? body as Record : {} + const primaryRaw = windowOf(data.rate_limit?.primary_window) + const secondaryRaw = windowOf(data.rate_limit?.secondary_window) + const primary = primaryRaw ?? secondaryRaw + const details: QuotaWindow[] = [] + if (primaryRaw) details.push(primaryRaw) + if (secondaryRaw && secondaryRaw !== primary) details.push(secondaryRaw) + else if (!primaryRaw && secondaryRaw) details.push(secondaryRaw) + if (Array.isArray(data.additional_rate_limits)) { + for (const additional of data.additional_rate_limits) { + if (!additional || typeof additional !== 'object' || typeof additional.limit_name !== 'string') continue + for (const key of ['primary_window', 'secondary_window'] as const) { + const raw = additional.rate_limit?.[key] + const base = windowOf(raw) + if (base && base.percent > 0) details.push({ ...base, label: `${additional.limit_name} · ${base.label}` }) + } + } + } + const rawBalance = data.credits?.balance + const balance = typeof rawBalance === 'number' ? rawBalance : typeof rawBalance === 'string' ? Number(rawBalance) : NaN + return { + provider: 'codex', connection: 'connected', primary, details, + planLabel: planLabel(data.plan_type), + footerLines: Number.isFinite(balance) && balance > 0 ? [`Credits remaining · $${balance.toFixed(2)}`] : [], + } +} + +async function refresh(auth: AuthDoc, deps: CodexDeps, signal?: AbortSignal): Promise { + const refreshToken = auth.tokens?.refresh_token + if (!refreshToken) return null + const response = await deps.fetch(TOKEN_ENDPOINT, { + method: 'POST', signal: quotaRequestSignal(signal), + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ client_id: CLIENT_ID, grant_type: 'refresh_token', refresh_token: refreshToken, scope: 'openid profile email' }), + }) + if (!response.ok) return null + const next = await response.json() as Record + if (typeof next.access_token !== 'string' || !next.access_token) return null + const latest = await readAuth(deps) + if (!latest || latest.auth_mode !== 'chatgpt') return null + latest.tokens = { + ...latest.tokens, + access_token: next.access_token, + ...(typeof next.refresh_token === 'string' ? { refresh_token: next.refresh_token } : {}), + ...(typeof next.id_token === 'string' ? { id_token: next.id_token } : {}), + } + latest.last_refresh = new Date(deps.now()).toISOString() + await deps.writeFile(deps.authPath, `${JSON.stringify(latest, null, 2)}\n`) + return latest +} + +async function usage(auth: AuthDoc, deps: CodexDeps, signal?: AbortSignal): Promise { + const token = auth.tokens?.access_token + if (!token) return null + const headers: Record = { Authorization: `Bearer ${token}`, Accept: 'application/json', 'User-Agent': 'CodeBurn' } + if (auth.tokens?.account_id) headers['ChatGPT-Account-Id'] = auth.tokens.account_id + return deps.fetch(USAGE_ENDPOINT, { method: 'GET', headers, signal: quotaRequestSignal(signal) }) +} + +export type CodexResult = { quota: QuotaProvider; retryAfterSeconds?: number } + +export async function fetchCodexQuota(options: Partial & { signal?: AbortSignal } = {}): Promise { + const deps = { ...defaults, ...options } + try { + let auth = await readAuth(deps) + if (!auth) return { quota: empty('disconnected') } + if (auth.auth_mode !== 'chatgpt') return { quota: empty('terminalFailure') } + if (!auth.tokens?.access_token) return { quota: empty('disconnected') } + + const refreshedAt = typeof auth.last_refresh === 'string' ? Date.parse(auth.last_refresh) : NaN + if (!Number.isFinite(refreshedAt) || deps.now() - refreshedAt > EIGHT_DAYS) { + const next = await refresh(auth, deps, options.signal) + if (next) auth = next + } + let response = await usage(auth, deps, options.signal) + if (!response) return { quota: empty('disconnected') } + if (response.status === 401) { + const reread = await readAuth(deps) + if (reread?.tokens?.access_token && reread.tokens.access_token !== auth.tokens?.access_token) auth = reread + else { + const next = await refresh(reread ?? auth, deps, options.signal) + if (!next) return { quota: empty('transientFailure') } + auth = next + } + response = await usage(auth, deps, options.signal) + if (!response) return { quota: empty('transientFailure') } + } + if (response.status === 429) { + const raw = response.headers.get('Retry-After') + let seconds = raw === null ? NaN : Number(raw) + if (!Number.isFinite(seconds) && raw) seconds = (Date.parse(raw) - deps.now()) / 1000 + return { quota: empty('transientFailure'), retryAfterSeconds: Math.max(Number.isFinite(seconds) ? Math.ceil(seconds) : 300, 60) } + } + if (!response.ok) return { quota: empty(response.status >= 400 && response.status < 500 ? 'terminalFailure' : 'transientFailure') } + return { quota: decodeCodexUsage(await response.json()) } + } catch (error) { + console.warn(`Codex quota unavailable: ${sanitizeError(error)}`) + return { quota: empty('transientFailure') } + } +} diff --git a/app/electron/quota/index.test.ts b/app/electron/quota/index.test.ts new file mode 100644 index 00000000..cf24c44e --- /dev/null +++ b/app/electron/quota/index.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it, vi } from 'vitest' + +import { QuotaService } from './index' +import type { QuotaProvider } from './types' + +const quota = (provider: 'claude' | 'codex'): QuotaProvider => ({ + provider, connection: 'connected', primary: null, details: [], planLabel: null, footerLines: [], +}) + +describe('QuotaService', () => { + it('persists provider 429 blocked-until and gates the next forced fetch', async () => { + const writes: string[] = [] + const claude = vi.fn(async () => ({ quota: quota('claude'), retryAfterSeconds: 60 })) + const codex = vi.fn(async () => ({ quota: quota('codex') })) + const service = new QuotaService({ + claude, codex, now: () => Date.parse('2026-07-12T00:00:00Z'), + readFile: vi.fn(async () => writes.at(-1) ?? null), + writeFile: vi.fn(async (_path, value) => { writes.push(value) }), + statePath: '/mock/backoff.json', + }) + await service.getQuota({ force: true }) + const saved = JSON.parse(writes[0]!) + expect(saved.claude).toBe('2026-07-12T00:01:00.000Z') + await service.getQuota({ force: true }) + expect(claude).toHaveBeenCalledTimes(1) + expect(codex).toHaveBeenCalledTimes(2) + }) + + it('single-flights simultaneous callers', async () => { + let release!: () => void + const pending = new Promise(resolve => { release = resolve }) + const claude = vi.fn(async () => { await pending; return { quota: quota('claude') } }) + const service = new QuotaService({ + claude, codex: vi.fn(async () => ({ quota: quota('codex') })), + readFile: vi.fn(async () => null), writeFile: vi.fn(async () => undefined), + }) + const first = service.getQuota({ force: true }) + const second = service.getQuota({ force: true }) + release() + expect(await first).toEqual(await second) + expect(claude).toHaveBeenCalledTimes(1) + }) +}) + diff --git a/app/electron/quota/index.ts b/app/electron/quota/index.ts new file mode 100644 index 00000000..01112269 --- /dev/null +++ b/app/electron/quota/index.ts @@ -0,0 +1,114 @@ +import os from 'node:os' +import path from 'node:path' + +import { fetchClaudeQuota } from './claude' +import { fetchCodexQuota } from './codex' +import { atomicWriteSecureFile, readSecureFile, sanitizeError } from './security' +import type { ProviderName, QuotaProvider } from './types' + +export type { QuotaProvider, QuotaWindow } from './types' +export { sanitizeError } from './security' + +type Blocked = Partial> +type FetchResult = { quota: QuotaProvider; retryAfterSeconds?: number } +type QuotaDeps = { + claude: (options: { signal: AbortSignal; allowKeychain: boolean }) => Promise + codex: (options: { signal: AbortSignal }) => Promise + statePath: string + readFile: typeof readSecureFile + writeFile: typeof atomicWriteSecureFile + now: () => number + refreshMs: number +} + +const defaultDeps: QuotaDeps = { + claude: fetchClaudeQuota, + codex: fetchCodexQuota, + statePath: path.join(os.homedir(), '.codeburn', 'quota-backoff.json'), + readFile: readSecureFile, + writeFile: atomicWriteSecureFile, + now: Date.now, + refreshMs: 2 * 60_000, +} + +function unavailable(provider: ProviderName, connection: QuotaProvider['connection']): QuotaProvider { + return { provider, connection, primary: null, details: [], planLabel: null, footerLines: [] } +} + +export class QuotaService { + private readonly deps: QuotaDeps + private cache: { at: number; value: QuotaProvider[] } | null = null + private flight: Promise | null = null + private generations: Record = { claude: 0, codex: 0 } + private controllers: Partial> = {} + + constructor(deps: Partial = {}) { this.deps = { ...defaultDeps, ...deps } } + + invalidate(provider: ProviderName): void { + this.generations[provider] += 1 + this.controllers[provider]?.abort() + this.controllers[provider] = undefined + this.cache = null + } + + async getQuota(options: { force?: boolean; allowClaudeKeychain?: boolean } = {}): Promise { + if (!options.force && this.cache && this.deps.now() - this.cache.at < this.deps.refreshMs) return this.cache.value + if (this.flight) return this.flight + this.flight = this.fetchAll(Boolean(options.allowClaudeKeychain)).finally(() => { this.flight = null }) + return this.flight + } + + private async readBlocked(): Promise { + try { + const raw = await this.deps.readFile(this.deps.statePath, 16 * 1024) + return raw ? JSON.parse(raw) as Blocked : {} + } catch (error) { + console.warn(`Quota backoff state unavailable: ${sanitizeError(error)}`) + return {} + } + } + + private async writeBlocked(blocked: Blocked): Promise { + try { await this.deps.writeFile(this.deps.statePath, `${JSON.stringify(blocked, null, 2)}\n`) } + catch (error) { console.warn(`Quota backoff state not saved: ${sanitizeError(error)}`) } + } + + private async fetchAll(allowClaudeKeychain: boolean): Promise { + const startingGenerations = { ...this.generations } + const prior = this.cache?.value ?? [] + const blocked = await this.readBlocked() + const run = async (provider: ProviderName): Promise => { + const retainOnFailure = (next: QuotaProvider): QuotaProvider => { + if (next.connection !== 'transientFailure') return next + const previous = prior.find(item => item.provider === provider) + return previous?.connection === 'connected' ? { ...previous, connection: 'transientFailure' } : next + } + const until = blocked[provider] ? Date.parse(blocked[provider]!) : NaN + if (Number.isFinite(until) && until > this.deps.now()) return retainOnFailure(unavailable(provider, 'transientFailure')) + const generation = this.generations[provider] + const controller = new AbortController() + this.controllers[provider] = controller + const result = provider === 'claude' + ? await this.deps.claude({ signal: controller.signal, allowKeychain: allowClaudeKeychain }) + : await this.deps.codex({ signal: controller.signal }) + if (generation !== this.generations[provider] || controller.signal.aborted) return unavailable(provider, 'disconnected') + if (result.retryAfterSeconds !== undefined) { + blocked[provider] = new Date(this.deps.now() + result.retryAfterSeconds * 1000).toISOString() + await this.writeBlocked(blocked) + } else if (blocked[provider]) { + delete blocked[provider] + await this.writeBlocked(blocked) + } + if (this.controllers[provider] === controller) this.controllers[provider] = undefined + return retainOnFailure(result.quota) + } + const value = await Promise.all([run('claude'), run('codex')]) + if (startingGenerations.claude === this.generations.claude && startingGenerations.codex === this.generations.codex) { + this.cache = { at: this.deps.now(), value } + } + return value + } +} + +export const quotaService = new QuotaService() +export const getQuota = (): Promise => quotaService.getQuota() diff --git a/app/electron/quota/security.ts b/app/electron/quota/security.ts new file mode 100644 index 00000000..f569a8ed --- /dev/null +++ b/app/electron/quota/security.ts @@ -0,0 +1,77 @@ +import { constants, type Stats } from 'node:fs' +import { chmod, lstat, mkdir, open, rename, unlink } from 'node:fs/promises' +import path from 'node:path' + +const NOFOLLOW = constants.O_NOFOLLOW ?? 0 + +export function sanitizeError(error: unknown): string { + const raw = error instanceof Error ? error.message : String(error) + return raw + .replace(/\0/g, '') + .replace(/Bearer\s+[^\s,;"']+/gi, 'Bearer [REDACTED]') + .replace(/sk-ant-[A-Za-z0-9_-]+/gi, '[REDACTED]') + .replace(/sk-[A-Za-z0-9_-]+/gi, '[REDACTED]') + .replace(/eyJ[A-Za-z0-9._-]+/g, '[REDACTED]') + .slice(0, 240) +} + +function assertSafeMode(stats: Stats, filePath: string): void { + if (!stats.isFile()) throw new Error(`Credential path is not a regular file: ${filePath}`) + if ((stats.mode & 0o077) !== 0) throw new Error(`Credential file permissions are too broad: ${filePath}`) +} + +export async function readSecureFile(filePath: string, maxBytes = 64 * 1024): Promise { + let before: Stats + try { + before = await lstat(filePath) + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') return null + throw error + } + if (before.isSymbolicLink()) throw new Error(`Refusing symbolic link: ${filePath}`) + assertSafeMode(before, filePath) + if (before.size > maxBytes) throw new Error(`Credential file exceeds ${maxBytes} bytes: ${filePath}`) + + const handle = await open(filePath, constants.O_RDONLY | NOFOLLOW) + try { + const after = await handle.stat() + assertSafeMode(after, filePath) + if (after.dev !== before.dev || after.ino !== before.ino) throw new Error(`Credential file changed while opening: ${filePath}`) + if (after.size > maxBytes) throw new Error(`Credential file exceeds ${maxBytes} bytes: ${filePath}`) + return await handle.readFile({ encoding: 'utf8' }) + } finally { + await handle.close() + } +} + +export async function atomicWriteSecureFile(filePath: string, contents: string): Promise { + await mkdir(path.dirname(filePath), { recursive: true, mode: 0o700 }) + const tempPath = path.join(path.dirname(filePath), `.${path.basename(filePath)}.${process.pid}.${Date.now()}.tmp`) + const handle = await open(tempPath, constants.O_WRONLY | constants.O_CREAT | constants.O_EXCL, 0o600) + try { + await handle.writeFile(contents, 'utf8') + await handle.sync() + } catch (error) { + await handle.close() + await unlink(tempPath).catch(() => undefined) + throw error + } + await handle.close() + await chmod(tempPath, 0o600) + try { + await rename(tempPath, filePath) + } catch (error) { + await unlink(tempPath).catch(() => undefined) + throw error + } +} + +export function quotaRequestSignal(parent?: AbortSignal): AbortSignal { + const timeout = AbortSignal.timeout(30_000) + return parent ? AbortSignal.any([parent, timeout]) : timeout +} + +export function fraction(value: unknown): number | null { + if (typeof value !== 'number' || !Number.isFinite(value)) return null + return Math.min(1, Math.max(0, value / 100)) +} diff --git a/app/electron/quota/types.ts b/app/electron/quota/types.ts new file mode 100644 index 00000000..4936f69b --- /dev/null +++ b/app/electron/quota/types.ts @@ -0,0 +1,17 @@ +export type QuotaWindow = { + label: string + percent: number + resetsAt: string | null +} + +export type QuotaProvider = { + provider: 'claude' | 'codex' + connection: 'connected' | 'disconnected' | 'loading' | 'stale' | 'transientFailure' | 'terminalFailure' + primary: QuotaWindow | null + details: QuotaWindow[] + planLabel: string | null + footerLines: string[] +} + +export type ProviderName = QuotaProvider['provider'] + diff --git a/app/renderer/lib/types.ts b/app/renderer/lib/types.ts index b61eed2c..7c64c1ee 100644 --- a/app/renderer/lib/types.ts +++ b/app/renderer/lib/types.ts @@ -19,6 +19,21 @@ export interface CliError { export type AliasRow = { from: string; to: string } export type ActionResult = { ok: boolean; stdout: string; stderr: string; code: number | null } +export type QuotaWindow = { + label: string + percent: number + resetsAt: string | null +} + +export type QuotaProvider = { + provider: 'claude' | 'codex' + connection: 'connected' | 'disconnected' | 'loading' | 'stale' | 'transientFailure' | 'terminalFailure' + primary: QuotaWindow | null + details: QuotaWindow[] + planLabel: string | null + footerLines: string[] +} + // ————— src/menubar-json.ts ————— export type DailyModelBreakdown = { @@ -468,6 +483,7 @@ export type CompareJsonReport = { // ————— IPC surface (preload contextBridge → window.codeburn) ————— export interface CodeburnBridge { + getQuota(): Promise getOverview(period: Period, provider: string, range?: DateRange): Promise getPlans(period: Period): Promise getActReport(): Promise