From 287791e2a98ceee771983dd6687549dab4686d5d Mon Sep 17 00:00:00 2001 From: iamtoruk Date: Sun, 23 Aug 2026 03:13:47 -0700 Subject: [PATCH] feat: live quota for Gemini, Copilot, and Antigravity in the desktop app Extends the existing borrowed-credential quota module (claude/codex) with three providers, per-provider settings toggles, and a provider display map: - gemini: Code Assist loadCodeAssist/retrieveUserQuota via the Gemini CLI's on-disk OAuth creds; optional env-gated token refresh; retired-tier responses degrade to terminalFailure with migration guidance - copilot: copilot_internal/user with editor-plugin headers; token read-only from hosts.json/apps.json with one re-read on 401; marked as an internal API that may drift, failures degrade honestly - antigravity: loopback-only Connect-RPC probe of the local language server (ps/lsof discovery via execFile, csrf token for the app variant, agy CLI accepted tokenless); TLS relaxation scoped to 127.0.0.1; no OAuth fallback - settings: per-provider quota switches (default on) persisted renderer-side, honored in the main-process service so disabled fetchers never run; IPC disabled-list filtered against the provider allowlist - sanitizeError extended for Google (ya29.) and GitHub (gh*_) token shapes Kimi audited, not duplicated: the macOS menubar already fetches Kimi quota live (KimiSubscriptionService.swift); the electron app not surfacing it is a pre-existing gap, unchanged here. --- app/electron/main.ts | 4 +- app/electron/preload.ts | 2 +- app/electron/quota/antigravity.test.ts | 154 ++++++++++++ app/electron/quota/antigravity.ts | 236 ++++++++++++++++++ app/electron/quota/copilot.test.ts | 145 +++++++++++ app/electron/quota/copilot.ts | 143 +++++++++++ app/electron/quota/gemini.test.ts | 149 +++++++++++ app/electron/quota/gemini.ts | 206 +++++++++++++++ app/electron/quota/index.test.ts | 86 +++++-- app/electron/quota/index.ts | 49 ++-- app/electron/quota/security.test.ts | 9 +- app/electron/quota/security.ts | 6 +- app/electron/quota/types.ts | 2 +- app/renderer/components/ConnectAffordance.tsx | 20 +- app/renderer/lib/providers.ts | 39 +++ app/renderer/lib/types.ts | 6 +- app/renderer/sections/Plans.test.tsx | 6 +- app/renderer/sections/Plans.tsx | 8 +- app/renderer/sections/Settings.test.tsx | 4 +- app/renderer/sections/Settings.tsx | 41 ++- 20 files changed, 1245 insertions(+), 70 deletions(-) create mode 100644 app/electron/quota/antigravity.test.ts create mode 100644 app/electron/quota/antigravity.ts create mode 100644 app/electron/quota/copilot.test.ts create mode 100644 app/electron/quota/copilot.ts create mode 100644 app/electron/quota/gemini.test.ts create mode 100644 app/electron/quota/gemini.ts create mode 100644 app/renderer/lib/providers.ts diff --git a/app/electron/main.ts b/app/electron/main.ts index 0147179f..19008845 100644 --- a/app/electron/main.ts +++ b/app/electron/main.ts @@ -344,8 +344,8 @@ export function createBridgeHandlers(deps: Deps = { spawnCli, spawnCliAction, re } return { - 'codeburn:getQuota': async (force?: boolean) => { - try { return { ok: true, value: await deps.getQuota({ force: Boolean(force) }) } } + 'codeburn:getQuota': async (force?: boolean, disabled?: string[]) => { + try { return { ok: true, value: await deps.getQuota({ force: Boolean(force), disabled }) } } catch (error) { return { ok: false, error: { kind: 'nonzero', message: sanitizeError(error) } } } }, 'codeburn:getOverview': getOverview, diff --git a/app/electron/preload.ts b/app/electron/preload.ts index f681c8a7..04501045 100644 --- a/app/electron/preload.ts +++ b/app/electron/preload.ts @@ -19,7 +19,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: (force?: boolean) => invoke('codeburn:getQuota', force), + getQuota: (force?: boolean, disabled?: string[]) => invoke('codeburn:getQuota', force, disabled), getOverview: (period: string, provider: string, range?: DateRange, configSource?: string | null, background?: boolean, scope?: string) => invoke('codeburn:getOverview', period, provider, range, configSource, background, scope), getTimeline: (period: string, provider: string, range?: DateRange) => invoke('codeburn:getTimeline', period, provider, range), getPlans: (period: string) => invoke('codeburn:getPlans', period), diff --git a/app/electron/quota/antigravity.test.ts b/app/electron/quota/antigravity.test.ts new file mode 100644 index 00000000..5768ba26 --- /dev/null +++ b/app/electron/quota/antigravity.test.ts @@ -0,0 +1,154 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { classifyProcessLine, decodeAntigravityStatus, decodeAntigravitySummary, fetchAntigravityQuota } from './antigravity' + +afterEach(() => vi.restoreAllMocks()) + +const appLine = '1234 /Applications/Antigravity.app/Contents/Resources/app/extensions/antigravity/bin/language_server_macos_arm --app_data_dir antigravity --csrf_token tok-123 --extension_server_port 54321' +const ideLine = '1235 /Applications/Antigravity IDE.app/.../extensions/antigravity/bin/language_server --app_data_dir antigravity-ide --csrf_token ide-tok' +const tokenlessAppLine = '1236 /usr/local/lib/language_server_macos --app_data_dir antigravity' +const cliLine = '1237 /opt/homebrew/bin/agy serve' + +describe('process classification', () => { + it('classifies the app language server with its CSRF token and fallback port', () => { + expect(classifyProcessLine(appLine)).toEqual({ pid: '1234', cli: false, csrf: 'tok-123', extPort: 54321 }) + }) + + it('skips tokenless app servers and IDE servers, accepts the agy CLI without a token', () => { + expect(classifyProcessLine(tokenlessAppLine)).toBeNull() + expect(classifyProcessLine(ideLine)).toBeNull() + const candidate = classifyProcessLine(cliLine) + expect(candidate).toMatchObject({ pid: '1237', cli: true }) + expect(candidate?.csrf).toBeUndefined() + }) + + it('ignores unrelated language servers and non-matching lines', () => { + expect(classifyProcessLine('999 /usr/bin/codeium_language_server --csrf_token x')).toBeNull() + expect(classifyProcessLine('998 vim /tmp/agy-notes.md')).toBeNull() + expect(classifyProcessLine('garbage line')).toBeNull() + }) +}) + +describe('payload decoding', () => { + it('decodes summary groups into joined windows', () => { + const windows = decodeAntigravitySummary({ + groups: [ + { + displayName: 'Gemini Models', + buckets: [ + { displayName: 'Weekly limit', remaining: { remainingFraction: 0.8 } }, + { displayName: 'Five hour limit', remaining: { remainingFraction: 0.25 } }, + { displayName: 'No fraction row' }, + ], + }, + { displayName: 'Claude and GPT models', buckets: [{ bucketId: 'claude_weekly', remaining: { remainingFraction: 1 } }] }, + ], + }) + expect(windows.map(row => row.label)).toEqual(['Gemini Models · Weekly limit', 'Gemini Models · Five hour limit', 'Claude and GPT models · claude_weekly']) + expect(windows[1]!.percent).toBeCloseTo(0.75) + expect(windows[2]!.percent).toBe(0) + }) + + it('decodes the legacy GetUserStatus quota rows with reset times', () => { + const windows = decodeAntigravityStatus({ + userStatus: { + cascadeModelConfigData: { + clientModelConfigs: [ + { modelName: 'gemini-2.5-pro', quotaInfo: { remainingFraction: 0.5, resetTime: 1_800_000_000 } }, + { modelName: 'claude-sonnet-4', quotaInfo: { remainingFraction: 0.9, resetTime: '2026-07-12T00:00:00Z' } }, + { modelName: 'no-quota-model' }, + ], + }, + }, + }) + expect(windows.map(row => row.label)).toEqual(['gemini-2.5-pro', 'claude-sonnet-4']) + expect(windows[0]!.resetsAt).toBe(new Date(1_800_000_000 * 1000).toISOString()) + expect(windows[1]!.resetsAt).toBe('2026-07-12T00:00:00.000Z') + }) + + it('returns no windows for garbage payloads', () => { + expect(decodeAntigravitySummary(null)).toEqual([]) + expect(decodeAntigravityStatus({ userStatus: {} })).toEqual([]) + }) +}) + +describe('Antigravity local probe', () => { + const psOutput = [appLine, cliLine].join('\n') + + it('probes the app server over loopback TLS then HTTP and renders the most constrained window first', async () => { + const calls: Array<[number, boolean, string]> = [] + const request = vi.fn(async (port: number, tls: boolean, pathName: string) => { + calls.push([port, tls, pathName]) + if (!tls) return null + if (pathName.includes('RetrieveUserQuotaSummary')) return { status: 200, text: JSON.stringify({ groups: [{ displayName: 'Gemini Models', buckets: [{ displayName: 'Weekly limit', remaining: { remainingFraction: 0.7 } }] }] }) } + return { status: 200, text: '{}' } + }) + const execFile = vi.fn(async (_file: string, args: string[]) => args[0] === '-ax' + ? { stdout: psOutput } + : { stdout: `python 1234 user 12u IPv4 0x1 0t0 TCP 127.0.0.1:60123 (LISTEN)\n` }) + const quota = await fetchAntigravityQuota({ execFile, request }) + expect(quota.connection).toBe('connected') + expect(quota.primary?.percent).toBeCloseTo(0.3) + expect(quota.primary?.label).toBe('Gemini Models · Weekly limit') + expect(calls.every(([port]) => port === 60123)).toBe(true) + expect(calls[0]).toEqual([60123, true, '/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary']) + }) + + it('sends the CSRF header for app servers and none for the agy CLI', async () => { + const headers: Array = [] + let call = 0 + const request = vi.fn(async (_port: number, _tls: boolean, _pathName: string, _body: string, csrf?: string) => { + headers.push(csrf) + call += 1 + return call <= 2 + ? null + : { status: 200, text: JSON.stringify({ groups: [{ displayName: 'Claude + GPT', buckets: [{ bucketId: 'weekly', remaining: { remainingFraction: 0.1 } }] }] }) } + }) + const execFile = vi.fn(async (_file: string, args: string[]) => args[0] === '-ax' + ? { stdout: cliLine } + : { stdout: `agy 1237 user 5u IPv4 0x1 0t0 TCP *:60555 (LISTEN)\n` }) + const quota = await fetchAntigravityQuota({ execFile, request }) + expect(quota.connection).toBe('connected') + expect(quota.planLabel).toBeNull() + expect(headers[0]).toBeUndefined() + }) + + it('falls back to GetUserStatus and lifts planName when the summary has no windows', async () => { + const request = vi.fn(async (_port: number, tls: boolean, pathName: string) => tls && pathName.includes('GetUserStatus') + ? { status: 200, text: JSON.stringify({ userStatus: { planName: 'AI Pro', cascadeModelConfigData: { clientModelConfigs: [{ modelName: 'gemini-2.5-pro', quotaInfo: { remainingFraction: 0.4 } }] } } }) } + : null) + const execFile = vi.fn(async (_file: string, args: string[]) => args[0] === '-ax' + ? { stdout: cliLine } + : { stdout: `agy 1237 user 5u IPv4 0x1 0t0 TCP *:60555 (LISTEN)\n` }) + const quota = await fetchAntigravityQuota({ execFile, request }) + expect(quota.planLabel).toBe('AI Pro') + expect(quota.details.map(row => row.label)).toEqual(['gemini-2.5-pro']) + }) + + it('reports disconnected when nothing local is listening', async () => { + const request = vi.fn() + const execFile = vi.fn(async () => ({ stdout: '' })) + const quota = await fetchAntigravityQuota({ execFile, request }) + expect(quota.connection).toBe('disconnected') + expect(request).not.toHaveBeenCalled() + }) + + it('reports disconnected when every port probe fails', async () => { + const request = vi.fn(async () => null) + const execFile = vi.fn(async (_file: string, args: string[]) => args[0] === '-ax' + ? { stdout: cliLine } + : { stdout: `agy 1237 user 5u IPv4 0x1 0t0 TCP *:60555 (LISTEN)\n` }) + const quota = await fetchAntigravityQuota({ execFile, request }) + expect(quota.connection).toBe('disconnected') + }) + + it('degrades an unexpected ps failure to transientFailure with sanitized diagnostics', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const execFile = vi.fn(async () => { throw new Error('ps exploded Bearer sk-secret eyJabc.def\0tail') }) + const quota = await fetchAntigravityQuota({ execFile }) + expect(quota.connection).toBe('transientFailure') + const logged = warn.mock.calls.flat().join(' ') + expect(logged).not.toMatch(/sk-secret|eyJabc|\0/) + expect(logged).toContain('[REDACTED]') + }) +}) diff --git a/app/electron/quota/antigravity.ts b/app/electron/quota/antigravity.ts new file mode 100644 index 00000000..d407d05b --- /dev/null +++ b/app/electron/quota/antigravity.ts @@ -0,0 +1,236 @@ +// Live Antigravity quota from LOCAL surfaces only — the Antigravity app's own +// language server or a signed-in `agy` CLI's embedded server (prior art: +// steipete/CodexBar docs/antigravity.md, which marks this protocol internal and +// experimental). No Google OAuth fallback in v1: when no local server answers, +// the provider reports disconnected and the UI shows its Connect affordance. +// +// Endpoints (localhost only, Connect-RPC JSON): +// - POST https://127.0.0.1:/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary +// (preferred; falls back to) +// - POST https://127.0.0.1:/exa.language_server_pb.LanguageServerService/GetUserStatus +// +// Discovery mirrors CodexBar: `ps` finds candidate processes (app language +// servers need their `--csrf_token`; the `agy` CLI needs none), then `lsof` +// lists each pid's listening TCP ports. Local HTTPS uses a self-signed cert, +// so TLS verification is relaxed ONLY for the 127.0.0.1 loopback probes. +import { execFile } from 'node:child_process' +import http from 'node:http' +import https from 'node:https' +import { promisify } from 'node:util' + +import { fraction, sanitizeError } from './security' +import type { QuotaProvider, QuotaWindow } from './types' + +const SUMMARY_PATH = '/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary' +const STATUS_PATH = '/exa.language_server_pb.LanguageServerService/GetUserStatus' +const LOCAL_TIMEOUT_MS = 3_000 + +export type ExecFileFn = (file: string, args: string[]) => Promise<{ stdout: string }> +export type LocalRequestFn = ( + port: number, + tls: boolean, + pathName: string, + body: string, + csrf?: string, +) => Promise<{ status: number; text: string } | null> + +export type AntigravityDeps = { + execFile: ExecFileFn + request: LocalRequestFn +} + +const execFileAsync: ExecFileFn = promisify(execFile) + +function postLocal(port: number, tls: boolean, pathName: string, body: string, csrf?: string): Promise<{ status: number; text: string } | null> { + return new Promise(resolve => { + const request = (tls ? https : http).request({ + host: '127.0.0.1', + port, + method: 'POST', + path: pathName, + ...(tls ? { rejectUnauthorized: false } : {}), + headers: { + 'Content-Type': 'application/json', + 'Connect-Protocol-Version': '1', + ...(csrf ? { 'X-Codeium-Csrf-Token': csrf } : {}), + }, + timeout: LOCAL_TIMEOUT_MS, + }, response => { + let text = '' + response.setEncoding('utf8') + response.on('data', chunk => { text += chunk }) + response.on('end', () => resolve({ status: response.statusCode ?? 0, text })) + }) + request.on('timeout', () => { request.destroy(); resolve(null) }) + request.on('error', () => resolve(null)) + request.end(body) + }) +} + +const defaults: AntigravityDeps = { + execFile: execFileAsync, + request: postLocal, +} + +function empty(connection: QuotaProvider['connection']): QuotaProvider { + return { provider: 'antigravity', connection, primary: null, details: [], planLabel: null, footerLines: [] } +} + +// Process kinds, mirroring CodexBar's classification. The app language server +// carries richer quota data than the IDE variant, so IDE matches are skipped; +// a CLI (`agy`) match is accepted because its tokenless server exposes the +// same summary payload. +const LANGUAGE_SERVER = /language[_-]server(_macos)?(_arm)?/ +const IDE_MARKER = /antigravity[-_]ide/i +const CLI_MARKER = /antigravity[-_]cli/i +const AGY_BINARY = /(^|\/|\s)agy(\s|$)/ +const APP_MARKER = /--app_data_dir[= ]"?antigravity(?![\w-])|\/antigravity\//i + +type Candidate = { pid: string; cli: boolean; csrf?: string; extPort?: number } + +function flagValue(line: string, flag: string): string | undefined { + const match = line.match(new RegExp(`${flag}[= ]([^\\s]+)`)) + return match?.[1] +} + +export function classifyProcessLine(line: string): Candidate | null { + const isServer = LANGUAGE_SERVER.test(line) + const isCli = !isServer && (CLI_MARKER.test(line) || AGY_BINARY.test(line)) + if (!isServer && !isCli) return null + const pid = line.trim().match(/^(\d+)/)?.[1] + if (!pid) return null + // A tokenless desktop language-server match is skipped so a later, valid + // server can be found; the CLI exposes no CSRF flag and needs none. IDE + // language servers are excluded too — their payloads lack the weekly groups. + if (!isCli && (!APP_MARKER.test(line) || IDE_MARKER.test(line))) return null + const csrf = flagValue(line, '--csrf_token') + if (!isCli && !csrf) return null + const extPort = Number(flagValue(line, '--extension_server_port')) + return { pid, cli: isCli, csrf, extPort: Number.isFinite(extPort) && extPort > 0 ? extPort : undefined } +} + +async function discoverCandidates(deps: AntigravityDeps): Promise { + const { stdout } = await deps.execFile('ps', ['-ax', '-o', 'pid=,command=']) + const seen = new Set() + const candidates: Candidate[] = [] + for (const line of stdout.split('\n')) { + const candidate = classifyProcessLine(line) + if (!candidate || seen.has(candidate.pid)) continue + seen.add(candidate.pid) + candidates.push(candidate) + } + // The app source ranks above the CLI: richer quota data beats availability. + return candidates.sort((a, b) => Number(a.cli) - Number(b.cli)) +} + +async function listeningPorts(deps: AntigravityDeps, pid: string): Promise { + try { + const { stdout } = await deps.execFile('lsof', ['-nP', '-iTCP', '-sTCP:LISTEN', '-a', '-p', pid]) + const ports = [...stdout.matchAll(/:(\d+)\s/g)].map(match => Number(match[1])) + return [...new Set(ports)].filter(port => Number.isFinite(port) && port > 0) + } catch { + return [] + } +} + +function resetTimeOf(value: unknown): string | null { + if (typeof value === 'number' && Number.isFinite(value)) { + return new Date(value > 1e12 ? value : value * 1000).toISOString() + } + if (typeof value === 'string' && !Number.isNaN(Date.parse(value))) return new Date(value).toISOString() + return null +} + +function windowOf(label: string, remainingFraction: unknown, resetTime?: unknown): QuotaWindow | null { + const remaining = fraction(typeof remainingFraction === 'number' ? (1 - remainingFraction) * 100 : NaN) + if (remaining === null) return null + return { label, percent: remaining, resetsAt: resetTimeOf(resetTime ?? null) } +} + +/** Preferred payload: two named quota groups of model buckets. */ +export function decodeAntigravitySummary(body: unknown): QuotaWindow[] { + const data = body && typeof body === 'object' ? body as Record : {} + const windows: QuotaWindow[] = [] + for (const group of Array.isArray(data.groups) ? data.groups : []) { + const groupName = typeof group?.displayName === 'string' ? group.displayName : '' + for (const bucket of Array.isArray(group?.buckets) ? group.buckets : []) { + const name = [groupName, typeof bucket?.displayName === 'string' ? bucket.displayName : bucket?.bucketId] + .filter(Boolean).join(' · ') + if (!name) continue + const window = windowOf(name, bucket?.remaining?.remainingFraction) + if (window) windows.push(window) + } + } + return windows +} + +/** Legacy payload: flat per-model quota rows under GetUserStatus. */ +export function decodeAntigravityStatus(body: unknown): QuotaWindow[] { + const data = body && typeof body === 'object' ? body as Record : {} + const configs = data.userStatus?.cascadeModelConfigData?.clientModelConfigs + const windows: QuotaWindow[] = [] + for (const config of Array.isArray(configs) ? configs : []) { + const name = typeof config?.modelName === 'string' ? config.modelName : '' + if (!name) continue + const window = windowOf(name, config?.quotaInfo?.remainingFraction, config?.quotaInfo?.resetTime) + if (window) windows.push(window) + } + return windows +} + +function planFromStatus(body: unknown): string | null { + const data = body && typeof body === 'object' ? body as Record : {} + const plan = data.planName ?? data.userStatus?.planName ?? data.account_plan + return typeof plan === 'string' && plan.trim() ? plan.trim() : null +} + +async function probePort(deps: AntigravityDeps, port: number, csrf?: string): Promise<{ windows: QuotaWindow[]; planLabel: string | null } | null> { + const body = JSON.stringify({}) + for (const tls of [true, false]) { + const summary = await deps.request(port, tls, SUMMARY_PATH, body, csrf) + if (summary?.status === 200) { + const windows = decodeAntigravitySummary(parseJson(summary.text)) + if (windows.length > 0) return { windows, planLabel: null } + } + const status = await deps.request(port, tls, STATUS_PATH, body, csrf) + if (status?.status === 200) { + const parsed = parseJson(status.text) + const windows = decodeAntigravityStatus(parsed) + if (windows.length > 0) return { windows, planLabel: planFromStatus(parsed) } + } + } + return null +} + +function parseJson(text: string): unknown { + try { return JSON.parse(text) } catch { return null } +} + +export async function fetchAntigravityQuota(deps: Partial = {}): Promise { + const resolved = { ...defaults, ...deps } + try { + const candidates = await discoverCandidates(resolved) + if (candidates.length === 0) return empty('disconnected') + for (const candidate of candidates) { + const ports = await listeningPorts(resolved, candidate.pid) + if (candidate.extPort !== undefined && !ports.includes(candidate.extPort)) ports.push(candidate.extPort) + for (const port of ports) { + const found = await probePort(resolved, port, candidate.csrf) + if (found) { + const windows = [...found.windows].sort((a, b) => b.percent - a.percent) + return { + provider: 'antigravity', connection: 'connected', + primary: windows[0] ?? null, + details: windows, + planLabel: found.planLabel, + footerLines: [], + } + } + } + } + return empty('disconnected') + } catch (error) { + console.warn(`Antigravity quota unavailable: ${sanitizeError(error)}`) + return empty('transientFailure') + } +} diff --git a/app/electron/quota/copilot.test.ts b/app/electron/quota/copilot.test.ts new file mode 100644 index 00000000..353297f2 --- /dev/null +++ b/app/electron/quota/copilot.test.ts @@ -0,0 +1,145 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { decodeCopilotUsage, fetchCopilotQuota } from './copilot' + +const hosts = JSON.stringify({ 'github.com': { user: 'octocat', oauth_token: 'gho_test-secret' } }) +const usageBody = { + copilot_plan: 'individual', + quota_snapshots: { + premium_interactions: { percent_remaining: 70 }, + chat: { percent_remaining: 100 }, + }, +} + +const okJson = (value: unknown) => new Response(JSON.stringify(value), { status: 200 }) + +afterEach(() => vi.restoreAllMocks()) + +describe('Copilot usage decode', () => { + it('decodes remaining-percent snapshots into used windows with the plan label', () => { + const quota = decodeCopilotUsage(usageBody) + expect(quota.connection).toBe('connected') + expect(quota.planLabel).toBe('Individual') + expect(quota.primary).toEqual({ label: 'Premium requests', percent: 0.3, resetsAt: null }) + expect(quota.details.map(row => row.label)).toEqual(['Premium requests', 'Chat']) + expect(quota.details.map(row => row.percent)).toEqual([0.3, 0]) + }) + + it('accepts camelCase spellings and promotes chat when premium is absent', () => { + const quota = decodeCopilotUsage({ + copilotPlan: 'business', + quotaSnapshots: { chat: { percentRemaining: 55 } }, + }) + expect(quota.planLabel).toBe('Business') + expect(quota.primary?.label).toBe('Chat') + expect(quota.primary?.percent).toBeCloseTo(0.45) + }) + + it('survives a malformed payload without usable snapshots', () => { + const quota = decodeCopilotUsage({ quota_snapshots: { chat: 'garbage' }, extra: true }) + expect(quota.connection).toBe('connected') + expect(quota.primary).toBeNull() + expect(quota.details).toEqual([]) + }) + + it('title-cases unknown plan tiers', () => { + expect(decodeCopilotUsage({ copilot_plan: 'for_educators' }).planLabel).toBe('Educators') + expect(decodeCopilotUsage({ copilot_plan: 'some_future_tier' }).planLabel).toBe('Some Future Tier') + expect(decodeCopilotUsage({}).planLabel).toBeNull() + }) +}) + +describe('Copilot quota fetch', () => { + it('returns disconnected without credentials and never fetches', async () => { + const fetchMock = vi.fn() + const result = await fetchCopilotQuota({ fetch: fetchMock, readFile: vi.fn(async () => null) }) + expect(result.quota.connection).toBe('disconnected') + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('reads hosts.json first and uses exact plugin headers', async () => { + const fetchMock = vi.fn(async () => okJson(usageBody)) + const result = await fetchCopilotQuota({ + fetch: fetchMock, + readFile: vi.fn(async (path: string) => path.endsWith('hosts.json') ? hosts : JSON.stringify({ 'Some App': { oauth_token: 'gho_wrong' } })), + }) + expect(result.quota.connection).toBe('connected') + const [url, init] = fetchMock.mock.calls[0]! as unknown as [string, RequestInit] + expect(url).toBe('https://api.github.com/copilot_internal/user') + expect(init.method).toBe('GET') + expect(init.headers).toEqual({ + Authorization: 'token gho_test-secret', + Accept: 'application/json', + 'Editor-Version': 'vscode/1.96.2', + 'Editor-Plugin-Version': 'copilot-chat/0.26.7', + 'User-Agent': 'GitHubCopilotChat/0.26.7', + 'X-Github-Api-Version': '2025-04-01', + }) + }) + + it('falls back to apps.json when hosts.json has no token', async () => { + const fetchMock = vi.fn(async () => okJson(usageBody)) + const result = await fetchCopilotQuota({ + fetch: fetchMock, + readFile: vi.fn(async (path: string) => path.endsWith('apps.json') ? JSON.stringify({ 'Visual Studio Code': { oauth_token: 'ghu_apps-token' } }) : '{}'), + }) + expect(result.quota.connection).toBe('connected') + const init = (fetchMock.mock.calls[0]! as unknown as [string, RequestInit])[1] + expect(init.headers).toMatchObject({ Authorization: 'token ghu_apps-token' }) + }) + + it('re-reads once on a 401 and adopts a rotated token', async () => { + let reads = 0 + const readFile = vi.fn(async () => { + reads += 1 + return JSON.stringify({ 'github.com': { oauth_token: reads === 1 ? 'gho_stale' : 'gho_rotated' } }) + }) + const fetchMock = vi.fn(async (_url: string | URL | RequestInfo, init?: RequestInit) => (init?.headers as Record).Authorization === 'token gho_rotated' + ? okJson(usageBody) + : new Response('', { status: 401 })) + const result = await fetchCopilotQuota({ fetch: fetchMock, readFile }) + expect(result.quota.connection).toBe('connected') + expect(fetchMock).toHaveBeenCalledTimes(2) + }) + + it('stays transientFailure when a 401 leaves the stored token unchanged', async () => { + const fetchMock = vi.fn(async () => new Response('', { status: 401 })) + const result = await fetchCopilotQuota({ fetch: fetchMock, readFile: vi.fn(async () => hosts) }) + expect(result.quota.connection).toBe('transientFailure') + // One probe only: re-reading found the same token, so a retry is pointless. + expect(fetchMock).toHaveBeenCalledTimes(1) + }) + + it('uses the Retry-After header for 429 backoff', async () => { + const fetchMock = vi.fn(async () => new Response('', { status: 429, headers: { 'Retry-After': '75' } })) + const result = await fetchCopilotQuota({ fetch: fetchMock, readFile: vi.fn(async () => hosts) }) + expect(result.retryAfterSeconds).toBe(75) + expect(result.quota.connection).toBe('transientFailure') + }) + + it('maps 5xx to transientFailure and other 4xx to terminalFailure', async () => { + const serverError = vi.fn(async () => new Response('', { status: 503 })) + const bad = await fetchCopilotQuota({ fetch: serverError, readFile: vi.fn(async () => hosts) }) + expect(bad.quota.connection).toBe('transientFailure') + const clientError = vi.fn(async () => new Response('', { status: 404 })) + const worse = await fetchCopilotQuota({ fetch: clientError, readFile: vi.fn(async () => hosts) }) + expect(worse.quota.connection).toBe('terminalFailure') + }) + + it('degrades a malformed success body instead of crashing the panel', async () => { + const fetchMock = vi.fn(async () => new Response('not json {', { status: 200 })) + const result = await fetchCopilotQuota({ fetch: fetchMock, readFile: vi.fn(async () => hosts) }) + expect(result.quota.connection).toBe('transientFailure') + expect(result.quota.primary).toBeNull() + }) + + 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('token gho_leak-secret eyJabc.def.ghi\0tail') }) + const result = await fetchCopilotQuota({ fetch: fetchMock, readFile: vi.fn(async () => hosts) }) + const logged = warn.mock.calls.flat().join(' ') + expect(result.quota).not.toHaveProperty('error') + expect(logged).not.toMatch(/gho_leak|eyJabc|\0/) + expect(logged).toContain('[REDACTED]') + }) +}) diff --git a/app/electron/quota/copilot.ts b/app/electron/quota/copilot.ts new file mode 100644 index 00000000..785998f5 --- /dev/null +++ b/app/electron/quota/copilot.ts @@ -0,0 +1,143 @@ +// Live GitHub Copilot quota via the editor plugins' internal usage endpoint. +// +// - GET https://api.github.com/copilot_internal/user +// Headers mirror the VS Code Copilot Chat plugin (Editor-Version / +// Editor-Plugin-Version / X-Github-Api-Version, prior art: +// steipete/CodexBar docs/copilot.md). This is an INTERNAL, UNDOCUMENTED +// API that may drift without notice; every failure must degrade to the +// normal connection states and never crash the panel. +// +// Credential: the GitHub OAuth token already on disk from a signed-in Copilot +// plugin — ~/.config/github-copilot/hosts.json (keyed by host) falling back to +// apps.json (keyed by app name). Read-only; no new storage. +import os from 'node:os' +import path from 'node:path' + +import { fraction, quotaRequestSignal, readSecureFile, sanitizeError } from './security' +import type { QuotaProvider, QuotaWindow } from './types' + +const USAGE_ENDPOINT = 'https://api.github.com/copilot_internal/user' +const HEADERS = { + Accept: 'application/json', + 'Editor-Version': 'vscode/1.96.2', + 'Editor-Plugin-Version': 'copilot-chat/0.26.7', + 'User-Agent': 'GitHubCopilotChat/0.26.7', + 'X-Github-Api-Version': '2025-04-01', +} as const + +type HostRecord = Record & { oauth_token?: unknown } + +export type CopilotDeps = { + fetch: typeof fetch + hostsPath: string + appsPath: string + readFile: typeof readSecureFile +} + +const defaults: CopilotDeps = { + fetch: globalThis.fetch, + hostsPath: path.join(os.homedir(), '.config', 'github-copilot', 'hosts.json'), + appsPath: path.join(os.homedir(), '.config', 'github-copilot', 'apps.json'), + readFile: readSecureFile, +} + +function empty(connection: QuotaProvider['connection']): QuotaProvider { + return { provider: 'copilot', connection, primary: null, details: [], planLabel: null, footerLines: [] } +} + +function tokenFromMap(raw: string): string | null { + const map = JSON.parse(raw) as Record + // hosts.json keys by host — prefer github.com; apps.json has no canonical + // key, so its first entry wins. Both store the token as `oauth_token`. + const preferred = map['github.com'] ?? Object.values(map)[0] + const token = preferred?.oauth_token + return typeof token === 'string' && token ? token : null +} + +async function credentialFromFiles(deps: CopilotDeps): Promise { + for (const filePath of [deps.hostsPath, deps.appsPath]) { + try { + const raw = await deps.readFile(filePath, 64 * 1024) + if (!raw) continue + const token = tokenFromMap(raw) + if (token) return token + } catch { + // A malformed or unreadable file falls through to the next candidate. + } + } + return null +} + +function windowOf(label: string, snapshot: unknown): QuotaWindow | null { + if (!snapshot || typeof snapshot !== 'object') return null + const row = snapshot as Record + // The API reports percent REMAINING (0..100); windows render percent USED. + const rawRemaining = row.percent_remaining ?? row.percentRemaining + const remaining = fraction(typeof rawRemaining === 'number' ? rawRemaining : NaN) + if (remaining === null) return null + // Round away float dust from the 1-remaining subtraction (1-0.7 !== 0.3). + return { label, percent: Number((1 - remaining).toFixed(6)), resetsAt: null } +} + +function planLabel(value: unknown): string | null { + if (typeof value !== 'string' || !value.trim()) return null + const lower = value.trim().toLowerCase() + const known: Record = { + free: 'Free', individual: 'Individual', pro: 'Pro', business: 'Business', + enterprise: 'Enterprise', for_educators: 'Educators', 'for-educators': 'Educators', + } + return known[lower] ?? lower.replace(/(^|[_-])\w/g, match => match.replace(/[_-]/, ' ').toUpperCase()) +} + +export function decodeCopilotUsage(body: unknown): QuotaProvider { + const data = body && typeof body === 'object' ? body as Record : {} + // Field names have shipped both camelCase and snake_case; read each alias + // rather than trusting one spelling. + const snapshots = data.quota_snapshots ?? data.quotaSnapshots + const premium = windowOf('Premium requests', snapshots?.premium_interactions ?? snapshots?.premiumInteractions) + const chat = windowOf('Chat', snapshots?.chat) + const details = [premium, chat].filter((row): row is QuotaWindow => row !== null) + return { + provider: 'copilot', connection: 'connected', primary: premium ?? chat, + details, + planLabel: planLabel(data.copilot_plan ?? data.copilotPlan), + footerLines: [], + } +} + +async function request(token: string, deps: CopilotDeps, parent?: AbortSignal): Promise { + return deps.fetch(USAGE_ENDPOINT, { + method: 'GET', signal: quotaRequestSignal(parent), + headers: { ...HEADERS, Authorization: `token ${token}` }, + }) +} + +export type CopilotResult = { quota: QuotaProvider; retryAfterSeconds?: number } + +export async function fetchCopilotQuota(options: Partial & { signal?: AbortSignal } = {}): Promise { + const deps = { ...defaults, ...options } + try { + let token = await credentialFromFiles(deps) + if (!token) return { quota: empty('disconnected') } + + let response = await request(token, deps, options.signal) + if (response.status === 401) { + // An active editor session rotates this token; re-read once before + // giving up so we don't report a failure the disk already fixed. + const reread = await credentialFromFiles(deps) + if (!reread || reread === token) return { quota: empty('transientFailure') } + token = reread + response = await request(token, deps, options.signal) + } + if (response.status === 429) { + const raw = response.headers.get('Retry-After') + const seconds = raw === null ? NaN : Number(raw) + 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: decodeCopilotUsage(await response.json()) } + } catch (error) { + console.warn(`Copilot quota unavailable: ${sanitizeError(error)}`) + return { quota: empty('transientFailure') } + } +} diff --git a/app/electron/quota/gemini.test.ts b/app/electron/quota/gemini.test.ts new file mode 100644 index 00000000..b10aff79 --- /dev/null +++ b/app/electron/quota/gemini.test.ts @@ -0,0 +1,149 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { decodeGeminiUsage, fetchGeminiQuota } from './gemini' + +const credential = JSON.stringify({ + access_token: 'ya29.test-secret', + refresh_token: '1//refresh-secret', + expiry_date: Date.now() + 3_600_000, +}) + +// A Google Workspace id_token carries the hosted-domain (`hd`) JWT claim. +const workspaceCredential = JSON.stringify({ + access_token: 'ya29.workspace-secret', + id_token: `eyJhbGciOiJub25lIn0.${Buffer.from(JSON.stringify({ hd: 'example.com' })).toString('base64url')}.sig`, + expiry_date: Date.now() + 3_600_000, +}) + +const quotaBody = { + buckets: [ + { modelId: 'gemini-2.5-flash', remainingFraction: 0.9, resetTime: '2026-07-13T00:00:00Z' }, + { modelId: 'gemini-2.5-pro', remainingFraction: 0.25, resetTime: '2026-07-12T18:00:00Z' }, + { modelId: 'gemini-2.5-lite', remainingFraction: 'garbage' }, + ], +} + +const okJson = (value: unknown) => new Response(JSON.stringify(value), { status: 200 }) + +afterEach(() => vi.unstubAllEnvs()) + +describe('Gemini usage decode', () => { + it('decodes buckets most-constrained first and derives used percent', () => { + const quota = decodeGeminiUsage(quotaBody) + expect(quota.connection).toBe('connected') + expect(quota.primary?.label).toBe('gemini-2.5-pro') + expect(quota.primary?.percent).toBeCloseTo(0.75) + expect(quota.details.map(row => row.label)).toEqual(['gemini-2.5-pro', 'gemini-2.5-flash']) + expect(quota.details[0]!.resetsAt).toBe('2026-07-12T18:00:00.000Z') + }) + + it('survives a malformed payload with no usable buckets', () => { + const quota = decodeGeminiUsage({ buckets: [null, 'x', {}], extra: true }) + expect(quota.connection).toBe('connected') + expect(quota.primary).toBeNull() + expect(quota.details).toEqual([]) + }) +}) + +describe('Gemini quota fetch', () => { + it('returns disconnected without credentials and never fetches', async () => { + const fetchMock = vi.fn() + const result = await fetchGeminiQuota({ fetch: fetchMock, readFile: vi.fn(async () => null) }) + expect(result.quota.connection).toBe('disconnected') + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('calls loadCodeAssist then retrieveUserQuota with the discovered project and exact headers', async () => { + const fetchMock = vi.fn(async (url: string | URL | RequestInfo) => String(url).includes('loadCodeAssist') + ? okJson({ currentTier: { name: 'free-tier' }, cloudaicompanionProject: 'gen-lang-client-1' }) + : okJson(quotaBody)) + const result = await fetchGeminiQuota({ fetch: fetchMock, readFile: vi.fn(async () => credential) }) + expect(result.quota.connection).toBe('connected') + expect(result.quota.planLabel).toBe('Free') + expect(fetchMock.mock.calls.every(call => String(call[0]).startsWith('https://cloudcode-pa.googleapis.com/v1internal:'))).toBe(true) + const [, quotaInit] = fetchMock.mock.calls[1]! as unknown as [string, RequestInit] + expect(JSON.parse(String(quotaInit.body))).toEqual({ project: 'gen-lang-client-1' }) + expect(quotaInit.headers).toMatchObject({ + Authorization: 'Bearer ya29.test-secret', + 'Content-Type': 'application/json', + 'User-Agent': 'CodeBurn', + }) + }) + + it('sends an empty project object when discovery yields none', async () => { + const fetchMock = vi.fn(async (url: string | URL | RequestInfo) => String(url).includes('loadCodeAssist') ? okJson({}) : okJson(quotaBody)) + await fetchGeminiQuota({ fetch: fetchMock, readFile: vi.fn(async () => credential) }) + const [, quotaInit] = fetchMock.mock.calls[1]! as unknown as [string, RequestInit] + expect(JSON.parse(String(quotaInit.body))).toEqual({}) + }) + + it('refreshes a stale token through Google OAuth using the documented env overrides', async () => { + vi.stubEnv('GEMINI_OAUTH_CLIENT_ID', 'client-id') + vi.stubEnv('GEMINI_OAUTH_CLIENT_SECRET', 'client-secret') + const stale = JSON.stringify({ ...JSON.parse(credential), expiry_date: Date.now() - 1000 }) + const fetchMock = vi.fn(async (url: string | URL | RequestInfo) => String(url).includes('oauth2.googleapis.com') + ? okJson({ access_token: 'ya29.refreshed' }) + : String(url).includes('loadCodeAssist') + ? okJson({ paidTier: { name: 'standard-tier' } }) + : okJson(quotaBody)) + const result = await fetchGeminiQuota({ fetch: fetchMock, readFile: vi.fn(async () => stale) }) + expect(result.quota.planLabel).toBe('Paid') + const [tokenUrl, tokenInit] = fetchMock.mock.calls[0]! as unknown as [string, RequestInit] + expect(tokenUrl).toBe('https://oauth2.googleapis.com/token') + expect(tokenInit.method).toBe('POST') + expect(String(tokenInit.body)).toContain('grant_type=refresh_token') + const init = (fetchMock.mock.calls.at(-1)! as unknown as [string, RequestInit])[1] + expect(init.headers).toMatchObject({ Authorization: 'Bearer ya29.refreshed' }) + }) + + it('uses the Retry-After header for 429 backoff', async () => { + const fetchMock = vi.fn(async () => new Response('', { status: 429, headers: { 'Retry-After': '90' } })) + const result = await fetchGeminiQuota({ fetch: fetchMock, readFile: vi.fn(async () => credential) }) + expect(result.retryAfterSeconds).toBe(90) + expect(result.quota.connection).toBe('transientFailure') + }) + + it('stays transientFailure when a 401 leaves the stored token unchanged', async () => { + const fetchMock = vi.fn(async () => new Response('', { status: 401 })) + const result = await fetchGeminiQuota({ fetch: fetchMock, readFile: vi.fn(async () => credential) }) + expect(result.quota.connection).toBe('transientFailure') + expect((fetchMock.mock.calls as unknown as Array<[string]>).every(call => call[0].includes('loadCodeAssist'))).toBe(true) + }) + + it('maps retired consumer tiers to terminalFailure with migration guidance', async () => { + const fetchMock = vi.fn(async () => new Response(JSON.stringify({ error: { code: 400, status: 'UNSUPPORTED_CLIENT', message: 'IneligibleTierError: use Antigravity' } }), { status: 400 })) + const result = await fetchGeminiQuota({ fetch: fetchMock, readFile: vi.fn(async () => credential) }) + expect(result.quota.connection).toBe('terminalFailure') + expect(result.quota.footerLines).toEqual(['Google retired Gemini CLI OAuth for this account tier — use Antigravity.']) + }) + + it('labels a Workspace account from the id_token hd claim, personal from the tier alone', async () => { + const assist = { currentTier: { name: 'free-tier' } } + const workspace = vi.fn(async (url: string | URL | RequestInfo) => String(url).includes('loadCodeAssist') ? okJson(assist) : okJson(quotaBody)) + const result = await fetchGeminiQuota({ fetch: workspace, readFile: vi.fn(async () => workspaceCredential) }) + expect(result.quota.planLabel).toBe('Workspace') + + const personal = vi.fn(async (url: string | URL | RequestInfo) => String(url).includes('loadCodeAssist') ? okJson(assist) : okJson(quotaBody)) + const free = await fetchGeminiQuota({ fetch: personal, readFile: vi.fn(async () => credential) }) + expect(free.quota.planLabel).toBe('Free') + }) + + it('maps 5xx to transientFailure and other 4xx to terminalFailure', async () => { + const serverError = vi.fn(async () => new Response('', { status: 503 })) + const bad = await fetchGeminiQuota({ fetch: serverError, readFile: vi.fn(async () => credential) }) + expect(bad.quota.connection).toBe('transientFailure') + const clientError = vi.fn(async () => new Response('', { status: 404 })) + const worse = await fetchGeminiQuota({ fetch: clientError, readFile: vi.fn(async () => credential) }) + expect(worse.quota.connection).toBe('terminalFailure') + }) + + 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 rawtoken ya29.leak eyJabc.def.ghi\0tail') }) + const result = await fetchGeminiQuota({ fetch: fetchMock, readFile: vi.fn(async () => credential) }) + const logged = warn.mock.calls.flat().join(' ') + expect(result.quota).not.toHaveProperty('error') + expect(logged).not.toMatch(/rawtoken|ya29\.leak|eyJabc|\0/) + expect(logged).toContain('[REDACTED]') + }) +}) diff --git a/app/electron/quota/gemini.ts b/app/electron/quota/gemini.ts new file mode 100644 index 00000000..3576b950 --- /dev/null +++ b/app/electron/quota/gemini.ts @@ -0,0 +1,206 @@ +// Live Gemini quota via the OAuth-backed Code Assist APIs the Gemini CLI +// itself calls (prior art: steipete/CodexBar docs/gemini.md, which derives the +// flow from the CLI's own traffic): +// +// - POST https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist +// body { metadata: { ideType: 'GEMINI_CLI', pluginType: 'GEMINI' } } +// → tier (plan label), cloudaicompanionProject (quota project) +// - POST https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota +// body { project } (or {} when unknown) → per-model quota buckets +// - POST https://oauth2.googleapis.com/token (only when the stored token is +// stale and GEMINI_OAUTH_CLIENT_ID/GEMINI_OAUTH_CLIENT_SECRET are set; +// the refreshed token stays in memory — the Gemini CLI owns its file) +// +// Credential: the Gemini CLI's own ~/.gemini/oauth_creds.json, read-only. +import os from 'node:os' +import path from 'node:path' + +import { fraction, quotaRequestSignal, readSecureFile, sanitizeError } from './security' +import type { QuotaProvider, QuotaWindow } from './types' + +const CODE_ASSIST_ENDPOINT = 'https://cloudcode-pa.googleapis.com/v1internal' +const TOKEN_ENDPOINT = 'https://oauth2.googleapis.com/token' + +type GeminiCredential = { + access_token?: string + refresh_token?: string + id_token?: string + expiry_date?: number +} + +export type GeminiDeps = { + fetch: typeof fetch + credentialPath: string + readFile: typeof readSecureFile + now: () => number +} + +const defaults: GeminiDeps = { + fetch: globalThis.fetch, + credentialPath: path.join(os.homedir(), '.gemini', 'oauth_creds.json'), + readFile: readSecureFile, + now: Date.now, +} + +function empty(connection: QuotaProvider['connection'], footerLines: string[] = []): QuotaProvider { + return { provider: 'gemini', connection, primary: null, details: [], planLabel: null, footerLines } +} + +function parseCredential(raw: string): GeminiCredential | null { + const parsed = JSON.parse(raw) as GeminiCredential + if (!parsed || typeof parsed.access_token !== 'string' || parsed.access_token.length === 0) return null + return parsed +} + +async function credentialFromFile(deps: GeminiDeps): Promise { + const raw = await deps.readFile(deps.credentialPath, 64 * 1024) + return raw ? parseCredential(raw) : null +} + +// Client credentials live inside the installed Gemini CLI bundle; scanning the +// install is out of scope here, so only the CLI's documented env overrides are +// honored. Without them an expired token is used as-is and a 401 degrades. +function clientCredentials(): { clientId: string; clientSecret: string } | null { + const clientId = process.env['GEMINI_OAUTH_CLIENT_ID'] + const clientSecret = process.env['GEMINI_OAUTH_CLIENT_SECRET'] + return clientId && clientSecret ? { clientId, clientSecret } : null +} + +async function refresh(credential: GeminiCredential, deps: GeminiDeps, signal?: AbortSignal): Promise { + const client = clientCredentials() + const refreshToken = credential.refresh_token + if (!client || !refreshToken) return null + const body = new URLSearchParams({ + client_id: client.clientId, + client_secret: client.clientSecret, + refresh_token: refreshToken, + grant_type: 'refresh_token', + }) + const response = await deps.fetch(TOKEN_ENDPOINT, { method: 'POST', signal: quotaRequestSignal(signal), body }) + if (!response.ok) return null + const next = await response.json() as Record + return typeof next.access_token === 'string' && next.access_token ? next.access_token : null +} + +// The `hd` claim (hosted domain) separates Google Workspace logins from +// personal ones, mirroring CodexBar's tier labeling. +function workspaceClaim(idToken: string | undefined): boolean { + if (!idToken) return false + const [, payload] = idToken.split('.') + if (!payload) return false + try { + const json = JSON.parse(Buffer.from(payload.replace(/-/g, '+').replace(/_/g, '/'), 'base64').toString('utf8')) as Record + return typeof json.hd === 'string' && json.hd.length > 0 + } catch { + return false + } +} + +function tierLabel(body: Record, credential: GeminiCredential): string | null { + // paidTier wins whenever present; otherwise whatever tier the account sits on. + const raw = body.paidTier?.name ?? body.currentTier?.name ?? body.currentTier?.id + if (typeof raw !== 'string' || !raw.trim()) return null + const lower = raw.trim().toLowerCase() + if (lower === 'standard-tier') return 'Paid' + if (lower === 'legacy-tier') return 'Legacy' + if (lower.includes('free') && workspaceClaim(credential.id_token)) return 'Workspace' + if (lower.includes('free')) return 'Free' + return raw.trim() +} + +function windowOf(bucket: Record): QuotaWindow | null { + const modelId = bucket.modelId + const remaining = fraction(typeof bucket.remainingFraction === 'number' ? bucket.remainingFraction * 100 : NaN) + if (typeof modelId !== 'string' || !modelId || remaining === null) return null + const resetTime = bucket.resetTime + const resetsAt = typeof resetTime === 'string' && !Number.isNaN(Date.parse(resetTime)) + ? new Date(resetTime).toISOString() : null + return { label: modelId, percent: 1 - remaining, resetsAt } +} + +export function decodeGeminiUsage(body: unknown): QuotaProvider { + const data = body && typeof body === 'object' ? body as Record : {} + const buckets = Array.isArray(data.buckets) ? data.buckets : [] + const details = buckets + .filter((bucket): bucket is Record => Boolean(bucket) && typeof bucket === 'object') + .map(windowOf) + .filter((window): window is QuotaWindow => window !== null) + .sort((a, b) => b.percent - a.percent) + return { + provider: 'gemini', connection: 'connected', + primary: details[0] ?? null, + details, + planLabel: null, + footerLines: [], + } +} + +async function post(token: string, path: string, payload: Record, deps: GeminiDeps, parent?: AbortSignal): Promise { + return deps.fetch(`${CODE_ASSIST_ENDPOINT}:${path}`, { + method: 'POST', signal: quotaRequestSignal(parent), + headers: { Authorization: `Bearer ${token}`, 'Content-Type': 'application/json', Accept: 'application/json', 'User-Agent': 'CodeBurn' }, + body: JSON.stringify(payload), + }) +} + +// Google flags retired consumer tiers (June 2026 shutdown of individual/AI +// Pro/Ultra OAuth access) with these sentinels instead of ordinary errors. +function migrationFooter(body: unknown): string[] { + const error = body && typeof body === 'object' ? (body as Record).error : null + const status = typeof error?.status === 'string' ? error.status : '' + const message = typeof error?.message === 'string' ? error.message : '' + const deprecated = status === 'UNSUPPORTED_CLIENT' || message.includes('IneligibleTierError') + return deprecated ? ['Google retired Gemini CLI OAuth for this account tier — use Antigravity.'] : [] +} + +export type GeminiResult = { quota: QuotaProvider; retryAfterSeconds?: number } + +export async function fetchGeminiQuota(options: Partial & { signal?: AbortSignal } = {}): Promise { + const deps = { ...defaults, ...options } + try { + let credential = await credentialFromFile(deps) + if (!credential) return { quota: empty('disconnected') } + + if ((credential.expiry_date ?? Infinity) - deps.now() <= 5 * 60_000) { + const accessToken = await refresh(credential, deps, options.signal) + if (accessToken) credential = { ...credential, access_token: accessToken } + } + const token = credential.access_token! + let response = await post(token, 'loadCodeAssist', { metadata: { ideType: 'GEMINI_CLI', pluginType: 'GEMINI' } }, deps, options.signal) + if (response.status === 401) { + const reread = await credentialFromFile(deps) + if (!reread || reread.access_token === credential.access_token) return { quota: empty('transientFailure') } + credential = reread + response = await post(credential.access_token!, 'loadCodeAssist', { metadata: { ideType: 'GEMINI_CLI', pluginType: 'GEMINI' } }, deps, options.signal) + } + if (response.status === 429) { + const raw = response.headers.get('Retry-After') + const seconds = raw === null ? NaN : Number(raw) + return { quota: empty('transientFailure'), retryAfterSeconds: Math.max(Number.isFinite(seconds) ? Math.ceil(seconds) : 300, 60) } + } + // Retired-tier sentinels ride inside the JSON error body of a non-200 + // response, so parse before branching on status. + const assist = await response.json().catch(() => ({})) as Record + const migration = migrationFooter(assist) + if (migration.length > 0) return { quota: empty('terminalFailure', migration) } + if (!response.ok) return { quota: empty(response.status >= 400 && response.status < 500 ? 'terminalFailure' : 'transientFailure') } + + const project = typeof assist.cloudaicompanionProject === 'string' && assist.cloudaicompanionProject + ? assist.cloudaicompanionProject : undefined + const quotaResponse = await post(credential.access_token!, 'retrieveUserQuota', project ? { project } : {}, deps, options.signal) + if (quotaResponse.status === 429) { + const raw = quotaResponse.headers.get('Retry-After') + const seconds = raw === null ? NaN : Number(raw) + return { quota: empty('transientFailure'), retryAfterSeconds: Math.max(Number.isFinite(seconds) ? Math.ceil(seconds) : 300, 60) } + } + if (!quotaResponse.ok) { + return { quota: empty(quotaResponse.status >= 400 && quotaResponse.status < 500 ? 'terminalFailure' : 'transientFailure') } + } + const quota = decodeGeminiUsage(await quotaResponse.json()) + quota.planLabel = tierLabel(assist, credential) + return { quota } + } catch (error) { + console.warn(`Gemini quota unavailable: ${sanitizeError(error)}`) + return { quota: empty('transientFailure') } + } +} diff --git a/app/electron/quota/index.test.ts b/app/electron/quota/index.test.ts index 7ff48f99..88a86487 100644 --- a/app/electron/quota/index.test.ts +++ b/app/electron/quota/index.test.ts @@ -1,13 +1,36 @@ import { describe, expect, it, vi } from 'vitest' import { QuotaService } from './index' -import type { QuotaProvider } from './types' +import type { ProviderName, QuotaProvider } from './types' -const quota = (provider: 'claude' | 'codex'): QuotaProvider => ({ +const quota = (provider: ProviderName): QuotaProvider => ({ provider, connection: 'connected', primary: null, details: [], planLabel: null, footerLines: [], }) +// Every construction stubs all five fetchers: a missing dep falls back to the +// real fetcher, which would touch disk or the network inside a test. +const noopFetchers = () => ({ + claude: vi.fn(async () => ({ quota: quota('claude') })), + codex: vi.fn(async () => ({ quota: quota('codex') })), + gemini: vi.fn(async () => ({ quota: quota('gemini') })), + copilot: vi.fn(async () => ({ quota: quota('copilot') })), + antigravity: vi.fn(async () => ({ quota: quota('antigravity') })), +}) + describe('QuotaService', () => { + it('fetches and returns every registered provider', async () => { + const fetchers = noopFetchers() + const service = new QuotaService({ + ...fetchers, now: () => 1000, + readFile: vi.fn(async () => null), writeFile: vi.fn(async () => undefined), + }) + const results = await service.getQuota({ force: true }) + expect(results.map(row => row.provider)).toEqual(['claude', 'codex', 'gemini', 'copilot', 'antigravity']) + for (const fetcher of Object.values(fetchers)) expect(fetcher).toHaveBeenCalledTimes(1) + // Antigravity is local-only; it must not receive keychain permission. + expect(fetchers.antigravity).toHaveBeenCalledWith({ signal: expect.any(AbortSignal), allowKeychain: false }) + }) + // The snap declares no Codex credential path, because the live gauge would // need write access to the Codex CLI's own auth.json to rotate the token. // Under $SNAP the Codex fetch must not run at all; Claude is unaffected. @@ -15,70 +38,85 @@ describe('QuotaService', () => { const previous = process.env['SNAP'] process.env['SNAP'] = '/snap/codeburn/current' try { - const claude = vi.fn(async () => ({ quota: quota('claude') })) - const codex = vi.fn(async () => ({ quota: quota('codex') })) + const fetchers = noopFetchers() const service = new QuotaService({ - claude, codex, now: () => Date.parse('2026-08-14T00:00:00Z'), + ...fetchers, now: () => Date.parse('2026-08-14T00:00:00Z'), readFile: vi.fn(async () => null), writeFile: vi.fn(async () => {}), statePath: '/mock/backoff.json', }) - const [claudeQuota, codexQuota] = await service.getQuota({ force: true }) - expect(codex).not.toHaveBeenCalled() - expect(claude).toHaveBeenCalledTimes(1) - expect(codexQuota?.connection).toBe('disconnected') - expect(claudeQuota?.connection).toBe('connected') + const results = await service.getQuota({ force: true }) + expect(fetchers.codex).not.toHaveBeenCalled() + expect(fetchers.claude).toHaveBeenCalledTimes(1) + expect(results.find(row => row.provider === 'codex')?.connection).toBe('disconnected') + expect(results.find(row => row.provider === 'claude')?.connection).toBe('connected') } finally { if (previous === undefined) delete process.env['SNAP'] else process.env['SNAP'] = previous } }) + it('omits disabled providers from polling and results', async () => { + const fetchers = noopFetchers() + const service = new QuotaService({ + ...fetchers, now: () => 1000, + readFile: vi.fn(async () => null), writeFile: vi.fn(async () => undefined), + }) + // Unknown names are ignored rather than throwing. + const results = await service.getQuota({ force: true, disabled: ['gemini', 'copilot', 'bogus' as ProviderName] }) + expect(results.map(row => row.provider)).toEqual(['claude', 'codex', 'antigravity']) + expect(fetchers.gemini).not.toHaveBeenCalled() + expect(fetchers.copilot).not.toHaveBeenCalled() + }) + 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 fetchers = noopFetchers() + fetchers.claude.mockImplementation(async () => ({ quota: quota('claude'), retryAfterSeconds: 60 })) + fetchers.gemini.mockImplementation(async () => ({ quota: quota('gemini'), retryAfterSeconds: 120 })) const service = new QuotaService({ - claude, codex, now: () => Date.parse('2026-07-12T00:00:00Z'), + ...fetchers, 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]!) + const saved = JSON.parse(writes.at(-1)!) expect(saved.claude).toBe('2026-07-12T00:01:00.000Z') + expect(saved.gemini).toBe('2026-07-12T00:02:00.000Z') await service.getQuota({ force: true }) - expect(claude).toHaveBeenCalledTimes(1) - expect(codex).toHaveBeenCalledTimes(2) + expect(fetchers.claude).toHaveBeenCalledTimes(1) + expect(fetchers.gemini).toHaveBeenCalledTimes(1) + expect(fetchers.codex).toHaveBeenCalledTimes(2) }) it('force re-fetches within the cache window by invalidating first', async () => { - const claude = vi.fn(async () => ({ quota: quota('claude') })) - const codex = vi.fn(async () => ({ quota: quota('codex') })) + const fetchers = noopFetchers() const service = new QuotaService({ - claude, codex, now: () => 1000, refreshMs: 120_000, + ...fetchers, now: () => 1000, refreshMs: 120_000, readFile: vi.fn(async () => null), writeFile: vi.fn(async () => undefined), }) await service.getQuota() await service.getQuota() // fresh cache, no re-fetch - expect(claude).toHaveBeenCalledTimes(1) + expect(fetchers.claude).toHaveBeenCalledTimes(1) await service.getQuota({ force: true }) // force clears the still-fresh cache - expect(claude).toHaveBeenCalledTimes(2) + expect(fetchers.claude).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 fetchers = noopFetchers() + fetchers.claude.mockImplementation(async () => { await pending; return { quota: quota('claude') } }) const service = new QuotaService({ - claude, codex: vi.fn(async () => ({ quota: quota('codex') })), + ...fetchers, 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) + expect(fetchers.claude).toHaveBeenCalledTimes(1) }) }) diff --git a/app/electron/quota/index.ts b/app/electron/quota/index.ts index 1fa0b7e9..a716a0bd 100644 --- a/app/electron/quota/index.ts +++ b/app/electron/quota/index.ts @@ -1,8 +1,11 @@ import os from 'node:os' import path from 'node:path' +import { fetchAntigravityQuota } from './antigravity' import { fetchClaudeQuota } from './claude' import { fetchCodexQuota } from './codex' +import { fetchCopilotQuota } from './copilot' +import { fetchGeminiQuota } from './gemini' import { atomicWriteSecureFile, readSecureFile, sanitizeError } from './security' import type { ProviderName, QuotaProvider } from './types' @@ -11,9 +14,13 @@ export { sanitizeError } from './security' type Blocked = Partial> type FetchResult = { quota: QuotaProvider; retryAfterSeconds?: number } +type FetcherOptions = { signal: AbortSignal; allowKeychain: boolean } type QuotaDeps = { - claude: (options: { signal: AbortSignal; allowKeychain: boolean }) => Promise - codex: (options: { signal: AbortSignal; allowKeychain: boolean }) => Promise + claude: (options: FetcherOptions) => Promise + codex: (options: FetcherOptions) => Promise + gemini: (options: FetcherOptions) => Promise + copilot: (options: FetcherOptions) => Promise + antigravity: (options: FetcherOptions) => Promise statePath: string readFile: typeof readSecureFile writeFile: typeof atomicWriteSecureFile @@ -21,9 +28,16 @@ type QuotaDeps = { refreshMs: number } +const PROVIDERS: ProviderName[] = ['claude', 'codex', 'gemini', 'copilot', 'antigravity'] + const defaultDeps: QuotaDeps = { claude: fetchClaudeQuota, codex: fetchCodexQuota, + gemini: fetchGeminiQuota, + copilot: fetchCopilotQuota, + // The Antigravity probe talks only to localhost surfaces (no credentials, + // no remote endpoints), so it ignores the abort/keychain options entirely. + antigravity: async () => ({ quota: await fetchAntigravityQuota() }), statePath: path.join(os.homedir(), '.codeburn', 'quota-backoff.json'), readFile: readSecureFile, writeFile: atomicWriteSecureFile, @@ -54,13 +68,13 @@ 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 generations: Record = Object.fromEntries(PROVIDERS.map(p => [p, 0])) as Record private controllers: Partial> = {} constructor(deps: Partial = {}) { this.deps = { ...defaultDeps, ...deps } } invalidate(provider?: ProviderName): void { - const providers: ProviderName[] = provider ? [provider] : ['claude', 'codex'] + const providers: ProviderName[] = provider ? [provider] : PROVIDERS for (const p of providers) { this.generations[p] += 1 this.controllers[p]?.abort() @@ -69,11 +83,13 @@ export class QuotaService { this.cache = null } - async getQuota(options: { force?: boolean; allowKeychain?: boolean } = {}): Promise { + async getQuota(options: { force?: boolean; allowKeychain?: boolean; disabled?: string[] } = {}): Promise { if (options.force) this.invalidate() 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.allowKeychain)).finally(() => { this.flight = null }) + // IPC names are untrusted strings; only known providers may be skipped. + const disabled = new Set((options.disabled ?? []).filter((p): p is ProviderName => PROVIDERS.includes(p as ProviderName))) + this.flight = this.fetchAll(Boolean(options.allowKeychain), disabled).finally(() => { this.flight = null }) return this.flight } @@ -92,7 +108,7 @@ export class QuotaService { catch (error) { console.warn(`Quota backoff state not saved: ${sanitizeError(error)}`) } } - private async fetchAll(allowKeychain: boolean): Promise { + private async fetchAll(allowKeychain: boolean, disabled: Set): Promise { const startingGenerations = { ...this.generations } const prior = this.cache?.value ?? [] const blocked = await this.readBlocked() @@ -112,9 +128,7 @@ export class QuotaService { 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 }) - : await this.deps.codex({ signal: controller.signal, allowKeychain }) + const result = await this.deps[provider]({ signal: controller.signal, allowKeychain }) 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() @@ -128,11 +142,12 @@ export class QuotaService { if (this.controllers[provider] === controller) this.controllers[provider] = undefined return retainOnFailure(result.quota) } - const value = await Promise.all([ - run('claude'), - codexQuotaSupported() ? run('codex') : Promise.resolve(unavailable('codex', 'disconnected')), - ]) - if (startingGenerations.claude === this.generations.claude && startingGenerations.codex === this.generations.codex) { + const value = await Promise.all(PROVIDERS.filter(provider => !disabled.has(provider)).map(provider => + provider === 'codex' && !codexQuotaSupported() + ? Promise.resolve(unavailable('codex', 'disconnected')) + : run(provider), + )) + if (PROVIDERS.every(p => startingGenerations[p] === this.generations[p])) { this.cache = { at: this.deps.now(), value } } return value @@ -144,5 +159,5 @@ export const quotaService = new QuotaService() // them on a user-initiated forced refresh (the Connect / Refresh affordance). // Background polls skip the keychain and lean on retainOnFailure to hold a // live connection steady between forced refreshes. -export const getQuota = (options: { force?: boolean } = {}): Promise => - quotaService.getQuota({ force: options.force, allowKeychain: Boolean(options.force) }) +export const getQuota = (options: { force?: boolean; disabled?: string[] } = {}): Promise => + quotaService.getQuota({ force: options.force, allowKeychain: Boolean(options.force), disabled: options.disabled }) diff --git a/app/electron/quota/security.test.ts b/app/electron/quota/security.test.ts index 6df3f0f5..11f9b3ec 100644 --- a/app/electron/quota/security.test.ts +++ b/app/electron/quota/security.test.ts @@ -1,6 +1,6 @@ import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' -import { readKeychainPassword } from './security' +import { readKeychainPassword, sanitizeError } from './security' // readKeychainPassword short-circuits off darwin; pin the platform so the // classification logic is exercised on any CI host. @@ -55,3 +55,10 @@ describe('readKeychainPassword', () => { expect(await readKeychainPassword('svc', ['alice', null], exec)).toEqual({ status: 'accessDenied' }) }) }) + +describe('sanitizeError', () => { + it('redacts the OAuth token shapes used by the quota providers', () => { + const sanitized = sanitizeError(new Error('Bearer abc123 ya29.google-token gho_ghub eyJwt.sig\0x sk-ant-one sk-two')) + expect(sanitized).toBe('[REDACTED] [REDACTED] [REDACTED] [REDACTED] [REDACTED] [REDACTED]') + }) +}) diff --git a/app/electron/quota/security.ts b/app/electron/quota/security.ts index 3637f06b..060a09de 100644 --- a/app/electron/quota/security.ts +++ b/app/electron/quota/security.ts @@ -81,9 +81,13 @@ 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(/Bearer\s+[^\s,;"']+/gi, '[REDACTED]') .replace(/sk-ant-[A-Za-z0-9_-]+/gi, '[REDACTED]') .replace(/sk-[A-Za-z0-9_-]+/gi, '[REDACTED]') + // Google (ya29.) and GitHub (gho_/ghu_/ghp_) OAuth token shapes used by + // the Gemini/Copilot quota providers. + .replace(/ya29\.[A-Za-z0-9._-]+/g, '[REDACTED]') + .replace(/gh[opusr]_[A-Za-z0-9_]+/g, '[REDACTED]') .replace(/eyJ[A-Za-z0-9._-]+/g, '[REDACTED]') .slice(0, 240) } diff --git a/app/electron/quota/types.ts b/app/electron/quota/types.ts index 981025b7..6e342998 100644 --- a/app/electron/quota/types.ts +++ b/app/electron/quota/types.ts @@ -5,7 +5,7 @@ export type QuotaWindow = { } export type QuotaProvider = { - provider: 'claude' | 'codex' + provider: 'claude' | 'codex' | 'gemini' | 'copilot' | 'antigravity' connection: 'connected' | 'disconnected' | 'accessDenied' | 'loading' | 'stale' | 'transientFailure' | 'terminalFailure' primary: QuotaWindow | null details: QuotaWindow[] diff --git a/app/renderer/components/ConnectAffordance.tsx b/app/renderer/components/ConnectAffordance.tsx index c58bb71f..fcaaca54 100644 --- a/app/renderer/components/ConnectAffordance.tsx +++ b/app/renderer/components/ConnectAffordance.tsx @@ -1,12 +1,18 @@ import { useState } from 'react' +import { PROVIDER_NAMES } from '../lib/providers' import type { QuotaProvider } from '../lib/types' // The exact terminal login command per provider. No interactive login is // attempted from the app — we only show the command to copy and a Refresh. -const LOGIN: Record = { +// Providers without a CLI login (Copilot signs in from an editor plugin; +// Antigravity is local-only) get a note instead of a command. +const LOGIN: Record = { claude: { command: 'claude', hint: 'then type /login' }, codex: { command: 'codex login' }, + gemini: { command: 'gemini', hint: 'then sign in when prompted' }, + copilot: { note: 'Sign in to GitHub Copilot in your editor (VS Code or JetBrains), then Refresh.' }, + antigravity: { note: 'Open Antigravity and sign in, then Refresh — quota comes from its local server only.' }, } /** Inline "Connect" affordance for a disconnected or access-denied provider: a @@ -18,7 +24,7 @@ export function ConnectAffordance({ provider, connection, onRefresh }: { onRefresh: () => void }) { const [open, setOpen] = useState(false) - const name = provider === 'claude' ? 'Claude' : 'Codex' + const name = PROVIDER_NAMES[provider] const message = connection === 'accessDenied' ? 'Keychain access needed: click Allow when macOS asks, then Refresh.' : `Not connected. Log in with the ${name} CLI.` @@ -30,8 +36,14 @@ export function ConnectAffordance({ provider, connection, onRefresh }: { {open && (
-

Sign in from a terminal, then Refresh:

-

{login.command}{login.hint ? {login.hint} : null}

+ {login.command ? ( + <> +

Sign in from a terminal, then Refresh:

+

{login.command}{login.hint ? {login.hint} : null}

+ + ) : ( +

{login.note}

+ )} {connection === 'accessDenied' &&

Already logged in? Click Allow when macOS asks for keychain access.

}
diff --git a/app/renderer/lib/providers.ts b/app/renderer/lib/providers.ts new file mode 100644 index 00000000..70ee8246 --- /dev/null +++ b/app/renderer/lib/providers.ts @@ -0,0 +1,39 @@ +import type { ProviderName } from './types' + +export const PROVIDER_NAMES: Record = { + claude: 'Claude', + codex: 'Codex', + gemini: 'Gemini', + copilot: 'Copilot', + antigravity: 'Antigravity', +} + +/** Company named in honest copy like "Anthropic rate limited the quota endpoint". */ +export const PROVIDER_OWNERS: Record = { + claude: 'Anthropic', + codex: 'OpenAI', + gemini: 'Google', + copilot: 'GitHub', + antigravity: 'Google', +} + +const ALL_PROVIDERS = Object.keys(PROVIDER_NAMES) as ProviderName[] +const DISABLED_KEY = 'codeburn.quotaDisabled' + +/** Display order for quota rows (matches the electron poll order). */ +export const QUOTA_PROVIDERS = Object.keys(PROVIDER_NAMES) as ProviderName[] + +export function readDisabledProviders(): ProviderName[] { + try { + const raw = globalThis.localStorage?.getItem(DISABLED_KEY) + if (!raw) return [] + const parsed: unknown = JSON.parse(raw) + return Array.isArray(parsed) ? parsed.filter((p): p is ProviderName => typeof p === 'string' && ALL_PROVIDERS.includes(p as ProviderName)) : [] + } catch { + return [] + } +} + +export function writeDisabledProviders(disabled: ProviderName[]): void { + try { globalThis.localStorage?.setItem(DISABLED_KEY, JSON.stringify(disabled)) } catch { /* storage can be unavailable in hardened contexts */ } +} diff --git a/app/renderer/lib/types.ts b/app/renderer/lib/types.ts index 82614631..813edd91 100644 --- a/app/renderer/lib/types.ts +++ b/app/renderer/lib/types.ts @@ -35,7 +35,7 @@ export type QuotaWindow = { } export type QuotaProvider = { - provider: 'claude' | 'codex' + provider: 'claude' | 'codex' | 'gemini' | 'copilot' | 'antigravity' connection: 'connected' | 'disconnected' | 'accessDenied' | 'loading' | 'stale' | 'transientFailure' | 'terminalFailure' primary: QuotaWindow | null details: QuotaWindow[] @@ -45,6 +45,8 @@ export type QuotaProvider = { rateLimited?: boolean } +export type ProviderName = QuotaProvider['provider'] + // ————— src/menubar-json.ts ————— export type DailyModelBreakdown = { @@ -668,7 +670,7 @@ export interface CodeburnBridge { getUpdateStatus(): Promise /** Subscribe to pushed update-availability status; returns an unsubscribe fn. */ onUpdateStatus(cb: (status: UpdateStatus) => void): () => void - getQuota(force?: boolean): Promise + getQuota(force?: boolean, disabled?: ProviderName[]): Promise // `background` (prefetch only) requests background CLI-spawn priority; optional // so an older preload that ignores it degrades to interactive priority. // `scope` selects local-device usage ('local', default) or paired-device diff --git a/app/renderer/sections/Plans.test.tsx b/app/renderer/sections/Plans.test.tsx index 18db5665..251679a7 100644 --- a/app/renderer/sections/Plans.test.tsx +++ b/app/renderer/sections/Plans.test.tsx @@ -223,11 +223,11 @@ describe('Plans', () => { const { rerender } = render() await screen.findByText('Max 20x') - expect(getQuota).toHaveBeenCalledWith(false) // mount is a steady poll + expect(getQuota).toHaveBeenCalledWith(false, []) // mount is a steady poll getQuota.mockClear() rerender() // manual refresh bumps the token - await waitFor(() => expect(getQuota).toHaveBeenCalledWith(true)) + await waitFor(() => expect(getQuota).toHaveBeenCalledWith(true, [])) getQuota.mockClear() rerender() // unchanged token must not re-force @@ -255,7 +255,7 @@ describe('Plans', () => { getQuota.mockClear() fireEvent.click(screen.getByRole('button', { name: 'Refresh' })) - await waitFor(() => expect(getQuota).toHaveBeenCalledWith(true)) + await waitFor(() => expect(getQuota).toHaveBeenCalledWith(true, [])) }) it('renders the honest rate-limited note on a 429 backoff, per provider owner', async () => { diff --git a/app/renderer/sections/Plans.tsx b/app/renderer/sections/Plans.tsx index 127944d2..b011b756 100644 --- a/app/renderer/sections/Plans.tsx +++ b/app/renderer/sections/Plans.tsx @@ -10,6 +10,7 @@ import { usePolled } from '../hooks/usePolled' import { formatConverted } from '../lib/format' import { codeburn } from '../lib/ipc' import { motionClass } from '../lib/motion' +import { PROVIDER_NAMES, PROVIDER_OWNERS, readDisabledProviders } from '../lib/providers' import type { JsonPlanSummary, Period, PlanId, PlanProvider, QuotaProvider, QuotaWindow, StatusJson } from '../lib/types' import type { SettingsPane } from './Settings' @@ -33,8 +34,7 @@ function fmtPct(n: number): string { /** Honest copy for a 429 backoff window (the upstream quota endpoint rate * limited us), replacing the generic "waiting" note. */ export function rateLimitedNote(provider: QuotaProvider['provider']): string { - const owner = provider === 'claude' ? 'Anthropic' : 'OpenAI' - return `${owner} rate limited the quota endpoint, retrying in a few minutes` + return `${PROVIDER_OWNERS[provider]} rate limited the quota endpoint, retrying in a few minutes` } function cycleEndDate(plan: JsonPlanSummary): Date | null { @@ -79,7 +79,7 @@ export function Plans({ period, refreshToken = 0, onNavigate, ready = true }: { const key = `${refreshToken}:${reconnectNonce}` const force = key !== lastForced.current lastForced.current = key - return codeburn.getQuota(force) + return codeburn.getQuota(force, readDisabledProviders()) }, [refreshToken, reconnectNonce]) const reconnect = () => setReconnectNonce(value => value + 1) const budgetReport = usePolled(() => codeburn.getPlans(period), [period, refreshToken], { enabled: ready }) @@ -146,7 +146,7 @@ function renderBudgetPlans(data: StatusJson | null, error: ReturnType void }) { - const providerName = quota.provider === 'claude' ? 'Claude' : 'Codex' + const providerName = PROVIDER_NAMES[quota.provider] return ( { expect(await screen.findByText('Detected subscriptions')).toBeInTheDocument() expect(screen.getByText('Max 20x')).toBeInTheDocument() expect(screen.getByText('Not connected. Log in with the Codex CLI.')).toBeInTheDocument() - expect(mocks.getQuota).toHaveBeenCalledWith(false) + expect(mocks.getQuota).toHaveBeenCalledWith(false, []) }) it('expands the DetectedRow Connect affordance and forces a keychain refresh', async () => { @@ -268,7 +268,7 @@ describe('Settings', () => { expect(screen.getByText('codex login')).toBeInTheDocument() mocks.getQuota.mockClear() await user.click(screen.getByRole('button', { name: 'Refresh' })) - await waitFor(() => expect(mocks.getQuota).toHaveBeenCalledWith(true)) + await waitFor(() => expect(mocks.getQuota).toHaveBeenCalledWith(true, [])) }) it('offers only non-OAuth budget presets; Claude and Codex are excluded', async () => { diff --git a/app/renderer/sections/Settings.tsx b/app/renderer/sections/Settings.tsx index 41ddcf46..7f19bde1 100644 --- a/app/renderer/sections/Settings.tsx +++ b/app/renderer/sections/Settings.tsx @@ -15,11 +15,12 @@ import { formatConverted, formatUsd } from '../lib/format' import { codeburn } from '../lib/ipc' import { shortcutLabel } from '../lib/platform' import { motionClass } from '../lib/motion' +import { PROVIDER_NAMES, QUOTA_PROVIDERS, readDisabledProviders, writeDisabledProviders } from '../lib/providers' import { REFRESH_OPTIONS, useRefreshCadence } from '../lib/refreshCadence' import { showToast } from '../lib/toast' import { ToastHost } from '../components/ToastHost' import { rateLimitedNote } from './Plans' -import type { ActionResult, AliasRow, ClaudeConfigSelector, CliError, CombinedUsage, DeviceScanResult, Identity, JsonPlanSummary, MenubarPayload, Period, PlanId, PlanProvider, PriceOverrideList, PriceOverrideRow, PriceRates, QuotaProvider, Scope, ShareStatus, StatusJson, TelemetryStatus } from '../lib/types' +import type { ActionResult, AliasRow, ClaudeConfigSelector, CliError, CombinedUsage, DeviceScanResult, Identity, JsonPlanSummary, MenubarPayload, Period, PlanId, PlanProvider, PriceOverrideList, PriceOverrideRow, PriceRates, ProviderName, QuotaProvider, Scope, ShareStatus, StatusJson, TelemetryStatus } from '../lib/types' export type SettingsPane = 'general' | 'providers' | 'aliases' | 'pricing' | 'plans' | 'devices' | 'export' | 'privacy' type Pane = SettingsPane @@ -327,16 +328,18 @@ function planSummaries(status: StatusJson): JsonPlanSummary[] { return status.plan ? [status.plan] : [] } -function DetectedRow({ quota, onReconnect }: { quota: QuotaProvider; onReconnect: () => void }) { - const name = quota.provider === 'claude' ? 'Claude' : 'Codex' +function DetectedRow({ quota, enabled, onToggle, onReconnect }: { quota: QuotaProvider; enabled: boolean; onToggle: () => void; onReconnect: () => void }) { return
- {name} - {quota.connection === 'disconnected' || quota.connection === 'accessDenied' + {PROVIDER_NAMES[quota.provider]} + {!enabled + ? Off + : quota.connection === 'disconnected' || quota.connection === 'accessDenied' ?
: quota.rateLimited ? {rateLimitedNote(quota.provider)} : {quota.planLabel ?? 'Connected'}} +
} @@ -345,13 +348,14 @@ function PlansPane({ period, refreshToken, onNavigate, onConfigMutated }: { peri // Steady poll serves cached quota (force=false); the Connect affordance's // Refresh forces a keychain-allowed fetch via the same path as Plans.tsx. const [reconnectNonce, setReconnectNonce] = useState(0) + const [disabledProviders, setDisabledProviders] = useState(() => readDisabledProviders()) const lastForced = useRef(`${refreshToken}:${reconnectNonce}`) const quota = usePolled(() => { const key = `${refreshToken}:${reconnectNonce}` const force = key !== lastForced.current lastForced.current = key - return codeburn.getQuota(force) - }, [refreshToken, reconnectNonce]) + return codeburn.getQuota(force, disabledProviders) + }, [refreshToken, reconnectNonce, disabledProviders]) const plans = usePolled(() => codeburn.getPlans(period), [period, refreshToken, nonce]) const [presetId, setPresetId] = useState(MANUAL_PLAN_PRESETS[0]!.id) const configured = plans.data ? planSummaries(plans.data) : [] @@ -363,6 +367,17 @@ function PlansPane({ period, refreshToken, onNavigate, onConfigMutated }: { peri const remove = (plan: JsonPlanSummary) => { void codeburn.resetPlan(plan.provider).then(finish) } + // Toggling a provider off stops polling it entirely (the main process never + // contacts its endpoints); toggling on forces a fresh fetch so the row + // repopulates immediately. + const toggleProvider = (provider: ProviderName) => { + const next = disabledProviders.includes(provider) + ? disabledProviders.filter(item => item !== provider) + : [...disabledProviders, provider] + writeDisabledProviders(next) + setDisabledProviders(next) + setReconnectNonce(value => value + 1) + } const add = () => { const preset = MANUAL_PLAN_PRESETS.find(item => item.id === presetId)! void codeburn.setPlan(preset.id, preset.provider).then(finish) @@ -373,7 +388,17 @@ function PlansPane({ period, refreshToken, onNavigate, onConfigMutated }: { peri
Detected subscriptions
- {quota.error && !quota.data ? : !quota.data ?

Detecting subscriptions…

: quota.data.length === 0 ?

No detectable subscriptions.

: quota.data.map(provider => setReconnectNonce(value => value + 1)} />)} + {quota.error && !quota.data ? : QUOTA_PROVIDERS.map(provider => { + const row = quota.data?.find(item => item.provider === provider) + if (!row && !disabledProviders.includes(provider)) return null + return toggleProvider(provider)} + onReconnect={() => setReconnectNonce(value => value + 1)} + /> + })}