From faf949835fc0b9f5958d17ba99b573caa0c649d8 Mon Sep 17 00:00:00 2001 From: Alexander Saavedra Date: Thu, 30 Apr 2026 01:55:32 +0200 Subject: [PATCH] feat(rate-limits): add Gemini and OpenCode Go usage tracking (#1023) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(rate-limits): add Gemini and OpenCode Go usage tracking - Extend provider union and RateLimitState with 'gemini' and 'opencode-go' - Add opencodeSessionCookie to GlobalSettings (password input in GeneralPane) - Add GeminiIcon and OpenCodeGoIcon to status bar icons - Implement gemini-usage-fetcher: OAuth creds from ~/.gemini/oauth_creds.json, auto-refresh expired tokens (including server-side 401 retry), loadCodeAssist project resolution, retrieveUserQuota with pro/flash bucket mapping - Implement opencode-go-usage-fetcher: cookie-based POST to opencode.ai/_server, two-step workspace+subscription fetch, regex parse of text/javascript response - Wire both fetchers into RateLimitService via Promise.allSettled for isolation - Add formatWindowLabel() replacing hardcoded '5h'/'wk' strings in StatusBar - Render Gemini and OpenCode Go segments in StatusBar and tooltip - 28 new tests across fetchers, service, and label formatter Closes #1022 * fix(persistence): additive merge for statusBarItems on load New providers (gemini, opencode-go) were invisible to users with existing settings because the persisted statusBarItems array (without the new entries) overwrote the defaults. Union saved items with current defaults so new providers appear automatically after upgrade without user intervention. * feat(rate-limits): add Gemini multi-bucket usage Restore provider auth wiring, preserve Gemini model buckets in detailed views, and keep compact status rendering summary-only. Avoid embedding OAuth client secrets by failing closed for auth.json refreshes. * fix(gemini-oauth): resolve symlinks and use known-path + bundle-dir extraction Replace the recursive directory walker with two targeted strategies: 1. Known paths: checks explicit Homebrew/Nix/npm layouts without walking. 2. Bundle dir: walks up from the binary to find package.json, then scans the bundle dir for hash-named oauth2 chunks. Also threads the refresh token as a plain string through tryRefreshTokenFromBundle so both the oauth_creds.json and auth.json paths share one refresh flow without coupling to either struct. Windows: uses 'where gemini' and splits on newlines for multi-result output. * fix(status-bar): fix Gemini bucket display — names, window size, model context Three issues fixed: 1. windowMinutes was computed as time-remaining-until-reset instead of the fixed window size. Gemini buckets are always 1-hour windows; use the constant 60 so labels read "93% Pro 1h" instead of "93% 47m". 2. Unknown bucket names now humanize gracefully. Unknown model IDs get the "gemini-" prefix stripped and title-cased ("gemini-3.0-ultra" → "3.0 Ultra") instead of showing "Unknown (gemini-3.0-ultra)". Added more known model mappings (2.0 Flash, 2.0 Flash Lite, 1.5 Pro, 1.5 Flash, Flash Lite). 3. Status bar segment now shows the most-constrained bucket name next to the percentage so users know which model is the binding constraint. Pure unit tests for getBucketName/deriveSessionSummary extracted to gemini-bucket-helpers.test.ts to keep gemini-usage-fetcher.test.ts under the 300-line lint limit. * feat(status-bar): show all Gemini buckets individually in status bar Instead of showing only the most-constrained bucket summary, render each bucket by name with its remaining percentage (e.g. "Flash 93% · 3.1 Pro Preview 7%"). Falls back to the session/weekly window display for providers without buckets. * fix(status-bar): show only Flash and 3.1 Pro Preview buckets in Gemini segment Filter to the two most relevant buckets (Flash + 3.1 Pro Preview) to avoid cluttering the bar. Falls back to session summary if neither bucket is present. * fix(icons): replace GeminiIcon with official 2025 multicolor gradient The previous icon used a simple linear gradient (blue→purple→red). The official Google Gemini 2025 icon uses 11 blurred ellipses (feGaussianBlur) stacked under an alpha mask to produce the characteristic multicolor glow effect. Each instance gets unique filter/mask IDs via a module-level counter to prevent ID collisions when the icon is rendered multiple times on the same page. * feat(opencode): add monthly limits, workspace override, and improved cookie handling * fix(opencode): show correct status when session cookie is missing * fix(rate-limits): discard stale data on unavailable and show errors in tooltip * fix(rate-limits): invalidate stale data when opencode config changes * fix(security): resolve vulnerabilities and harden rate-limit fetchers * fix: address critical security, performance and concurrency issues identified during code review * fix(build): resolve claude rate-limit export and broken type definitions * fix(build): remove unused code in claude-fetcher after rebase * Fix claude fetcher * fix: resolve synchronization, performance, and robustness issues in AI providers - Convert synchronous file I/O to asynchronous to prevent main process blocking - Implement concurrency limiting in recursive directory copying to avoid EMFILE errors - Persist refreshed Gemini OAuth tokens to disk and improve project ID resolution - Fix OpenCode Go fetcher to correctly handle workspace overrides and robustly resolve IDs - Refine scraping regexes to handle nested objects and improve resilience - Update status bar bucket names and fix Gemini model mapping typos - Ensure proper handling of async operations in background services * feat(gemini): deduplicate quota buckets and update model mappings * fix(rate-limits): preserve stale data for Gemini and OpenCode Go on fetch errors * fix: remove out-of-scope Linear changes and restore files to upstream/main * fix(opencode-go): handle React Flight wire format and duplicate keys The opencode.ai page uses React Server Components. Usage keys like monthlyUsage appear twice: once as `key:$R[N]={...}` with real data and once as `key:null` inside a billing component. Render order varies, so on refresh the null could appear first, causing monthly to vanish or show 100% from a sibling sub-object. Replace flat regex extraction with extractUsageBlock, which: - Iterates all occurrences of each key - Skips null assignments (no { in 30-char window after colon) - Handles the $R[N]= token between colon and opening brace - Validates usagePercent + resetInSec as direct fields before accepting Add regression tests using the real React Flight HTML format. * fix(rate-limits): address PR review feedback from nwparker Scope reverts (out-of-scope changes removed): - Revert codex-accounts async refactor (fs-utils.ts, service.ts) - Revert filesystem-auth.ts path-auth reordering - Revert filesystem-mutations.ts bundled changes - Revert repos.ts handler move; keep only -- separator fix - Revert relay/fs-handler-git-fallback.ts RegExp try/catch - Revert fs-handler.ts stat→lstat, useFileDeletion.ts !isRemote guard - Revert persistence.ts statusBarItems additive merge Feature fixes: - Validate opencodeWorkspaceId override with ^(wrk|wk)_[A-Za-z0-9]+$ before URL interpolation - Remove internal 5-min cache from fetchGeminiRateLimits (service polls already) - Make saveGeminiCredentials atomic using tmp-file + rename pattern - Add fetchGeneration counter to OpenCode to discard stale mid-flight results - Add console.warn on safeStorage decrypt failure in persistence.ts - Add geminiCliOAuthEnabled opt-in setting (default: false) with UI toggle and risk disclosure - Use per-candidate AbortControllers in opencode-go fetcher - Add fragility comments to brace-depth parsers in opencode-go fetcher Refactor (keep under 300-line lint limit): - Extract gemini-bucket-formatting.ts from gemini-usage-fetcher.ts - Extract opencode-go-page-scraper.ts from opencode-go-usage-fetcher.ts Co-authored-by: Orca * fix(rate-limits): restore OpenCode Go workspace ID regex + cleanup - opencode-go-usage-fetcher: regex accidentally shipped as `\\s*` (matches literal backslash-s) instead of `\s*` during the review-feedback commit, breaking workspace ID extraction when no override is configured — parseWorkspaceIds always returned []. 10 tests were failing as a result. - gemini-usage-fetcher: drop unused `_force` parameter left over from the removed 5-min internal cache; update service.ts call site. - codex-accounts test fixtures: add missing `geminiCliOAuthEnabled` to createSettings() to fix pre-existing typecheck errors introduced when the new setting was added. Co-authored-by: Orca * test(rate-limits): stabilize Gemini fetcher tests, silence max-lines - Rewire Gemini fetcher tests to mock the CLI-credential extractor at the module boundary instead of stubbing every fs/child_process call. The extractor is a self-contained dependency with a simple async contract, and was previously being reached through sync-fs mocks that stopped matching after the extractor was refactored to node:fs/promises — three tests were silently failing as a result. With this change, all 17 Gemini fetcher tests pass. - Replace the "proceeds with empty projectId" test (which asserted that an empty projectId still hits the quota API — the current fetcher correctly short-circuits to an actionable error instead) with a test that documents the new behavior, plus an explicit test for the geminiCliOAuthEnabled=false unavailable path. - Add a max-lines disable pragma to service.test.ts matching the one already on service.ts, so the rate-limit fetch-ordering contract and its tests remain reviewable as a single unit. Co-authored-by: Orca * fix(codex-accounts): revert runtime-home-service async refactor The review explicitly asked to drop the codex async refactor from this PR (fires floating promises from the constructor via void, and the prepare* helpers no longer block on sync completion — that is a behavior change around account prep racing with launch that deserves its own design discussion, not a drive-by). Earlier reverts removed fs-utils.ts and service.ts but left runtime-home-service.ts untouched, which was causing the matching runtime-home-service.test.ts suite (11 tests) to fail in CI. This restores the file to main so the suite is green and the PR stays scoped to the Gemini / OpenCode Go feature. Co-authored-by: Orca --------- Co-authored-by: Orca Co-authored-by: Neil <4138956+nwparker@users.noreply.github.com> --- .../runtime-home-service.test.ts | 3 + src/main/codex-accounts/service.test.ts | 3 + src/main/index.ts | 1 + src/main/ipc/repos.ts | 4 +- src/main/linear/client.ts | 1 + src/main/persistence.ts | 64 +++- src/main/rate-limits/claude-pty.ts | 14 +- .../gemini-bucket-formatting.test.ts | 65 ++++ .../rate-limits/gemini-bucket-formatting.ts | 84 +++++ .../rate-limits/gemini-cli-oauth-extractor.ts | 262 +++++++++++++ src/main/rate-limits/gemini-oauth-sources.ts | 177 +++++++++ .../gemini-usage-fetcher.fallback.test.ts | 193 ++++++++++ .../gemini-usage-fetcher.test-fixtures.ts | 56 +++ .../rate-limits/gemini-usage-fetcher.test.ts | 219 +++++++++++ src/main/rate-limits/gemini-usage-fetcher.ts | 245 ++++++++++++ .../rate-limits/opencode-go-page-scraper.ts | 172 +++++++++ .../opencode-go-usage-fetcher.test.ts | 350 ++++++++++++++++++ .../rate-limits/opencode-go-usage-fetcher.ts | 256 +++++++++++++ src/main/rate-limits/service.test.ts | 199 +++++++++- src/main/rate-limits/service.ts | 153 ++++++-- .../src/components/settings/GeneralPane.tsx | 127 +++++++ .../src/components/settings/general-search.ts | 24 ++ .../src/components/status-bar/StatusBar.tsx | 91 ++++- .../src/components/status-bar/icons.tsx | 283 ++++++++++++++ .../src/components/status-bar/tooltip.test.ts | 135 +++++++ .../src/components/status-bar/tooltip.tsx | 93 ++++- .../src/lib/window-label-formatter.test.ts | 40 ++ .../src/lib/window-label-formatter.ts | 30 ++ src/renderer/src/store/slices/rate-limits.ts | 2 +- src/shared/constants.ts | 5 + src/shared/rate-limit-types.ts | 12 +- src/shared/types.ts | 18 +- 32 files changed, 3314 insertions(+), 67 deletions(-) create mode 100644 src/main/rate-limits/gemini-bucket-formatting.test.ts create mode 100644 src/main/rate-limits/gemini-bucket-formatting.ts create mode 100644 src/main/rate-limits/gemini-cli-oauth-extractor.ts create mode 100644 src/main/rate-limits/gemini-oauth-sources.ts create mode 100644 src/main/rate-limits/gemini-usage-fetcher.fallback.test.ts create mode 100644 src/main/rate-limits/gemini-usage-fetcher.test-fixtures.ts create mode 100644 src/main/rate-limits/gemini-usage-fetcher.test.ts create mode 100644 src/main/rate-limits/gemini-usage-fetcher.ts create mode 100644 src/main/rate-limits/opencode-go-page-scraper.ts create mode 100644 src/main/rate-limits/opencode-go-usage-fetcher.test.ts create mode 100644 src/main/rate-limits/opencode-go-usage-fetcher.ts create mode 100644 src/renderer/src/components/status-bar/tooltip.test.ts create mode 100644 src/renderer/src/lib/window-label-formatter.test.ts create mode 100644 src/renderer/src/lib/window-label-formatter.ts diff --git a/src/main/codex-accounts/runtime-home-service.test.ts b/src/main/codex-accounts/runtime-home-service.test.ts index 5555b64c2d6..d5ce0b49525 100644 --- a/src/main/codex-accounts/runtime-home-service.test.ts +++ b/src/main/codex-accounts/runtime-home-service.test.ts @@ -86,6 +86,9 @@ function createSettings(overrides: Partial = {}): GlobalSettings defaultTaskSource: 'github', defaultRepoSelection: null, defaultLinearTeamSelection: null, + opencodeSessionCookie: '', + opencodeWorkspaceId: '', + geminiCliOAuthEnabled: false, agentCmdOverrides: {}, terminalMacOptionAsAlt: 'false', terminalMacOptionAsAltMigrated: true, diff --git a/src/main/codex-accounts/service.test.ts b/src/main/codex-accounts/service.test.ts index ce193feb02e..8b9c340d532 100644 --- a/src/main/codex-accounts/service.test.ts +++ b/src/main/codex-accounts/service.test.ts @@ -80,6 +80,9 @@ function createSettings(overrides: Partial = {}): GlobalSettings defaultTaskSource: 'github', defaultRepoSelection: null, defaultLinearTeamSelection: null, + opencodeSessionCookie: '', + opencodeWorkspaceId: '', + geminiCliOAuthEnabled: false, agentCmdOverrides: {}, terminalMacOptionAsAlt: 'false', terminalMacOptionAsAltMigrated: true, diff --git a/src/main/index.ts b/src/main/index.ts index bf9092e5568..ab60f0fac25 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -336,6 +336,7 @@ app.whenReady().then(async () => { claudeAccounts = new ClaudeAccountService(store, rateLimits, claudeRuntimeAuth) rateLimits.setCodexHomePathResolver(() => codexRuntimeHome!.prepareForRateLimitFetch()) rateLimits.setClaudeAuthPreparationResolver(() => claudeRuntimeAuth!.prepareForRateLimitFetch()) + rateLimits.setSettingsResolver(() => store!.getSettings()) runtime = new OrcaRuntimeService(store, stats) starNag = new StarNagService(store, stats) starNag.start() diff --git a/src/main/ipc/repos.ts b/src/main/ipc/repos.ts index 66fe074ca2a..b59aa8fb7ca 100644 --- a/src/main/ipc/repos.ts +++ b/src/main/ipc/repos.ts @@ -282,7 +282,9 @@ export function registerRepoHandlers(mainWindow: BrowserWindow, store: Store): v // Why: clone destination may be a WSL path (e.g. user picks a WSL // directory). Use the parent destination as the cwd so the runner // detects WSL and routes through wsl.exe. - const proc = gitSpawn(['clone', '--progress', args.url, clonePath], { + // Why: use the '--' separator to isolate the URL argument and prevent + // malicious URLs from being interpreted as git flags (command injection). + const proc = gitSpawn(['clone', '--progress', '--', args.url, clonePath], { cwd: args.destination, stdio: ['ignore', 'ignore', 'pipe'] }) diff --git a/src/main/linear/client.ts b/src/main/linear/client.ts index 5f0da2fc443..04854e3423f 100644 --- a/src/main/linear/client.ts +++ b/src/main/linear/client.ts @@ -232,6 +232,7 @@ export async function testConnection(): Promise< if (!token) { return { ok: false, error: 'No API key stored.' } } + try { const client = new LinearClient({ apiKey: token }) const me = await client.viewer diff --git a/src/main/persistence.ts b/src/main/persistence.ts index 02d54a10b9d..e2b8fd1cb61 100644 --- a/src/main/persistence.ts +++ b/src/main/persistence.ts @@ -1,7 +1,7 @@ /* eslint-disable max-lines -- Why: persistence keeps schema defaults, migration, load/save, and flush logic in one file so the full storage contract is reviewable as a unit instead of being scattered across modules. */ -import { app } from 'electron' +import { app, safeStorage } from 'electron' import { readFileSync, writeFileSync, mkdirSync, existsSync, renameSync, unlinkSync } from 'fs' import { writeFile, rename, mkdir, rm } from 'fs/promises' import { join, dirname } from 'path' @@ -19,6 +19,35 @@ import { } from '../shared/constants' import { parseWorkspaceSession } from '../shared/workspace-session-schema' +function encrypt(plaintext: string): string { + if (!plaintext || !safeStorage.isEncryptionAvailable()) { + return plaintext + } + try { + return safeStorage.encryptString(plaintext).toString('base64') + } catch (err) { + console.error('[persistence] Encryption failed:', err) + return plaintext + } +} + +function decrypt(ciphertext: string): string { + if (!ciphertext || !safeStorage.isEncryptionAvailable()) { + return ciphertext + } + try { + return safeStorage.decryptString(Buffer.from(ciphertext, 'base64')) + } catch { + // Why: if decryption fails, it likely means the value was stored as + // plaintext (pre-encryption build) or the OS keychain changed. Fall + // back to the raw string so users don't lose their cookie after upgrade. + console.warn( + '[persistence] safeStorage decryption failed — returning ciphertext as-is. Possible keychain reset.' + ) + return ciphertext + } +} + // Why: the data-file path must not be a module-level constant. Module-level // code runs at import time — before configureDevUserDataPath() redirects the // userData path in index.ts — so a constant would capture the default (non-dev) @@ -76,6 +105,13 @@ export class Store { if (existsSync(dataFile)) { const raw = readFileSync(dataFile, 'utf-8') const parsed = JSON.parse(raw) as PersistedState + + // Why: opencodeSessionCookie is stored encrypted on disk via safeStorage. + // Decrypt at the load boundary so the rest of the app sees plaintext. + if (parsed.settings?.opencodeSessionCookie) { + parsed.settings.opencodeSessionCookie = decrypt(parsed.settings.opencodeSessionCookie) + } + // Merge with defaults in case new fields were added const defaults = getDefaultPersistedState(homedir()) // Why: before the layout-aware 'auto' mode shipped (issue #903), @@ -185,12 +221,23 @@ export class Store { const dir = dirname(dataFile) await mkdir(dir, { recursive: true }).catch(() => {}) const tmpFile = `${dataFile}.${process.pid}.${Date.now()}.${Math.random().toString(16).slice(2)}.tmp` + + // Why: opencodeSessionCookie must be encrypted on disk. Clone state so + // the in-memory this.state stays plaintext for the rest of the app. + const stateToSave = { + ...this.state, + settings: { + ...this.state.settings, + opencodeSessionCookie: encrypt(this.state.settings.opencodeSessionCookie) + } + } + // Why: wrap write+rename in try/finally-on-error so any failure (ENOSPC, // ENFILE, EIO, permission) removes the tmp file rather than leaving a // multi-megabyte orphan behind. Successful rename consumes the tmp file. let renamed = false try { - await writeFile(tmpFile, JSON.stringify(this.state, null, 2), 'utf-8') + await writeFile(tmpFile, JSON.stringify(stateToSave, null, 2), 'utf-8') // Why: if flush() ran while this async write was in-flight, it bumped // writeGeneration and already wrote the latest state synchronously. // Renaming this stale tmp file would overwrite the fresh data. @@ -215,12 +262,23 @@ export class Store { mkdirSync(dir, { recursive: true }) } const tmpFile = `${dataFile}.${process.pid}.${Date.now()}.${Math.random().toString(16).slice(2)}.tmp` + + // Why: opencodeSessionCookie must be encrypted on disk. Clone state so + // the in-memory this.state stays plaintext for the rest of the app. + const stateToSave = { + ...this.state, + settings: { + ...this.state.settings, + opencodeSessionCookie: encrypt(this.state.settings.opencodeSessionCookie) + } + } + // Why: mirror the async path — on any failure between writeFileSync and // renameSync, remove the tmp file so crashes during shutdown don't leak // orphans into userData. let renamed = false try { - writeFileSync(tmpFile, JSON.stringify(this.state, null, 2), 'utf-8') + writeFileSync(tmpFile, JSON.stringify(stateToSave, null, 2), 'utf-8') renameSync(tmpFile, dataFile) renamed = true } finally { diff --git a/src/main/rate-limits/claude-pty.ts b/src/main/rate-limits/claude-pty.ts index 1ae686cebf9..53a4b4e5504 100644 --- a/src/main/rate-limits/claude-pty.ts +++ b/src/main/rate-limits/claude-pty.ts @@ -4,6 +4,7 @@ import type { ClaudeRuntimeAuthPreparation } from '../claude-accounts/runtime-au import { applyClaudeEnvPatch } from '../claude-accounts/environment' const PTY_TIMEOUT_MS = 25_000 +const MAX_OUTPUT_LENGTH = 100_000 // 100KB buffer limit // --------------------------------------------------------------------------- // PTY fallback — spawn interactive `claude`, send `/usage`, parse the TUI @@ -141,12 +142,11 @@ export async function fetchViaPty(options?: { const claudeCommand = resolveClaudeCommand() // Why: node-pty cannot spawn .cmd/.bat batch scripts directly on Windows — - // those need cmd.exe as an interpreter. resolveClaudeCommand() may also fall - // back to bare 'claude' when it can't locate the binary on disk, yet cmd.exe - // can still find claude.cmd via PATHEXT. Always route through cmd.exe on win32. + // those need cmd.exe as an interpreter. Always route through cmd.exe on win32 + // and ensure the command path is properly quoted if it contains spaces. const isWin32 = process.platform === 'win32' const spawnFile = isWin32 ? 'cmd.exe' : claudeCommand - const spawnArgs = isWin32 ? ['/c', claudeCommand] : [] + const spawnArgs = isWin32 ? ['/c', `"${claudeCommand}"`] : [] const spawnEnv = applyClaudeEnvPatch( { ...process.env, TERM: 'xterm-256color' } as Record, @@ -210,7 +210,7 @@ export async function fetchViaPty(options?: { return } enterInterval = setInterval(() => { - if (!resolved) { + if (!resolved && !stopDetected) { term.write('\r') } }, 800) @@ -266,6 +266,10 @@ export async function fetchViaPty(options?: { const onDataDisposable = term.onData((data) => { output += data + // Why: prevent memory exhaustion if the CLI process floods output + if (output.length > MAX_OUTPUT_LENGTH) { + output = output.slice(-MAX_OUTPUT_LENGTH) + } // eslint-disable-next-line no-control-regex const cleanChunk = data.replace(/\x1b\[[0-9;]*[a-zA-Z]/g, '') diff --git a/src/main/rate-limits/gemini-bucket-formatting.test.ts b/src/main/rate-limits/gemini-bucket-formatting.test.ts new file mode 100644 index 00000000000..b806fbef902 --- /dev/null +++ b/src/main/rate-limits/gemini-bucket-formatting.test.ts @@ -0,0 +1,65 @@ +import { describe, expect, it } from 'vitest' +import type { RateLimitBucket } from '../../shared/rate-limit-types' +import { getBucketName, deriveSessionSummary } from './gemini-bucket-formatting' + +describe('getBucketName', () => { + it('maps known model IDs to stable names', () => { + expect(getBucketName('gemini-2.5-pro')).toBe('Pro') + expect(getBucketName('gemini-2.5-flash')).toBe('Flash') + expect(getBucketName('gemini-2.5-flash-lite')).toBe('Flash Lite') + expect(getBucketName('gemini-2.0-flash-lite')).toBe('2.0 Flash Lite') + expect(getBucketName('gemini-2.0-flash')).toBe('2.0 Flash') + }) + + it('humanizes unknown model IDs by stripping the gemini- prefix', () => { + expect(getBucketName('gemini-3.0-ultra')).toBe('3.0 Ultra') + expect(getBucketName('gemini-experimental')).toBe('Exp') + expect(getBucketName('some-random-id')).toBe('Some Random Id') + }) +}) + +describe('deriveSessionSummary', () => { + it('returns null for empty buckets', () => { + expect(deriveSessionSummary([])).toBeNull() + }) + + it('picks the most constrained bucket (highest usedPercent) as session summary', () => { + const buckets: RateLimitBucket[] = [ + { name: 'Pro', usedPercent: 30, windowMinutes: 60, resetsAt: null, resetDescription: null }, + { name: 'Flash', usedPercent: 80, windowMinutes: 60, resetsAt: null, resetDescription: null }, + { + name: 'Flash Lite', + usedPercent: 10, + windowMinutes: 60, + resetsAt: null, + resetDescription: null + } + ] + const summary = deriveSessionSummary(buckets) + expect(summary).not.toBeNull() + expect(summary!.usedPercent).toBe(80) + expect(summary!.windowMinutes).toBe(60) + }) + + it('preserves reset metadata from the most constrained bucket', () => { + const buckets: RateLimitBucket[] = [ + { + name: 'Pro', + usedPercent: 30, + windowMinutes: 60, + resetsAt: 1000, + resetDescription: '2:00 PM' + }, + { + name: 'Flash', + usedPercent: 80, + windowMinutes: 60, + resetsAt: 2000, + resetDescription: '3:00 PM' + } + ] + const summary = deriveSessionSummary(buckets) + expect(summary!.resetsAt).toBe(2000) + expect(summary!.resetDescription).toBe('3:00 PM') + }) +}) diff --git a/src/main/rate-limits/gemini-bucket-formatting.ts b/src/main/rate-limits/gemini-bucket-formatting.ts new file mode 100644 index 00000000000..86cf9cea763 --- /dev/null +++ b/src/main/rate-limits/gemini-bucket-formatting.ts @@ -0,0 +1,84 @@ +import type { RateLimitBucket, RateLimitWindow } from '../../shared/rate-limit-types' + +const MODEL_ID_TO_BUCKET_NAME: Record = { + 'gemini-3.1-pro': '3.1 Pro', + 'gemini-3.1-flash': '3.1 Flash', + 'gemini-3.1-flash-lite': '3.1 Flash Lite', + 'gemini-3.0-pro': '3.0 Pro', + 'gemini-3.0-flash': '3.0 Flash', + 'gemini-2.5-pro': 'Pro', + 'gemini-2.5-flash': 'Flash', + 'gemini-2.5-flash-lite': 'Flash Lite', + 'gemini-2.0-pro': '2.0 Pro', + 'gemini-2.0-flash': '2.0 Flash', + 'gemini-2.0-flash-lite': '2.0 Flash Lite', + 'gemini-1.5-pro': '1.5 Pro', + 'gemini-1.5-flash': '1.5 Flash', + 'gemini-exp': 'Exp', + 'gemini-experimental': 'Exp' +} + +function humanizeModelId(modelId: string): string { + const withoutPrefix = modelId.replace(/^gemini-/i, '') + return withoutPrefix + .split('-') + .map((part) => (part.length > 0 ? part[0]!.toUpperCase() + part.slice(1) : part)) + .join(' ') +} + +export function getBucketName(modelId: string): string { + return MODEL_ID_TO_BUCKET_NAME[modelId] ?? humanizeModelId(modelId) +} + +export function buildRateLimitBucket(b: { + remainingFraction: number + resetTime: string + modelId: string +}): RateLimitBucket { + const usedPercent = Math.min(100, Math.max(0, Math.round((1 - b.remainingFraction) * 100))) + const resetsAtTime = new Date(b.resetTime).getTime() + return { + name: getBucketName(b.modelId), + usedPercent, + windowMinutes: 60, + resetsAt: !isNaN(resetsAtTime) ? resetsAtTime : null, + resetDescription: null + } +} + +export function deduplicateBuckets( + buckets: (RateLimitBucket & { modelId: string })[] +): RateLimitBucket[] { + const result: (RateLimitBucket & { modelId: string })[] = [] + const seenKeys = new Map() + for (const b of buckets) { + const key = `${b.usedPercent}-${b.resetsAt}` + const existingIndex = seenKeys.get(key) + if (existingIndex === undefined) { + seenKeys.set(key, result.length) + result.push(b) + continue + } + const existing = result[existingIndex]! + const existingInMap = existing.modelId in MODEL_ID_TO_BUCKET_NAME + const currentInMap = b.modelId in MODEL_ID_TO_BUCKET_NAME + if ( + (currentInMap && !existingInMap) || + (currentInMap === existingInMap && b.name.length < existing.name.length) + ) { + result[existingIndex] = b + } + } + return result.map(({ modelId: _id, ...rest }) => rest) +} + +export function deriveSessionSummary(buckets: RateLimitBucket[]): RateLimitWindow | null { + if (buckets.length === 0) { + return null + } + const mostConstrained = buckets.reduce((worst, bucket) => { + return bucket.usedPercent > worst.usedPercent ? bucket : worst + }) + const { name: _name, ...window } = mostConstrained + return window +} diff --git a/src/main/rate-limits/gemini-cli-oauth-extractor.ts b/src/main/rate-limits/gemini-cli-oauth-extractor.ts new file mode 100644 index 00000000000..80e841e7adf --- /dev/null +++ b/src/main/rate-limits/gemini-cli-oauth-extractor.ts @@ -0,0 +1,262 @@ +import { exec } from 'node:child_process' +import { access, readdir, readFile, realpath } from 'node:fs/promises' +import { promisify } from 'node:util' +import { homedir } from 'node:os' +import path from 'node:path' + +const execAsync = promisify(exec) + +async function fileExists(filePath: string): Promise { + try { + await access(filePath) + return true + } catch { + return false + } +} + +// The oauth2.js relative path inside a @google/gemini-cli-core package. +const OAUTH2_SUBPATH = path.join('dist', 'src', 'code_assist', 'oauth2.js') + +async function resolveGeminiBinary(): Promise { + const whichCmd = process.platform === 'win32' ? 'where gemini' : 'which gemini' + try { + const { stdout } = await execAsync(whichCmd, { encoding: 'utf-8' }) + const fromPath = stdout.trim().split(/\r?\n/)[0] + if (fromPath && (await fileExists(fromPath))) { + return fromPath + } + } catch { + // ignore which/where failure + } + + // Why: on macOS/Linux GUI apps, the PATH might not include the binary. + // Checking common installation prefixes as fallbacks. + if (process.platform !== 'win32') { + const fallbacks = [ + '/usr/local/bin/gemini', + '/opt/homebrew/bin/gemini', + path.join(homedir(), '.local', 'bin', 'gemini'), + path.join(homedir(), 'bin', 'gemini') + ] + for (const candidate of fallbacks) { + if (await fileExists(candidate)) { + return candidate + } + } + } + + return null +} + +// Why: on all platforms the gemini binary may be a symlink (e.g. Homebrew's bin/ +// symlinks into Cellar). We must resolve it before deriving sibling paths — otherwise +// dirname points to the symlink directory, not the real installation root. +async function resolveSymlink(filePath: string): Promise { + try { + return await realpath(filePath) + } catch { + return filePath + } +} + +function parseOAuthCredentials(content: string): { clientId: string; clientSecret: string } | null { + const idMatch = content.match(/OAUTH_CLIENT_ID\s*=\s*['"]([^'"]+)['"]/)?.[1] + const secretMatch = content.match(/OAUTH_CLIENT_SECRET\s*=\s*['"]([^'"]+)['"]/)?.[1] + if (idMatch && secretMatch) { + return { clientId: idMatch, clientSecret: secretMatch } + } + return null +} + +async function tryReadCredentials( + filePath: string +): Promise<{ clientId: string; clientSecret: string } | null> { + try { + const content = await readFile(filePath, 'utf-8') + return parseOAuthCredentials(content) + } catch { + return null + } +} + +// Why: these are the known stable layouts for every major Gemini CLI install method. +// Checking explicit paths is fast and avoids walking the entire directory tree. +async function extractFromKnownPaths( + realGeminiPath: string +): Promise<{ clientId: string; clientSecret: string } | null> { + const binDir = path.dirname(realGeminiPath) + const baseDir = path.dirname(binDir) + + const candidates = [ + // Homebrew: bin -> Cellar//bin, real files live under libexec/lib + path.join( + baseDir, + 'libexec', + 'lib', + 'node_modules', + '@google', + 'gemini-cli', + 'node_modules', + '@google', + 'gemini-cli-core', + OAUTH2_SUBPATH + ), + // Homebrew alternate (some versions skip the extra nesting) + path.join( + baseDir, + 'lib', + 'node_modules', + '@google', + 'gemini-cli', + 'node_modules', + '@google', + 'gemini-cli-core', + OAUTH2_SUBPATH + ), + // Nix package layout + path.join( + baseDir, + 'share', + 'gemini-cli', + 'node_modules', + '@google', + 'gemini-cli-core', + OAUTH2_SUBPATH + ), + // npm/bun global install: gemini-cli-core is a sibling of gemini-cli + path.join(baseDir, '..', 'gemini-cli-core', OAUTH2_SUBPATH), + // npm nested inside gemini-cli + path.join(baseDir, 'node_modules', '@google', 'gemini-cli-core', OAUTH2_SUBPATH) + ] + + for (const candidate of candidates) { + const creds = await tryReadCredentials(path.normalize(candidate)) + if (creds) { + return creds + } + } + + return null +} + +// Why: newer Gemini CLI versions (>=0.38) ship everything bundled into hash-named +// chunks with no oauth2.js source file. Scanning the bundle dir for the credential +// constants is the only reliable fallback for those installs. +async function extractFromBundleDir( + geminiCliPackageRoot: string +): Promise<{ clientId: string; clientSecret: string } | null> { + const bundleDir = path.join(geminiCliPackageRoot, 'bundle') + if (!(await fileExists(bundleDir))) { + return null + } + + let entries: string[] + try { + entries = (await readdir(bundleDir)).filter((f) => f.endsWith('.js')) + } catch { + return null + } + + for (const entry of entries) { + const creds = await tryReadCredentials(path.join(bundleDir, entry)) + if (creds) { + return creds + } + } + + return null +} + +// Resolves the gemini-cli package root directory by walking up the directory +// tree from the real binary path, looking for package.json with the right name, +// or the global Node layout under lib/node_modules. +async function findGeminiPackageRoot(realGeminiPath: string): Promise { + const MAX_ASCENTS = 8 + let current = path.dirname(realGeminiPath) + + for (let i = 0; i <= MAX_ASCENTS; i++) { + const pkgJson = path.join(current, 'package.json') + if (await fileExists(pkgJson)) { + try { + const raw = await readFile(pkgJson, 'utf-8') + const pkg = JSON.parse(raw) as { name?: string } + if (pkg.name === '@google/gemini-cli') { + return current + } + } catch { + // malformed package.json — keep walking + } + } + + // Global Node layout: /lib/node_modules/@google/gemini-cli + const globalPkg = path.join( + current, + 'lib', + 'node_modules', + '@google', + 'gemini-cli', + 'package.json' + ) + if (await fileExists(globalPkg)) { + return path.join(current, 'lib', 'node_modules', '@google', 'gemini-cli') + } + + // Windows global install layout: /node_modules/@google/gemini-cli + const windowsGlobalPkg = path.join( + current, + 'node_modules', + '@google', + 'gemini-cli', + 'package.json' + ) + if (await fileExists(windowsGlobalPkg)) { + return path.join(current, 'node_modules', '@google', 'gemini-cli') + } + + const parent = path.dirname(current) + if (parent === current) { + break + } + current = parent + } + + return null +} + +export async function extractOAuthClientCredentials(): Promise<{ + clientId: string + clientSecret: string +} | null> { + const geminiPath = await resolveGeminiBinary() + if (!geminiPath) { + return null + } + + const realPath = await resolveSymlink(geminiPath) + + // 1. Known static paths (fast, covers most installs with source layout) + const fromKnown = await extractFromKnownPaths(realPath) + if (fromKnown) { + return fromKnown + } + + // 2. Walk up to find the package root, then try source layout + bundle dir + const packageRoot = await findGeminiPackageRoot(realPath) + if (packageRoot) { + const fromSource = + (await tryReadCredentials( + path.join(packageRoot, 'node_modules', '@google', 'gemini-cli-core', OAUTH2_SUBPATH) + )) ?? (await tryReadCredentials(path.join(packageRoot, OAUTH2_SUBPATH))) + if (fromSource) { + return fromSource + } + + const fromBundle = await extractFromBundleDir(packageRoot) + if (fromBundle) { + return fromBundle + } + } + + return null +} diff --git a/src/main/rate-limits/gemini-oauth-sources.ts b/src/main/rate-limits/gemini-oauth-sources.ts new file mode 100644 index 00000000000..fec013ce745 --- /dev/null +++ b/src/main/rate-limits/gemini-oauth-sources.ts @@ -0,0 +1,177 @@ +import { readFile, writeFile, rename } from 'node:fs/promises' +import { homedir } from 'node:os' +import path from 'node:path' +import { net } from 'electron' +import { extractOAuthClientCredentials } from './gemini-cli-oauth-extractor' + +const API_TIMEOUT_MS = 10_000 +const OAUTH_CREDS_PATH = path.join(homedir(), '.gemini', 'oauth_creds.json') +const GOOGLE_TOKEN_URL = 'https://oauth2.googleapis.com/token' +const LOAD_CODE_ASSIST_URL = 'https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist' + +export type GeminiCredentials = { + access_token: string + refresh_token: string + expiry_date: number +} + +export type GoogleAuthEntry = { + type: 'oauth' + access: string + expires: number + refresh: string +} + +type AuthJson = { + google?: GoogleAuthEntry + 'opencode-go'?: { type: 'api'; key: string } +} + +export async function readAuthJson(): Promise { + const candidates = [ + process.env.APPDATA ? path.join(process.env.APPDATA, 'opencode', 'auth.json') : null, + process.env.XDG_DATA_HOME + ? path.join(process.env.XDG_DATA_HOME, 'opencode', 'auth.json') + : null, + path.join(homedir(), '.local', 'share', 'opencode', 'auth.json'), + path.join(homedir(), 'Library', 'Application Support', 'opencode', 'auth.json') + ].filter((candidate): candidate is string => candidate !== null) + + for (const candidate of candidates) { + try { + const raw = await readFile(candidate, 'utf-8') + return JSON.parse(raw) as AuthJson + } catch (err) { + if (err && typeof err === 'object' && 'code' in err && err.code === 'ENOENT') { + continue + } + throw err + } + } + + return null +} + +export async function readGeminiCredentials(): Promise { + try { + const raw = await readFile(OAUTH_CREDS_PATH, 'utf-8') + const parsed = JSON.parse(raw) as unknown + if ( + parsed && + typeof parsed === 'object' && + 'access_token' in parsed && + typeof parsed.access_token === 'string' && + 'refresh_token' in parsed && + typeof parsed.refresh_token === 'string' && + 'expiry_date' in parsed && + typeof parsed.expiry_date === 'number' + ) { + return parsed as GeminiCredentials + } + return null + } catch (err) { + if (err && typeof err === 'object' && 'code' in err && err.code === 'ENOENT') { + return null + } + throw err + } +} + +export async function saveGeminiCredentials(creds: GeminiCredentials): Promise { + const tmpPath = `${OAUTH_CREDS_PATH}.${process.pid}.tmp` + await writeFile(tmpPath, JSON.stringify(creds, null, 2), 'utf-8') + await rename(tmpPath, OAUTH_CREDS_PATH) +} + +export type RefreshTokenResult = { + accessToken: string | null + newRefreshToken: string | null + expiresIn?: number +} + +export async function refreshAccessToken( + refreshToken: string, + clientId: string, + clientSecret: string +): Promise { + const controller = new AbortController() + const timeout = setTimeout(() => controller.abort(), API_TIMEOUT_MS) + + try { + const res = await net.fetch(GOOGLE_TOKEN_URL, { + method: 'POST', + headers: { 'Content-Type': 'application/x-www-form-urlencoded' }, + body: new URLSearchParams({ + client_id: clientId, + client_secret: clientSecret, + refresh_token: refreshToken, + grant_type: 'refresh_token' + }).toString(), + signal: controller.signal + }) + + if (!res.ok) { + return { accessToken: null, newRefreshToken: null } + } + + const data = (await res.json()) as { + access_token?: string + refresh_token?: string + expires_in?: number + } + return { + accessToken: typeof data.access_token === 'string' ? data.access_token : null, + newRefreshToken: typeof data.refresh_token === 'string' ? data.refresh_token : null, + expiresIn: typeof data.expires_in === 'number' ? data.expires_in : undefined + } + } finally { + clearTimeout(timeout) + } +} + +export async function loadProjectId(accessToken: string): Promise { + const controller = new AbortController() + const timeout = setTimeout(() => controller.abort(), API_TIMEOUT_MS) + + try { + const res = await net.fetch(LOAD_CODE_ASSIST_URL, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer ${accessToken}` + }, + body: JSON.stringify({ metadata: { ideType: 'GEMINI_CLI', pluginType: 'GEMINI' } }), + signal: controller.signal + }) + + if (!res.ok) { + throw new Error(`Failed to load Gemini project ID (HTTP ${res.status})`) + } + + const data = (await res.json()) as { cloudaicompanionProject?: string } + if (typeof data.cloudaicompanionProject !== 'string') { + throw new Error('Gemini project ID not found in API response') + } + return data.cloudaicompanionProject + } finally { + clearTimeout(timeout) + } +} + +// Why: accepts a plain refresh token string so both the oauth_creds.json path +// (GeminiCredentials) and the auth.json path (pipe-split string) can share +// the same bundle credential extraction without coupling to either struct. +export async function tryRefreshTokenFromBundle( + refreshToken: string, + allowCliOAuth = true +): Promise { + if (!allowCliOAuth) { + return null + } + const clientCreds = await extractOAuthClientCredentials() + if (!clientCreds) { + return null + } + + return refreshAccessToken(refreshToken, clientCreds.clientId, clientCreds.clientSecret) +} diff --git a/src/main/rate-limits/gemini-usage-fetcher.fallback.test.ts b/src/main/rate-limits/gemini-usage-fetcher.fallback.test.ts new file mode 100644 index 00000000000..0d978a14f30 --- /dev/null +++ b/src/main/rate-limits/gemini-usage-fetcher.fallback.test.ts @@ -0,0 +1,193 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + expiredCreds, + makeResponse, + quotaResponse, + validCreds +} from './gemini-usage-fetcher.test-fixtures' + +const { readFileMock, extractCredsMock, netFetchMock } = vi.hoisted(() => ({ + readFileMock: vi.fn(), + extractCredsMock: vi.fn(), + netFetchMock: vi.fn() +})) + +// Why: mock the CLI-credential extractor at the module boundary. The extractor +// is a self-contained dependency with a simple async contract (returns a +// { clientId, clientSecret } record or null). Mocking it here keeps these +// tests focused on the oauth_creds.json refresh → loadCodeAssist → quota +// flow rather than on filesystem plumbing, and avoids having to keep pace +// with extractor internals when they change. +vi.mock('./gemini-cli-oauth-extractor', () => ({ + extractOAuthClientCredentials: extractCredsMock +})) + +vi.mock('node:fs/promises', () => ({ + readFile: readFileMock, + // Why: saveGeminiCredentials is exercised on the refresh path. The atomic + // tmp+rename write has no observable side effect in these tests, so the + // stubs just resolve. + writeFile: vi.fn().mockResolvedValue(undefined), + rename: vi.fn().mockResolvedValue(undefined) +})) + +vi.mock('electron', () => ({ + net: { fetch: netFetchMock } +})) + +import { fetchGeminiRateLimits } from './gemini-usage-fetcher' + +describe('fetchGeminiRateLimits fallback oauth creds', () => { + beforeEach(() => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-04-24T12:00:00.000Z')) + readFileMock.mockReset() + extractCredsMock.mockReset() + netFetchMock.mockReset() + extractCredsMock.mockResolvedValue({ + clientId: 'client-id-123', + clientSecret: 'client-secret-456' + }) + }) + + it('falls back to oauth_creds.json when auth.json has no google key', async () => { + readFileMock.mockImplementation(async (filePath: string) => { + if (filePath.includes('auth.json')) { + return JSON.stringify({ 'opencode-go': { type: 'api', key: 'k' } }) + } + if (filePath.includes('oauth_creds.json')) { + return JSON.stringify(validCreds) + } + throw { code: 'ENOENT' } + }) + netFetchMock + .mockResolvedValueOnce(makeResponse({ cloudaicompanionProject: 'proj-123' })) + .mockResolvedValueOnce(makeResponse(quotaResponse)) + + const result = await fetchGeminiRateLimits(true) + + expect(result.status).toBe('ok') + expect(result.error).toBeNull() + expect(result.session).not.toBeNull() + }) + + it('falls back to oauth_creds.json and resolves project via loadCodeAssist', async () => { + readFileMock.mockImplementation(async (filePath: string) => { + if (filePath.includes('auth.json')) { + return JSON.stringify({}) + } + if (filePath.includes('oauth_creds.json')) { + return JSON.stringify(validCreds) + } + throw { code: 'ENOENT' } + }) + netFetchMock + .mockResolvedValueOnce(makeResponse({ cloudaicompanionProject: 'cli-proj-456' })) + .mockResolvedValueOnce(makeResponse(quotaResponse)) + + const result = await fetchGeminiRateLimits(true) + + expect(result.status).toBe('ok') + expect(result.error).toBeNull() + expect(result.session).not.toBeNull() + + const quotaCall = netFetchMock.mock.calls.find( + (call) => typeof call[0] === 'string' && call[0].includes('retrieveUserQuota') + ) + expect(quotaCall).toBeDefined() + const quotaBody = JSON.parse((quotaCall![1] as RequestInit).body as string) + expect(quotaBody.project).toBe('cli-proj-456') + }) + + it('refreshes via bundled client credentials when expiry passed', async () => { + readFileMock.mockImplementation(async (filePath: string) => { + if (filePath.includes('auth.json')) { + return JSON.stringify({}) + } + if (filePath.includes('oauth_creds.json')) { + return JSON.stringify(expiredCreds) + } + throw { code: 'ENOENT' } + }) + netFetchMock + .mockResolvedValueOnce( + makeResponse({ access_token: 'bundle-refreshed-token', expires_in: 3600 }) + ) + .mockResolvedValueOnce(makeResponse({ cloudaicompanionProject: 'cli-proj-456' })) + .mockResolvedValueOnce(makeResponse(quotaResponse)) + + const result = await fetchGeminiRateLimits(true) + + expect(result.status).toBe('ok') + expect(result.error).toBeNull() + expect(result.session).not.toBeNull() + + const refreshCall = netFetchMock.mock.calls.find( + (call) => typeof call[0] === 'string' && call[0].includes('oauth2.googleapis.com') + ) + expect(refreshCall).toBeDefined() + const refreshBody = new URLSearchParams((refreshCall![1] as RequestInit).body as string) + expect(refreshBody.get('client_id')).toBe('client-id-123') + expect(refreshBody.get('client_secret')).toBe('client-secret-456') + }) + + it('returns error when oauth_creds.json token expired and bundle refresh fails', async () => { + readFileMock.mockImplementation(async (filePath: string) => { + if (filePath.includes('auth.json')) { + return JSON.stringify({}) + } + if (filePath.includes('oauth_creds.json')) { + return JSON.stringify(expiredCreds) + } + throw { code: 'ENOENT' } + }) + // Simulate: no Gemini CLI installed, so the extractor returns null and + // tryRefreshTokenFromBundle can't obtain client credentials to refresh. + extractCredsMock.mockResolvedValue(null) + + const result = await fetchGeminiRateLimits(true) + + expect(result.status).toBe('error') + expect(result.error).toContain('Token refresh failed') + expect(result.session).toBeNull() + expect(result.weekly).toBeNull() + }) + + it('returns error when loadCodeAssist cannot resolve a project for oauth_creds path', async () => { + // Why: when the fallback (oauth_creds.json) path has no project embedded + // and loadCodeAssist fails, we surface a clear "project ID not found" + // error rather than silently posting an empty project to the quota API — + // an empty project causes a 400 that looks like an auth failure. + readFileMock.mockImplementation(async (filePath: string) => { + if (filePath.includes('auth.json')) { + return JSON.stringify({}) + } + if (filePath.includes('oauth_creds.json')) { + return JSON.stringify(validCreds) + } + throw { code: 'ENOENT' } + }) + netFetchMock.mockResolvedValueOnce(makeResponse('Internal Server Error', 500)) + + const result = await fetchGeminiRateLimits(true) + + expect(result.status).toBe('error') + expect(result.error).toContain('Gemini project ID not found') + }) + + it('returns unavailable when geminiCliOAuthEnabled=false and no google entry in auth.json', async () => { + readFileMock.mockImplementation(async (filePath: string) => { + if (filePath.includes('auth.json')) { + return JSON.stringify({ 'opencode-go': { type: 'api', key: 'k' } }) + } + throw { code: 'ENOENT' } + }) + + const result = await fetchGeminiRateLimits(false) + + expect(result.status).toBe('unavailable') + expect(result.error).toContain('disabled') + // No network calls should have been made. + expect(netFetchMock).not.toHaveBeenCalled() + }) +}) diff --git a/src/main/rate-limits/gemini-usage-fetcher.test-fixtures.ts b/src/main/rate-limits/gemini-usage-fetcher.test-fixtures.ts new file mode 100644 index 00000000000..d37053b3458 --- /dev/null +++ b/src/main/rate-limits/gemini-usage-fetcher.test-fixtures.ts @@ -0,0 +1,56 @@ +export function makeResponse(body: unknown, status = 200): Response { + return { + ok: status >= 200 && status < 300, + status, + json: async () => body, + text: async () => (typeof body === 'string' ? body : JSON.stringify(body)) + } as Response +} + +export function makeDirent(name: string, isDir: boolean) { + return { + name, + isDirectory: () => isDir, + isFile: () => !isDir + } +} + +export const authJsonGoogle = { + google: { + type: 'oauth', + access: 'auth-json-access-token', + expires: new Date('2026-04-24T13:00:00.000Z').getTime(), + refresh: 'refresh-token-abc|proj-123|managed-456' + } +} as const + +export const authJsonGoogleExpired = { + google: { + type: 'oauth', + access: 'expired-access-token', + expires: new Date('2026-04-24T11:00:00.000Z').getTime(), + refresh: 'refresh-token-abc|proj-123|managed-456' + } +} as const + +export const validCreds = { + access_token: 'valid-token', + refresh_token: 'refresh-token', + expiry_date: new Date('2026-04-24T13:00:00.000Z').getTime() +} + +export const expiredCreds = { + access_token: 'expired-token', + refresh_token: 'refresh-token', + expiry_date: new Date('2026-04-24T11:00:00.000Z').getTime() +} + +export const oauth2JsContent = ` + const OAUTH_CLIENT_ID = 'client-id-123'; + const OAUTH_CLIENT_SECRET = 'client-secret-456'; +` + +export const quotaResponse = [ + { remainingFraction: 0.75, resetTime: '2026-04-24T13:00:00.000Z', modelId: 'gemini-2.5-pro' }, + { remainingFraction: 0.9, resetTime: '2026-04-24T14:00:00.000Z', modelId: 'gemini-2.5-flash' } +] diff --git a/src/main/rate-limits/gemini-usage-fetcher.test.ts b/src/main/rate-limits/gemini-usage-fetcher.test.ts new file mode 100644 index 00000000000..17ff0553b5d --- /dev/null +++ b/src/main/rate-limits/gemini-usage-fetcher.test.ts @@ -0,0 +1,219 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { + authJsonGoogle, + authJsonGoogleExpired, + makeResponse, + quotaResponse +} from './gemini-usage-fetcher.test-fixtures' + +const { readFileMock, extractCredsMock, netFetchMock } = vi.hoisted(() => ({ + readFileMock: vi.fn(), + extractCredsMock: vi.fn(), + netFetchMock: vi.fn() +})) + +// Why: mock the extractor at the module boundary rather than re-routing every +// child_process/fs call. The extractor is a self-contained dependency with a +// simple async contract; mocking it directly keeps tests focused on the +// fetcher's refresh/quota logic rather than on filesystem plumbing that has +// already been integration-tested elsewhere. +vi.mock('./gemini-cli-oauth-extractor', () => ({ + extractOAuthClientCredentials: extractCredsMock +})) + +vi.mock('node:fs/promises', () => ({ + readFile: readFileMock, + writeFile: vi.fn().mockResolvedValue(undefined), + rename: vi.fn().mockResolvedValue(undefined) +})) +vi.mock('electron', () => ({ net: { fetch: netFetchMock } })) + +import { fetchGeminiRateLimits } from './gemini-usage-fetcher' + +describe('fetchGeminiRateLimits', () => { + beforeEach(() => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-04-24T12:00:00.000Z')) + readFileMock.mockReset() + extractCredsMock.mockReset() + netFetchMock.mockReset() + netFetchMock.mockImplementation((url: string) => { + if (url.includes('loadCodeAssist')) { + return Promise.resolve(makeResponse({ cloudaicompanionProject: 'proj-123' })) + } + if (url.includes('token')) { + return Promise.resolve(makeResponse({ access_token: 'new-token', expires_in: 3600 })) + } + return Promise.resolve(makeResponse({ error: `Unhandled fetch to ${url}` }, 500)) + }) + // Default: no CLI installed, refresh path cannot find client credentials. + extractCredsMock.mockResolvedValue(null) + readFileMock.mockRejectedValue({ code: 'ENOENT' }) + }) + + const setupAuthJsonValid = () => { + readFileMock.mockImplementation(async (p: string) => { + if (p.includes('auth.json')) { + return JSON.stringify(authJsonGoogle) + } + throw { code: 'ENOENT' } + }) + } + const setupAuthJsonExpired = () => { + readFileMock.mockImplementation(async (p: string) => { + if (p.includes('auth.json')) { + return JSON.stringify(authJsonGoogleExpired) + } + throw { code: 'ENOENT' } + }) + } + + it('returns unavailable when no credentials exist', async () => { + const result = await fetchGeminiRateLimits(true) + expect(result.status).toBe('unavailable') + }) + + it('returns quota via auth.json', async () => { + setupAuthJsonValid() + netFetchMock.mockImplementation((url: string) => { + if (url.includes('retrieveUserQuota')) { + return Promise.resolve(makeResponse(quotaResponse)) + } + if (url.includes('loadCodeAssist')) { + return Promise.resolve(makeResponse({ cloudaicompanionProject: 'proj-123' })) + } + return Promise.resolve(makeResponse({}, 404)) + }) + const result = await fetchGeminiRateLimits(true) + expect(result.status).toBe('ok') + expect(result.buckets).toHaveLength(2) + }) + + it('deduplicates buckets', async () => { + setupAuthJsonValid() + netFetchMock.mockImplementation((url: string) => { + if (url.includes('retrieveUserQuota')) { + return Promise.resolve( + makeResponse([ + { + remainingFraction: 0.82, + resetTime: '2026-04-24T13:00:00.000Z', + modelId: 'gemini-1.5-flash' + }, + { + remainingFraction: 0.82, + resetTime: '2026-04-24T13:00:00.000Z', + modelId: 'gemini-3-flash-preview' + } + ]) + ) + } + if (url.includes('loadCodeAssist')) { + return Promise.resolve(makeResponse({ cloudaicompanionProject: 'proj-123' })) + } + return Promise.resolve(makeResponse({}, 404)) + }) + const result = await fetchGeminiRateLimits(true) + expect(result.status).toBe('ok') + expect(result.buckets).toHaveLength(1) + expect(result.buckets![0].name).toBe('1.5 Flash') + }) + + it('handles empty bucket list', async () => { + setupAuthJsonValid() + netFetchMock.mockImplementation((url: string) => { + if (url.includes('retrieveUserQuota')) { + return Promise.resolve(makeResponse([])) + } + if (url.includes('loadCodeAssist')) { + return Promise.resolve(makeResponse({ cloudaicompanionProject: 'proj-123' })) + } + return Promise.resolve(makeResponse({}, 404)) + }) + const result = await fetchGeminiRateLimits(true) + expect(result.status).toBe('ok') + expect(result.buckets).toEqual([]) + }) + + it('returns error when token refresh fails', async () => { + vi.useRealTimers() + setupAuthJsonExpired() + const result = await fetchGeminiRateLimits(true) + expect(result.status).toBe('error') + expect(result.error).toContain('Token refresh failed') + vi.useFakeTimers() + }) + + it('handles wrapped buckets response', async () => { + setupAuthJsonValid() + netFetchMock.mockImplementation((url: string) => { + if (url.includes('retrieveUserQuota')) { + return Promise.resolve(makeResponse({ buckets: quotaResponse })) + } + if (url.includes('loadCodeAssist')) { + return Promise.resolve(makeResponse({ cloudaicompanionProject: 'proj-123' })) + } + return Promise.resolve(makeResponse({}, 404)) + }) + const result = await fetchGeminiRateLimits(true) + expect(result.status).toBe('ok') + expect(result.session?.usedPercent).toBe(25) + }) + + it('filters out NaN buckets', async () => { + setupAuthJsonValid() + netFetchMock.mockImplementation((url: string) => { + if (url.includes('retrieveUserQuota')) { + return Promise.resolve( + makeResponse([ + { + remainingFraction: NaN, + resetTime: '2026-04-24T13:00:00.000Z', + modelId: 'gemini-1.5-pro' + }, + { + remainingFraction: 0.9, + resetTime: '2026-04-24T13:00:00.000Z', + modelId: 'gemini-1.5-flash' + } + ]) + ) + } + if (url.includes('loadCodeAssist')) { + return Promise.resolve(makeResponse({ cloudaicompanionProject: 'proj-123' })) + } + return Promise.resolve(makeResponse({}, 404)) + }) + const result = await fetchGeminiRateLimits(true) + expect(result.status).toBe('ok') + expect(result.buckets).toHaveLength(1) + }) + + it('retries refresh on 401', async () => { + setupAuthJsonValid() + extractCredsMock.mockResolvedValue({ clientId: 'cid', clientSecret: 'csec' }) + let quotaCallCount = 0 + netFetchMock.mockImplementation((url: string) => { + if (url.includes('retrieveUserQuota')) { + quotaCallCount += 1 + if (quotaCallCount === 1) { + return Promise.resolve(makeResponse({ error: 'Unauthenticated' }, 401)) + } + return Promise.resolve(makeResponse(quotaResponse)) + } + if (url.includes('token')) { + return Promise.resolve( + makeResponse({ access_token: 'retried-token', expires_in: 3600 }) + ) + } + if (url.includes('loadCodeAssist')) { + return Promise.resolve(makeResponse({ cloudaicompanionProject: 'proj-123' })) + } + return Promise.resolve(makeResponse({}, 404)) + }) + const result = await fetchGeminiRateLimits(true) + expect(result.status).toBe('ok') + // The second quota call should have been made with the refreshed token. + expect(quotaCallCount).toBe(2) + }) +}) diff --git a/src/main/rate-limits/gemini-usage-fetcher.ts b/src/main/rate-limits/gemini-usage-fetcher.ts new file mode 100644 index 00000000000..95b710ef91a --- /dev/null +++ b/src/main/rate-limits/gemini-usage-fetcher.ts @@ -0,0 +1,245 @@ +import { net } from 'electron' +import type { ProviderRateLimits } from '../../shared/rate-limit-types' +import { + loadProjectId, + readAuthJson, + readGeminiCredentials, + saveGeminiCredentials, + tryRefreshTokenFromBundle, + type GeminiCredentials, + type GoogleAuthEntry +} from './gemini-oauth-sources' +import { + buildRateLimitBucket, + deduplicateBuckets, + deriveSessionSummary +} from './gemini-bucket-formatting' + +const API_TIMEOUT_MS = 10_000 +const RETRIEVE_QUOTA_URL = 'https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota' + +type QuotaBucket = { remainingFraction: number; resetTime: string; modelId: string } + +function isQuotaBucket(o: unknown): o is QuotaBucket { + return ( + typeof o === 'object' && + o !== null && + typeof (o as QuotaBucket).remainingFraction === 'number' && + Number.isFinite((o as QuotaBucket).remainingFraction) && + typeof (o as QuotaBucket).resetTime === 'string' && + typeof (o as QuotaBucket).modelId === 'string' + ) +} + +function parseQuotaResponse(data: unknown): QuotaBucket[] { + let rawBuckets: unknown[] = [] + if (Array.isArray(data)) { + rawBuckets = data + } else if (data && typeof data === 'object' && 'buckets' in data && Array.isArray(data.buckets)) { + rawBuckets = data.buckets + } + return rawBuckets.filter((b) => isQuotaBucket(b)) +} + +async function fetchQuota(accessToken: string, projectId: string): Promise { + const controller = new AbortController() + const timeout = setTimeout(() => { + controller.abort() + }, API_TIMEOUT_MS) + try { + const res = await net.fetch(RETRIEVE_QUOTA_URL, { + method: 'POST', + headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${accessToken}` }, + body: JSON.stringify({ project: projectId }), + signal: controller.signal + }) + if (!res.ok) { + return { + provider: 'gemini', + session: null, + weekly: null, + updatedAt: Date.now(), + error: `Quota fetch failed (${res.status})`, + status: 'error' + } + } + const data = (await res.json()) as unknown + const buckets = deduplicateBuckets( + parseQuotaResponse(data).map((b) => ({ ...buildRateLimitBucket(b), modelId: b.modelId })) + ) + return { + provider: 'gemini', + session: deriveSessionSummary(buckets), + weekly: null, + buckets, + updatedAt: Date.now(), + error: null, + status: 'ok' + } + } finally { + clearTimeout(timeout) + } +} + +async function fetchViaAuthJson( + auth: GoogleAuthEntry, + geminiCliOAuthEnabled = false +): Promise { + let accessToken = auth.access + const refreshToken = (auth.refresh || '').split('|')[0] ?? '' + if (auth.expires < Date.now() || !accessToken) { + const refreshResult = await tryRefreshTokenFromBundle(refreshToken, geminiCliOAuthEnabled) + if (!refreshResult?.accessToken) { + return { + provider: 'gemini', + session: null, + weekly: null, + updatedAt: Date.now(), + error: 'Token refresh failed', + status: 'error' + } + } + accessToken = refreshResult.accessToken + } + let effectiveProjectId = '' + try { + effectiveProjectId = await loadProjectId(accessToken) + } catch { + effectiveProjectId = + (auth.refresh || '').split('|')[1] || (auth.refresh || '').split('|')[2] || '' + } + if (!effectiveProjectId) { + return { + provider: 'gemini', + session: null, + weekly: null, + updatedAt: Date.now(), + error: 'Gemini project ID not found', + status: 'error' + } + } + const result = await fetchQuota(accessToken, effectiveProjectId) + if (result.status === 'error' && result.error?.includes('401')) { + const refreshResult = await tryRefreshTokenFromBundle(refreshToken, geminiCliOAuthEnabled) + if (refreshResult?.accessToken) { + const newProjectId = await loadProjectId(refreshResult.accessToken).catch(() => { + return effectiveProjectId + }) + return fetchQuota(refreshResult.accessToken, newProjectId) + } + } + return result +} + +async function fetchViaOauthCreds( + creds: GeminiCredentials, + geminiCliOAuthEnabled = false +): Promise { + let accessToken = creds.access_token + let currentCreds = creds + if (creds.expiry_date < Date.now()) { + const refreshResult = await tryRefreshTokenFromBundle( + creds.refresh_token, + geminiCliOAuthEnabled + ) + if (!refreshResult?.accessToken) { + return { + provider: 'gemini', + session: null, + weekly: null, + updatedAt: Date.now(), + error: 'Token refresh failed', + status: 'error' + } + } + accessToken = refreshResult.accessToken + currentCreds = { + ...creds, + access_token: accessToken, + expiry_date: refreshResult.expiresIn + ? Date.now() + refreshResult.expiresIn * 1000 + : creds.expiry_date + } + await saveGeminiCredentials(currentCreds) + } + const projectId = await loadProjectId(accessToken).catch(() => { + return '' + }) + if (!projectId) { + return { + provider: 'gemini', + session: null, + weekly: null, + updatedAt: Date.now(), + error: 'Gemini project ID not found', + status: 'error' + } + } + const result = await fetchQuota(accessToken, projectId) + if (result.status === 'error' && result.error?.includes('401')) { + const refreshResult = await tryRefreshTokenFromBundle( + currentCreds.refresh_token, + geminiCliOAuthEnabled + ) + if (refreshResult?.accessToken) { + const newProjectId = await loadProjectId(refreshResult.accessToken).catch(() => { + return '' + }) + if (newProjectId) { + await saveGeminiCredentials({ + ...currentCreds, + access_token: refreshResult.accessToken, + expiry_date: refreshResult.expiresIn + ? Date.now() + refreshResult.expiresIn * 1000 + : currentCreds.expiry_date + }) + return fetchQuota(refreshResult.accessToken, newProjectId) + } + } + } + return result +} + +export async function fetchGeminiRateLimits( + geminiCliOAuthEnabled = false +): Promise { + try { + const authJson = await readAuthJson() + const result = + authJson?.google?.type === 'oauth' + ? await fetchViaAuthJson(authJson.google, geminiCliOAuthEnabled) + : await (async () => { + if (!geminiCliOAuthEnabled) { + return { + provider: 'gemini', + session: null, + weekly: null, + updatedAt: Date.now(), + error: 'Gemini CLI OAuth is disabled in settings', + status: 'unavailable' + } as ProviderRateLimits + } + const creds = await readGeminiCredentials() + return !creds + ? ({ + provider: 'gemini', + session: null, + weekly: null, + updatedAt: Date.now(), + error: 'Gemini CLI credentials not found', + status: 'unavailable' + } as ProviderRateLimits) + : await fetchViaOauthCreds(creds, geminiCliOAuthEnabled) + })() + return result + } catch (err) { + return { + provider: 'gemini', + session: null, + weekly: null, + updatedAt: Date.now(), + error: err instanceof Error ? err.message : 'Unknown error', + status: 'error' + } + } +} diff --git a/src/main/rate-limits/opencode-go-page-scraper.ts b/src/main/rate-limits/opencode-go-page-scraper.ts new file mode 100644 index 00000000000..0f092da6db3 --- /dev/null +++ b/src/main/rate-limits/opencode-go-page-scraper.ts @@ -0,0 +1,172 @@ +// Why: the opencode.ai page is rendered with React Server Components. The +// embedded JS uses a wire format where object references look like: +// key:$R[28]={field:value,...} +// rather than plain `key:{field:value,...}`. A single key (e.g. monthlyUsage) +// can appear multiple times — once with real data and once as `null` inside a +// different component's props. We must find the occurrence that is an object +// with both usagePercent and resetInSec, not the null one. + +/** + * Finds the brace-balanced object block assigned to `key` anywhere in `text`. + * Skips React Flight assignment tokens (e.g. `$R[N]=`) between the colon and + * the opening brace. Returns the first block that contains `usagePercent` AND + * `resetInSec` as direct numeric properties (not nested), so that placeholder + * `null` occurrences and billing-context duplicates are ignored. + */ +function extractUsageBlock(text: string, key: string): string | null { + // Match every occurrence of `key:` (with optional $R[N]= assignment) + // Why: React Flight wire format embeds object references between the colon + // and the literal brace, so we skip over any `$R[N]=` tokens to reach `{`. + const keyRegex = new RegExp(`\\b${key}\\b\\s*:`, 'g') + let keyMatch: RegExpExecArray | null + + while ((keyMatch = keyRegex.exec(text)) !== null) { + // Scan forward from after the colon to find the opening `{`, + // allowing for the `$R[N]=` token or plain whitespace in between. + // We only scan a short window so we don't accidentally land on the + // next occurrence of the key. + const searchStart = keyMatch.index + keyMatch[0].length + const searchWindow = text.slice(searchStart, searchStart + 30) + const braceOffset = searchWindow.indexOf('{') + if (braceOffset === -1) { + // This occurrence has no object (e.g. `monthlyUsage:null`) — skip. + continue + } + + const openBrace = searchStart + braceOffset + // Extract the balanced block + // Why: this brace-depth parser does not skip string literals. React Flight's + // current format does not emit raw { } inside strings, but this is a scraper + // against HTML we don't control — treat as fragile. + let depth = 0 + let block: string | null = null + for (let i = openBrace; i < text.length; i++) { + if (text[i] === '{') { + depth++ + } else if (text[i] === '}') { + depth-- + if (depth === 0) { + block = text.slice(openBrace, i + 1) + break + } + } + } + + if (!block) { + continue + } + + // Verify this block has both required numeric fields as direct properties + // (depth 1 within the block). This rejects billing/plan objects that share + // the key name but lack usage data. + if ( + hasDirectNumericField(block, 'usagePercent') && + hasDirectNumericField(block, 'resetInSec') + ) { + return block + } + } + + return null +} + +/** + * Returns true if `fieldName` exists as a direct (depth-1) numeric property + * of the object string `objText`. + */ +function hasDirectNumericField(objText: string, fieldName: string): boolean { + return extractTopLevelNumber(objText, fieldName) !== null +} + +/** + * Extracts a numeric field at depth 1 of `objText` — ignores the same field + * inside nested sub-objects. + * Why: without depth tracking, a regex matches the first occurrence regardless + * of nesting, returning wrong values when a sub-object contains the same name. + */ +function extractTopLevelNumber(objText: string, fieldName: string): number | null { + const fieldRegex = new RegExp(`\\b${fieldName}\\b\\s*:\\s*(-?[0-9]+(?:\\.[0-9]+)?)`) + // Why: this brace-depth parser does not skip string literals. React Flight's + // current format does not emit raw { } inside strings, but this is a scraper + // against HTML we don't control — treat as fragile. + let depth = 0 + + for (let i = 0; i < objText.length; i++) { + const ch = objText[i] + if (ch === '{') { + depth++ + continue + } + if (ch === '}') { + depth-- + continue + } + + // Only match at depth 1 (direct property of the root object). + if (depth === 1) { + const slice = objText.slice(i, i + fieldName.length + 30) + const m = fieldRegex.exec(slice) + if (m && m.index === 0) { + const n = Number.parseFloat(m[1]) + return Number.isFinite(n) ? n : null + } + } + } + return null +} + +type ParsedSubscription = { + rollingUsagePercent: number + weeklyUsagePercent: number + monthlyUsagePercent: number | null + rollingResetInSec: number + weeklyResetInSec: number + monthlyResetInSec: number | null +} + +export function parseSubscriptionFromPageText(text: string): ParsedSubscription | null { + // Why: OpenCode usage is scraped from HTML-embedded JS (React Flight wire + // format). Defensive size check prevents runaway parsing on unexpected payloads. + if (!text || text.length > 10_000_000) { + return null + } + + // Find the first occurrence of each usage key that has both usagePercent and + // resetInSec as direct numeric fields. This skips null occurrences and + // billing-context duplicates that use the same key name without usage data. + const rollingBlock = extractUsageBlock(text, 'rollingUsage') + const weeklyBlock = extractUsageBlock(text, 'weeklyUsage') + const monthlyBlock = extractUsageBlock(text, 'monthlyUsage') + + const rollingPercent = + rollingBlock !== null ? extractTopLevelNumber(rollingBlock, 'usagePercent') : null + const rollingReset = + rollingBlock !== null ? extractTopLevelNumber(rollingBlock, 'resetInSec') : null + const weeklyPercent = + weeklyBlock !== null ? extractTopLevelNumber(weeklyBlock, 'usagePercent') : null + const weeklyReset = weeklyBlock !== null ? extractTopLevelNumber(weeklyBlock, 'resetInSec') : null + + if ( + rollingPercent === null || + rollingReset === null || + weeklyPercent === null || + weeklyReset === null + ) { + return null + } + + const monthlyPercent = + monthlyBlock !== null ? extractTopLevelNumber(monthlyBlock, 'usagePercent') : null + const monthlyReset = + monthlyBlock !== null ? extractTopLevelNumber(monthlyBlock, 'resetInSec') : null + + return { + rollingUsagePercent: Math.min(100, Math.max(0, rollingPercent)), + weeklyUsagePercent: Math.min(100, Math.max(0, weeklyPercent)), + monthlyUsagePercent: + monthlyPercent !== null ? Math.min(100, Math.max(0, monthlyPercent)) : null, + rollingResetInSec: rollingReset, + weeklyResetInSec: weeklyReset, + monthlyResetInSec: monthlyReset + } +} diff --git a/src/main/rate-limits/opencode-go-usage-fetcher.test.ts b/src/main/rate-limits/opencode-go-usage-fetcher.test.ts new file mode 100644 index 00000000000..e6057c52571 --- /dev/null +++ b/src/main/rate-limits/opencode-go-usage-fetcher.test.ts @@ -0,0 +1,350 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const netFetchMock = vi.hoisted(() => vi.fn()) + +vi.mock('electron', () => ({ + net: { fetch: netFetchMock } +})) + +import { fetchOpenCodeGoRateLimits, normalizeCookieInput } from './opencode-go-usage-fetcher' + +const WORKSPACES_SERVER_ID = 'def39973159c7f0483d8793a822b8dbb10d067e12c65455fcb4608459ba0234f' + +function makeResponse(body: string, status = 200): Response { + return { + ok: status >= 200 && status < 300, + status, + text: async () => body + } as Response +} + +// Real React Flight wire format from opencode.ai — keys like `monthlyUsage` +// appear multiple times: once with actual data (as `$R[N]={...}`) and once as +// `null` inside a billing-context object. The parser must pick the data one. +const USAGE_PAGE_WITH_MONTHLY = ` + +` + +const USAGE_PAGE_NO_MONTHLY = ` + +` + +const WORKSPACES_RESPONSE = 'id: "wrk_TESTWORKSPACEID123"' + +describe('fetchOpenCodeGoRateLimits', () => { + beforeEach(() => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-04-24T12:00:00.000Z')) + netFetchMock.mockReset() + }) + + it('returns unavailable when cookie is empty', async () => { + const result = await fetchOpenCodeGoRateLimits('') + + expect(result.status).toBe('unavailable') + expect(result.provider).toBe('opencode-go') + expect(result.session).toBeNull() + expect(result.weekly).toBeNull() + expect(result.monthly).toBeNull() + expect(result.error).toBe('Session cookie not configured') + expect(netFetchMock).not.toHaveBeenCalled() + }) + + it('returns unavailable when cookie is only whitespace', async () => { + const result = await fetchOpenCodeGoRateLimits(' ') + + expect(result.status).toBe('unavailable') + expect(netFetchMock).not.toHaveBeenCalled() + }) + + it('returns error when cookie has no auth or __Host-auth name', async () => { + const result = await fetchOpenCodeGoRateLimits('session=abc123; other=xyz') + + expect(result.status).toBe('error') + expect(result.error).toMatch(/No auth cookie found/) + expect(netFetchMock).not.toHaveBeenCalled() + }) + + describe('normalizeCookieInput', () => { + it('returns empty string unchanged', () => { + expect(normalizeCookieInput('')).toBe('') + expect(normalizeCookieInput(' ')).toBe('') + }) + + it('wraps a bare token as auth=', () => { + expect(normalizeCookieInput('Fe26.2**abc123')).toBe('auth=Fe26.2**abc123') + }) + + it('leaves auth=... unchanged', () => { + expect(normalizeCookieInput('auth=Fe26.2**abc123')).toBe('auth=Fe26.2**abc123') + }) + + it('leaves __Host-auth=... unchanged', () => { + expect(normalizeCookieInput('__Host-auth=token')).toBe('__Host-auth=token') + }) + + it('leaves multi-pair cookie headers unchanged', () => { + expect(normalizeCookieInput('auth=tok; other=val')).toBe('auth=tok; other=val') + }) + + it('trims surrounding whitespace before wrapping', () => { + expect(normalizeCookieInput(' Fe26.2**abc ')).toBe('auth=Fe26.2**abc') + }) + + it('does not wrap unknown or malformed tokens', () => { + expect(normalizeCookieInput('invalid token format')).toBe('invalid token format') + expect(normalizeCookieInput('{}')).toBe('{}') + expect(normalizeCookieInput('{"token":"abc"}')).toBe('{"token":"abc"}') + }) + }) + + it('accepts a bare token (auto-wraps to auth=)', async () => { + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse(USAGE_PAGE_WITH_MONTHLY)) + + const result = await fetchOpenCodeGoRateLimits('Fe26.2**baretoken') + + expect(result.status).toBe('ok') + // Cookie sent to the server must be auth=, not the bare value. + expect(netFetchMock.mock.calls[0][1].headers.Cookie).toBe('auth=Fe26.2**baretoken') + }) + + it('uses GET /_server?id= with correct headers for workspaces', async () => { + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse(USAGE_PAGE_WITH_MONTHLY)) + + await fetchOpenCodeGoRateLimits('auth=mytoken') + + expect(netFetchMock).toHaveBeenNthCalledWith( + 1, + `https://opencode.ai/_server?id=${WORKSPACES_SERVER_ID}`, + expect.objectContaining({ + method: 'GET', + headers: expect.objectContaining({ + Cookie: 'auth=mytoken', + 'X-Server-Id': WORKSPACES_SERVER_ID + }) + }) + ) + }) + + it('fetches usage from /workspace//go after resolving workspace ID', async () => { + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse(USAGE_PAGE_WITH_MONTHLY)) + + await fetchOpenCodeGoRateLimits('auth=mytoken') + + expect(netFetchMock).toHaveBeenNthCalledWith( + 2, + 'https://opencode.ai/workspace/wrk_TESTWORKSPACEID123/go', + expect.objectContaining({ method: 'GET' }) + ) + }) + + it('returns ok with session, weekly, and monthly windows', async () => { + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse(USAGE_PAGE_WITH_MONTHLY)) + + const now = Date.now() + const result = await fetchOpenCodeGoRateLimits('auth=mytoken') + + expect(result.status).toBe('ok') + expect(result.error).toBeNull() + + expect(result.session).toEqual({ + usedPercent: 30, + windowMinutes: 300, + resetsAt: now + 7200 * 1000, + resetDescription: null + }) + expect(result.weekly).toEqual({ + usedPercent: 51, + windowMinutes: 10080, + resetsAt: now + 259200 * 1000, + resetDescription: null + }) + expect(result.monthly).toEqual({ + usedPercent: 89, + windowMinutes: 43200, + resetsAt: now + 1296000 * 1000, + resetDescription: null + }) + }) + + it('returns ok with null monthly when monthlyUsage is absent', async () => { + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse(USAGE_PAGE_NO_MONTHLY)) + + const result = await fetchOpenCodeGoRateLimits('auth=mytoken') + + expect(result.status).toBe('ok') + expect(result.session?.usedPercent).toBe(10) + expect(result.weekly?.usedPercent).toBe(20) + expect(result.monthly).toBeNull() + }) + + it('caps usedPercent at 100 and floors at 0', async () => { + const page = ` + rollingUsage: { usagePercent: 150, resetInSec: 3600 } + weeklyUsage: { usagePercent: -5, resetInSec: 86400 } + ` + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse(page)) + + const result = await fetchOpenCodeGoRateLimits('auth=token') + + expect(result.status).toBe('ok') + expect(result.session?.usedPercent).toBe(100) + expect(result.weekly?.usedPercent).toBe(0) + }) + + it('parses React Flight wire format with $R[N]= assignment tokens', async () => { + // Real format from opencode.ai — keys have $R[N]= between the colon and brace. + const page = ` + rollingUsage:$R[21]={status:"ok",resetInSec:1337,usagePercent:42}, + weeklyUsage:$R[22]={status:"ok",resetInSec:86400,usagePercent:68} + ` + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse(page)) + + const result = await fetchOpenCodeGoRateLimits('auth=token') + + expect(result.status).toBe('ok') + expect(result.session?.usedPercent).toBe(42) + expect(result.weekly?.usedPercent).toBe(68) + }) + + it('skips null occurrences and finds the real data block for monthlyUsage', async () => { + // Regression: on refresh, monthlyUsage:null appeared BEFORE the real + // monthlyUsage:$R[N]={usagePercent:89,...} in a different component's props. + // Parser must skip the null and find the data block. + const page = ` + rollingUsage:$R[21]={status:"ok",resetInSec:18000,usagePercent:0}, + weeklyUsage:$R[22]={status:"ok",resetInSec:57781,usagePercent:51}, + monthlyUsage:null,timeMonthlyUsageUpdated:null, + monthlyUsage:$R[28]={status:"ok",resetInSec:1214779,usagePercent:89} + ` + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse(page)) + + const result = await fetchOpenCodeGoRateLimits('auth=token') + + expect(result.status).toBe('ok') + expect(result.monthly?.usedPercent).toBe(89) + expect(result.monthly?.resetsAt).toBe(Date.now() + 1214779 * 1000) + }) + + it('returns null monthly when all monthlyUsage occurrences are null', async () => { + const page = ` + rollingUsage:$R[21]={status:"ok",resetInSec:3600,usagePercent:10}, + weeklyUsage:$R[22]={status:"ok",resetInSec:86400,usagePercent:20}, + monthlyUsage:null,timeMonthlyUsageUpdated:null + ` + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse(page)) + + const result = await fetchOpenCodeGoRateLimits('auth=token') + + expect(result.status).toBe('ok') + expect(result.monthly).toBeNull() + }) + + it('skips workspace lookup when workspaceIdOverride is provided', async () => { + netFetchMock.mockResolvedValueOnce(makeResponse(USAGE_PAGE_WITH_MONTHLY)) + + const result = await fetchOpenCodeGoRateLimits('auth=mytoken', 'wrk_OVERRIDE123') + + expect(netFetchMock).toHaveBeenCalledTimes(1) + expect(netFetchMock).toHaveBeenCalledWith( + 'https://opencode.ai/workspace/wrk_OVERRIDE123/go', + expect.anything() + ) + expect(result.status).toBe('ok') + }) + + it('filters cookie to auth name only', async () => { + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse(USAGE_PAGE_WITH_MONTHLY)) + + await fetchOpenCodeGoRateLimits('session=secret; auth=realtoken; tracking=xyz') + + const firstCall = netFetchMock.mock.calls[0] + expect(firstCall[1].headers.Cookie).toBe('auth=realtoken') + }) + + it('returns error on 404 from workspaces fetch', async () => { + netFetchMock.mockResolvedValueOnce(makeResponse('Not Found', 404)) + + const result = await fetchOpenCodeGoRateLimits('auth=mytoken') + + expect(result.status).toBe('error') + expect(result.error).toBe('Workspaces fetch failed (404)') + expect(result.session).toBeNull() + }) + + it('returns error on 401 from workspaces fetch', async () => { + netFetchMock.mockResolvedValueOnce(makeResponse('Unauthorized', 401)) + + const result = await fetchOpenCodeGoRateLimits('auth=mytoken') + + expect(result.status).toBe('error') + expect(result.error).toBe('Workspaces fetch failed (401)') + }) + + it('returns error when no workspace ID found in response', async () => { + netFetchMock.mockResolvedValueOnce(makeResponse('no workspace id here')) + + const result = await fetchOpenCodeGoRateLimits('auth=mytoken') + + expect(result.status).toBe('error') + expect(result.error).toMatch(/No workspace ID found/) + }) + + it('returns error on non-ok usage page response', async () => { + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse('Not Found', 404)) + + const result = await fetchOpenCodeGoRateLimits('auth=mytoken') + + expect(result.status).toBe('error') + expect(result.error).toBe('Usage page fetch failed (404)') + }) + + it('returns error when usage data cannot be parsed from page', async () => { + netFetchMock + .mockResolvedValueOnce(makeResponse(WORKSPACES_RESPONSE)) + .mockResolvedValueOnce(makeResponse('no usage data here')) + + const result = await fetchOpenCodeGoRateLimits('auth=mytoken') + + expect(result.status).toBe('error') + expect(result.error).toBe('Could not parse usage data from page') + }) + + it('never logs the cookie in error messages', async () => { + netFetchMock.mockRejectedValueOnce(new Error('network timeout')) + + const result = await fetchOpenCodeGoRateLimits('auth=secret123') + + expect(result.status).toBe('error') + expect(result.error).toBe('network timeout') + expect(result.error).not.toContain('secret123') + }) +}) diff --git a/src/main/rate-limits/opencode-go-usage-fetcher.ts b/src/main/rate-limits/opencode-go-usage-fetcher.ts new file mode 100644 index 00000000000..26c50d8e6bd --- /dev/null +++ b/src/main/rate-limits/opencode-go-usage-fetcher.ts @@ -0,0 +1,256 @@ +import { net } from 'electron' +import { randomUUID } from 'crypto' +import type { ProviderRateLimits, RateLimitWindow } from '../../shared/rate-limit-types' +import { parseSubscriptionFromPageText } from './opencode-go-page-scraper' + +const OPENCODE_BASE_URL = 'https://opencode.ai' +const OPENCODE_SERVER_URL = 'https://opencode.ai/_server' +const API_TIMEOUT_MS = 15_000 + +// Server-function hash for the workspaces endpoint — stable identifier used by +// the opencode.ai SST/TanStack router server-fn protocol. +const WORKSPACES_SERVER_ID = 'def39973159c7f0483d8793a822b8dbb10d067e12c65455fcb4608459ba0234f' + +// Only these cookie names carry session auth on opencode.ai. Sending unrelated +// cookies pollutes the header and can expose sensitive data from other sites. +const AUTH_COOKIE_NAMES = new Set(['auth', '__Host-auth']) + +// Why: users may paste just the token value (e.g. "Fe26.2**...") instead of +// the full cookie header ("auth=Fe26.2**..."). Auto-wrapping avoids a confusing +// silent failure where the cookie looks non-empty but contains no auth name. +export function normalizeCookieInput(raw: string): string { + const trimmed = raw.trim() + if (!trimmed) { + return trimmed + } + // Already a valid cookie header: has multiple pairs or starts with known name. + if (trimmed.includes(';') || /^(?:auth|__Host-auth)=/i.test(trimmed)) { + return trimmed + } + // Only wrap if it looks like an Iron Session seal (starts with Fe26.2**) + // or a reasonably structured bare token (alphanumeric with dots/dashes). + // Otherwise, leave it alone to fail predictably instead of sending malformed auth. + if (trimmed.startsWith('Fe26.2**') || /^[a-zA-Z0-9.\-_]+$/.test(trimmed)) { + return `auth=${trimmed}` + } + return trimmed +} + +function filterAuthCookie(raw: string): string { + return raw + .split(';') + .map((p) => p.trim()) + .filter((pair) => { + const eq = pair.indexOf('=') + if (eq < 0) { + return false + } + return AUTH_COOKIE_NAMES.has(pair.slice(0, eq).trim()) + }) + .join('; ') +} + +function parseWorkspaceIds(text: string): string[] { + // Match id:"wrk_..." or id: "wrk_..." patterns in JS-serialized output. + // Why: Workspace IDs follow a 'wrk_xxx' or 'wk_xxx' pattern. Using a + // more specific regex with word boundaries avoids picking up unrelated + // object properties that might match a generic ID pattern. + const ids: string[] = [] + const workspaceIdRegex = /\bid\s*:\s*["']((?:wrk|wk)_[a-zA-Z0-9]+)["']/g + for (const match of text.matchAll(workspaceIdRegex)) { + const id = match[1] + if (id && !ids.includes(id)) { + ids.push(id) + } + } + return ids +} + +function makeWindow( + usedPercent: number, + resetInSec: number, + windowMinutes: number +): RateLimitWindow { + return { + usedPercent, + windowMinutes, + resetsAt: Date.now() + resetInSec * 1000, + resetDescription: null + } +} + +export async function fetchOpenCodeGoRateLimits( + cookie: string, + workspaceIdOverride?: string +): Promise { + // Normalize before any guard — bare tokens become auth=. + const normalizedCookie = normalizeCookieInput(cookie) + + if (!normalizedCookie) { + return { + provider: 'opencode-go', + session: null, + weekly: null, + monthly: null, + updatedAt: Date.now(), + error: 'Session cookie not configured', + status: 'unavailable' + } + } + + // Filter to only auth cookies — avoids sending unrelated session data. + const cookieHeader = filterAuthCookie(normalizedCookie) + if (!cookieHeader) { + return { + provider: 'opencode-go', + session: null, + weekly: null, + monthly: null, + updatedAt: Date.now(), + error: 'No auth cookie found — paste the full Cookie header from opencode.ai DevTools', + status: 'error' + } + } + + // Step 1: resolve workspace IDs to try. + let ids: string[] = [] + const override = workspaceIdOverride?.trim() + + if (override) { + if (!/^(wrk|wk)_[A-Za-z0-9]+$/.test(override)) { + return { + provider: 'opencode-go', + session: null, + weekly: null, + monthly: null, + updatedAt: Date.now(), + error: 'Invalid workspace ID format: must match ^(wrk|wk)_[A-Za-z0-9]+$', + status: 'error' + } + } + ids = [override] + } else { + const workspacesController = new AbortController() + const workspacesTimeout = setTimeout(() => workspacesController.abort(), API_TIMEOUT_MS) + try { + // The /_server endpoint uses SST server-function protocol: GET with ?id= + // and X-Server-Id / X-Server-Instance headers for routing. + const instanceId = `server-fn:${randomUUID()}` + const workspacesUrl = `${OPENCODE_SERVER_URL}?id=${WORKSPACES_SERVER_ID}` + const workspacesRes = await net.fetch(workspacesUrl, { + method: 'GET', + headers: { + Cookie: cookieHeader, + 'X-Server-Id': WORKSPACES_SERVER_ID, + 'X-Server-Instance': instanceId, + Accept: 'text/javascript, application/json;q=0.9, */*;q=0.8', + Origin: OPENCODE_BASE_URL, + Referer: OPENCODE_BASE_URL + }, + signal: workspacesController.signal + }) + + if (!workspacesRes.ok) { + return { + provider: 'opencode-go', + session: null, + weekly: null, + monthly: null, + updatedAt: Date.now(), + error: `Workspaces fetch failed (${workspacesRes.status})`, + status: 'error' + } + } + + const workspacesText = await workspacesRes.text() + ids = parseWorkspaceIds(workspacesText) + } catch (err) { + const message = err instanceof Error ? err.message : 'Unknown error' + return { + provider: 'opencode-go', + session: null, + weekly: null, + monthly: null, + updatedAt: Date.now(), + error: message, + status: 'error' + } + } finally { + clearTimeout(workspacesTimeout) + } + } + + if (ids.length === 0) { + return { + provider: 'opencode-go', + session: null, + weekly: null, + monthly: null, + updatedAt: Date.now(), + error: 'No workspace ID found — set a Workspace ID override in settings', + status: 'error' + } + } + + // Step 2: Robust workspace resolution. Try each candidate ID until one returns 200 OK + // and valid usage data. Each candidate gets its own AbortController so a slow or + // hung candidate cannot starve the rest. + let lastError = '' + for (const candidateId of ids) { + const candidateController = new AbortController() + const candidateTimeout = setTimeout(() => candidateController.abort(), API_TIMEOUT_MS) + try { + const usagePageUrl = `${OPENCODE_BASE_URL}/workspace/${candidateId}/go` + const pageRes = await net.fetch(usagePageUrl, { + method: 'GET', + headers: { + Cookie: cookieHeader, + Accept: 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8', + Origin: OPENCODE_BASE_URL, + Referer: OPENCODE_BASE_URL + }, + signal: candidateController.signal + }) + + if (!pageRes.ok) { + lastError = `Usage page fetch failed (${pageRes.status})` + continue + } + + const pageText = await pageRes.text() + const parsed = parseSubscriptionFromPageText(pageText) + if (parsed) { + const monthly = + parsed.monthlyUsagePercent !== null && parsed.monthlyResetInSec !== null + ? makeWindow(parsed.monthlyUsagePercent, parsed.monthlyResetInSec, 43200) // 30d + : null + + return { + provider: 'opencode-go', + session: makeWindow(parsed.rollingUsagePercent, parsed.rollingResetInSec, 300), + weekly: makeWindow(parsed.weeklyUsagePercent, parsed.weeklyResetInSec, 10080), + monthly, + updatedAt: Date.now(), + error: null, + status: 'ok' + } + } + lastError = 'Could not parse usage data from page' + } catch (err) { + const message = err instanceof Error ? err.message : 'Unknown error' + lastError = message + } finally { + clearTimeout(candidateTimeout) + } + } + + return { + provider: 'opencode-go', + session: null, + weekly: null, + monthly: null, + updatedAt: Date.now(), + error: lastError || 'Could not parse usage data from any available workspace', + status: 'error' + } +} diff --git a/src/main/rate-limits/service.test.ts b/src/main/rate-limits/service.test.ts index 346539b00d8..0146ea88b41 100644 --- a/src/main/rate-limits/service.test.ts +++ b/src/main/rate-limits/service.test.ts @@ -1,8 +1,14 @@ +/* eslint-disable max-lines -- Why: these tests mirror the fetch ordering, +stale-data handling, account-switch generation, and OpenCode config-change +semantics covered in service.ts, which already carries the same pragma. +Keeping them in one file makes the ordering contract reviewable as a unit. */ import { beforeEach, describe, expect, it, vi } from 'vitest' import type { ProviderRateLimits } from '../../shared/rate-limit-types' import { RateLimitService } from './service' import { fetchClaudeRateLimits } from './claude-fetcher' import { fetchCodexRateLimits } from './codex-fetcher' +import { fetchGeminiRateLimits } from './gemini-usage-fetcher' +import { fetchOpenCodeGoRateLimits } from './opencode-go-usage-fetcher' vi.mock('./claude-fetcher', () => ({ fetchClaudeRateLimits: vi.fn() @@ -12,6 +18,14 @@ vi.mock('./codex-fetcher', () => ({ fetchCodexRateLimits: vi.fn() })) +vi.mock('./gemini-usage-fetcher', () => ({ + fetchGeminiRateLimits: vi.fn() +})) + +vi.mock('./opencode-go-usage-fetcher', () => ({ + fetchOpenCodeGoRateLimits: vi.fn() +})) + type Deferred = { promise: Promise resolve: (value: T) => void @@ -26,7 +40,7 @@ function deferred(): Deferred { } function okProvider( - provider: 'claude' | 'codex', + provider: 'claude' | 'codex' | 'gemini' | 'opencode-go', usedPercent: number, updatedAt = Date.now() ): ProviderRateLimits { @@ -45,7 +59,10 @@ function okProvider( } } -function errorProvider(provider: 'claude' | 'codex', message: string): ProviderRateLimits { +function errorProvider( + provider: 'claude' | 'codex' | 'gemini' | 'opencode-go', + message: string +): ProviderRateLimits { return { provider, session: null, @@ -63,6 +80,8 @@ function serviceInternals(service: RateLimitService): { fetchAll: () => Promise< describe('RateLimitService', () => { beforeEach(() => { vi.clearAllMocks() + vi.mocked(fetchGeminiRateLimits).mockResolvedValue(okProvider('gemini', 0, Date.now())) + vi.mocked(fetchOpenCodeGoRateLimits).mockResolvedValue(okProvider('opencode-go', 0, Date.now())) }) it('does not refetch Claude when a Codex account switch is queued during fetchAll', async () => { @@ -176,4 +195,180 @@ describe('RateLimitService', () => { expect(fetchClaudeRateLimits).toHaveBeenCalledTimes(2) expect(fetchCodexRateLimits).toHaveBeenCalledTimes(2) }) + + it('fetches Gemini and OpenCode Go alongside Claude and Codex', async () => { + const service = new RateLimitService() + service.setSettingsResolver(() => ({ + opencodeSessionCookie: 'session=abc123', + opencodeWorkspaceId: '' + })) + + vi.mocked(fetchClaudeRateLimits).mockResolvedValueOnce(okProvider('claude', 10, Date.now())) + vi.mocked(fetchCodexRateLimits).mockResolvedValueOnce(okProvider('codex', 20, Date.now())) + vi.mocked(fetchGeminiRateLimits).mockResolvedValueOnce(okProvider('gemini', 30, Date.now())) + vi.mocked(fetchOpenCodeGoRateLimits).mockResolvedValueOnce( + okProvider('opencode-go', 40, Date.now()) + ) + + await service.refresh() + + expect(fetchClaudeRateLimits).toHaveBeenCalledTimes(1) + expect(fetchCodexRateLimits).toHaveBeenCalledTimes(1) + expect(fetchGeminiRateLimits).toHaveBeenCalledTimes(1) + expect(fetchOpenCodeGoRateLimits).toHaveBeenCalledTimes(1) + expect(fetchOpenCodeGoRateLimits).toHaveBeenCalledWith('session=abc123', undefined) + + const state = service.getState() + expect(state.claude?.status).toBe('ok') + expect(state.claude?.session?.usedPercent).toBe(10) + expect(state.codex?.status).toBe('ok') + expect(state.codex?.session?.usedPercent).toBe(20) + expect(state.gemini?.status).toBe('ok') + expect(state.gemini?.session?.usedPercent).toBe(30) + expect(state.opencodeGo?.status).toBe('ok') + expect(state.opencodeGo?.session?.usedPercent).toBe(40) + }) + + it('preserves Gemini buckets through getState after fetch', async () => { + const service = new RateLimitService() + + const geminiWithBuckets: ProviderRateLimits = { + provider: 'gemini', + session: { usedPercent: 80, windowMinutes: 300, resetsAt: null, resetDescription: null }, + weekly: null, + buckets: [ + { + name: 'Pro', + usedPercent: 30, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + }, + { + name: 'Flash', + usedPercent: 80, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + } + ], + updatedAt: Date.now(), + error: null, + status: 'ok' + } + + vi.mocked(fetchClaudeRateLimits).mockResolvedValueOnce(okProvider('claude', 10, Date.now())) + vi.mocked(fetchCodexRateLimits).mockResolvedValueOnce(okProvider('codex', 20, Date.now())) + vi.mocked(fetchGeminiRateLimits).mockResolvedValueOnce(geminiWithBuckets) + vi.mocked(fetchOpenCodeGoRateLimits).mockResolvedValueOnce( + okProvider('opencode-go', 0, Date.now()) + ) + + await service.refresh() + + const state = service.getState() + expect(state.gemini?.buckets).toHaveLength(2) + expect(state.gemini?.buckets![0].name).toBe('Pro') + expect(state.gemini?.buckets![1].name).toBe('Flash') + // Why: session summary is derived from bucket data and must match the most constrained bucket. + expect(state.gemini?.session?.usedPercent).toBe(80) + }) + + it('isolates provider failures so one error does not block others', async () => { + const service = new RateLimitService() + service.setSettingsResolver(() => ({ opencodeSessionCookie: '', opencodeWorkspaceId: '' })) + + vi.mocked(fetchClaudeRateLimits).mockRejectedValueOnce(new Error('claude down')) + vi.mocked(fetchCodexRateLimits).mockResolvedValueOnce(okProvider('codex', 20, Date.now())) + vi.mocked(fetchGeminiRateLimits).mockRejectedValueOnce(new Error('gemini down')) + vi.mocked(fetchOpenCodeGoRateLimits).mockResolvedValueOnce( + okProvider('opencode-go', 40, Date.now()) + ) + + await service.refresh() + + const state = service.getState() + expect(state.claude?.status).toBe('error') + expect(state.claude?.error).toBe('claude down') + expect(state.codex?.status).toBe('ok') + expect(state.gemini?.status).toBe('error') + expect(state.gemini?.error).toBe('gemini down') + expect(state.opencodeGo?.status).toBe('ok') + }) + + it('discards stale data when a provider becomes unavailable', async () => { + const service = new RateLimitService() + let cookie = 'session=valid' + service.setSettingsResolver(() => ({ opencodeSessionCookie: cookie, opencodeWorkspaceId: '' })) + + // 1. Success fetch + vi.mocked(fetchClaudeRateLimits).mockResolvedValue(okProvider('claude', 10, Date.now())) + vi.mocked(fetchCodexRateLimits).mockResolvedValue(okProvider('codex', 20, Date.now())) + vi.mocked(fetchGeminiRateLimits).mockResolvedValue(okProvider('gemini', 30, Date.now())) + vi.mocked(fetchOpenCodeGoRateLimits).mockResolvedValue( + okProvider('opencode-go', 40, Date.now()) + ) + + await service.refresh() + expect(service.getState().opencodeGo?.session?.usedPercent).toBe(40) + + // 2. Clear cookie -> should become unavailable and LOSE the 40% data + cookie = '' + vi.mocked(fetchOpenCodeGoRateLimits).mockResolvedValue({ + provider: 'opencode-go', + session: null, + weekly: null, + monthly: null, + updatedAt: Date.now(), + error: 'Session cookie not configured', + status: 'unavailable' + }) + + await service.refresh() + const state = service.getState() + expect(state.opencodeGo?.status).toBe('unavailable') + expect(state.opencodeGo?.session).toBeNull() + expect(state.opencodeGo?.error).toBe('Session cookie not configured') + }) + + it('discards stale data when Workspace ID override is changed', async () => { + const service = new RateLimitService() + let workspaceId = 'wrk_A' + service.setSettingsResolver(() => ({ + opencodeSessionCookie: 'session=valid', + opencodeWorkspaceId: workspaceId + })) + + // 1. Success fetch for Workspace A + vi.mocked(fetchOpenCodeGoRateLimits).mockResolvedValue( + okProvider('opencode-go', 40, Date.now()) + ) + await service.refresh() + expect(service.getState().opencodeGo?.session?.usedPercent).toBe(40) + + // 2. Change Workspace ID to B -> old data from A should be discarded + workspaceId = 'wrk_B' + vi.mocked(fetchOpenCodeGoRateLimits).mockResolvedValue( + okProvider('opencode-go', 10, Date.now()) + ) + await service.refresh() + expect(service.getState().opencodeGo?.session?.usedPercent).toBe(10) + + // 3. Clear Workspace ID (automatic) but it fails -> should show error, NOT stale data from B + workspaceId = '' + vi.mocked(fetchOpenCodeGoRateLimits).mockResolvedValue({ + provider: 'opencode-go', + session: null, + weekly: null, + monthly: null, + updatedAt: Date.now(), + error: 'No workspace ID found', + status: 'error' + }) + await service.refresh() + const state = service.getState() + expect(state.opencodeGo?.status).toBe('error') + expect(state.opencodeGo?.session).toBeNull() + expect(state.opencodeGo?.error).toBe('No workspace ID found') + }) }) diff --git a/src/main/rate-limits/service.ts b/src/main/rate-limits/service.ts index 0810c41acd1..2b7a53dfb0c 100644 --- a/src/main/rate-limits/service.ts +++ b/src/main/rate-limits/service.ts @@ -6,6 +6,8 @@ import type { RateLimitState, ProviderRateLimits } from '../../shared/rate-limit import { fetchClaudeRateLimits } from './claude-fetcher' import { fetchCodexRateLimits } from './codex-fetcher' import type { ClaudeRuntimeAuthPreparation } from '../claude-accounts/runtime-auth-service' +import { fetchGeminiRateLimits } from './gemini-usage-fetcher' +import { fetchOpenCodeGoRateLimits } from './opencode-go-usage-fetcher' // Why: quota state does not need near-real-time polling, and a less aggressive // default reduces avoidable Claude /usage pressure. We intentionally use a @@ -15,7 +17,7 @@ const MIN_REFETCH_MS = 30 * 1000 // 30 seconds — debounce rapid refresh reques const STALE_THRESHOLD_MS = 10 * 60 * 1000 // 10 minutes — after this, stale data is dropped export class RateLimitService { - private state: RateLimitState = { claude: null, codex: null } + private state: RateLimitState = { claude: null, codex: null, gemini: null, opencodeGo: null } private pollInterval: number = DEFAULT_POLL_MS private timer: ReturnType | null = null private lastFetchAt = 0 @@ -28,8 +30,17 @@ export class RateLimitService { private fetchIdleResolvers: (() => void)[] = [] private codexFetchGeneration = 0 private claudeFetchGeneration = 0 + private opencodeFetchGeneration = 0 + private lastOpencodeConfigHash = '' private codexHomePathResolver: (() => string | null) | null = null private claudeAuthPreparationResolver: (() => Promise) | null = null + private settingsResolver: + | (() => { + opencodeSessionCookie: string + opencodeWorkspaceId: string + geminiCliOAuthEnabled?: boolean + }) + | null = null constructor() {} @@ -41,6 +52,15 @@ export class RateLimitService { this.claudeAuthPreparationResolver = resolver } + setSettingsResolver( + resolver: () => { + opencodeSessionCookie: string + opencodeWorkspaceId: string + geminiCliOAuthEnabled?: boolean + } + ): void { + this.settingsResolver = resolver + } attach(mainWindow: BrowserWindow): void { this.detachWindowListeners?.() this.mainWindow = mainWindow @@ -305,7 +325,7 @@ export class RateLimitService { private withFetchingStatus( current: ProviderRateLimits | null, - provider: 'claude' | 'codex' + provider: 'claude' | 'codex' | 'gemini' | 'opencode-go' ): ProviderRateLimits { if (!current) { return { @@ -328,38 +348,96 @@ export class RateLimitService { const codexProvenance = codexHomePath ? `managed:${codexHomePath}` : 'system' const codexGeneration = this.codexFetchGeneration const previousState = this.state + const settings = this.settingsResolver?.() + const cookie = settings?.opencodeSessionCookie ?? '' + const workspaceIdOverride = settings?.opencodeWorkspaceId ?? '' + const geminiCliOAuthEnabled = settings?.geminiCliOAuthEnabled ?? false - // Mark both providers as fetching while keeping previous data visible. + // Detect if configuration changed — if it did, we must discard any stale + // data because it belongs to a different session/workspace. + const currentConfigHash = `${cookie}|${workspaceIdOverride}` + const opencodeConfigChanged = currentConfigHash !== this.lastOpencodeConfigHash + if (opencodeConfigChanged) { + this.lastOpencodeConfigHash = currentConfigHash + this.opencodeFetchGeneration += 1 + } + const opencodeGeneration = this.opencodeFetchGeneration + + // Mark all providers as fetching while keeping previous data visible. // Codex account changes clear Codex separately before this method is // called, so ordinary refreshes still preserve the current values. this.updateState({ + ...previousState, claude: this.withFetchingStatus(previousState.claude, 'claude'), - codex: this.withFetchingStatus(previousState.codex, 'codex') + codex: this.withFetchingStatus(previousState.codex, 'codex'), + gemini: this.withFetchingStatus(previousState.gemini, 'gemini'), + opencodeGo: opencodeConfigChanged + ? this.withFetchingStatus(null, 'opencode-go') + : this.withFetchingStatus(previousState.opencodeGo, 'opencode-go') }) - const [claude, codex] = await Promise.all([ - fetchClaudeRateLimits({ authPreparation: claudeAuthPreparation }).catch( - (err): ProviderRateLimits => ({ - provider: 'claude', - session: null, - weekly: null, - updatedAt: Date.now(), - error: err instanceof Error ? err.message : 'Unknown error', - status: 'error' - }) - ), - fetchCodexRateLimits({ codexHomePath }).catch( - (err): ProviderRateLimits => ({ - provider: 'codex', - session: null, - weekly: null, - updatedAt: Date.now(), - error: err instanceof Error ? err.message : 'Unknown error', - status: 'error' - }) - ) + const [claudeResult, codexResult, geminiResult, opencodeGoResult] = await Promise.allSettled([ + fetchClaudeRateLimits({ authPreparation: claudeAuthPreparation }), + fetchCodexRateLimits({ codexHomePath }), + fetchGeminiRateLimits(geminiCliOAuthEnabled), + fetchOpenCodeGoRateLimits(cookie, workspaceIdOverride || undefined) ]) + const claude = + claudeResult.status === 'fulfilled' + ? claudeResult.value + : ({ + provider: 'claude', + session: null, + weekly: null, + updatedAt: Date.now(), + error: + claudeResult.reason instanceof Error ? claudeResult.reason.message : 'Unknown error', + status: 'error' + } satisfies ProviderRateLimits) + + const codex = + codexResult.status === 'fulfilled' + ? codexResult.value + : ({ + provider: 'codex', + session: null, + weekly: null, + updatedAt: Date.now(), + error: + codexResult.reason instanceof Error ? codexResult.reason.message : 'Unknown error', + status: 'error' + } satisfies ProviderRateLimits) + + const gemini = + geminiResult.status === 'fulfilled' + ? geminiResult.value + : ({ + provider: 'gemini', + session: null, + weekly: null, + updatedAt: Date.now(), + error: + geminiResult.reason instanceof Error ? geminiResult.reason.message : 'Unknown error', + status: 'error' + } satisfies ProviderRateLimits) + + const opencodeGo = + opencodeGoResult.status === 'fulfilled' + ? opencodeGoResult.value + : ({ + provider: 'opencode-go', + session: null, + weekly: null, + monthly: null, + updatedAt: Date.now(), + error: + opencodeGoResult.reason instanceof Error + ? opencodeGoResult.reason.message + : 'Unknown error', + status: 'error' + } satisfies ProviderRateLimits) + const latestCodexHomePath = this.codexHomePathResolver?.() ?? null const latestClaudeAuthPreparation = await this.claudeAuthPreparationResolver?.() const latestClaudeProvenance = latestClaudeAuthPreparation?.provenance ?? 'system' @@ -368,16 +446,26 @@ export class RateLimitService { codexGeneration === this.codexFetchGeneration && codexProvenance === latestCodexProvenance const shouldApplyClaude = claudeGeneration === this.claudeFetchGeneration && claudeProvenance === latestClaudeProvenance + const shouldApplyOpencode = opencodeGeneration === this.opencodeFetchGeneration // Why: account switches can race in-flight Codex fetches. Only apply a // Codex result if both the selected-account provenance and the request // generation still match, otherwise an old account could overwrite the // newly selected account's quota state. this.updateState({ + ...previousState, claude: shouldApplyClaude ? this.applyStalePolicy(claude, previousState.claude) : this.state.claude, - codex: shouldApplyCodex ? this.applyStalePolicy(codex, previousState.codex) : this.state.codex + codex: shouldApplyCodex + ? this.applyStalePolicy(codex, previousState.codex) + : this.state.codex, + gemini: this.applyStalePolicy(gemini, previousState.gemini), + opencodeGo: shouldApplyOpencode + ? opencodeConfigChanged + ? opencodeGo + : this.applyStalePolicy(opencodeGo, previousState.opencodeGo) + : this.state.opencodeGo }) this.lastFetchAt = Date.now() @@ -464,7 +552,18 @@ export class RateLimitService { return fresh } - const previousHasData = Boolean(previous?.session || previous?.weekly) + // Explicitly unavailable — user likely cleared a setting. Discard any stale + // data so the UI reflects that the provider is now disabled/unconfigured. + if (fresh.status === 'unavailable') { + return fresh + } + + const previousHasData = Boolean( + previous?.session || + previous?.weekly || + previous?.monthly || + (previous?.buckets && previous.buckets.length > 0) + ) // No previous data to fall back on if (!previous || !previousHasData) { diff --git a/src/renderer/src/components/settings/GeneralPane.tsx b/src/renderer/src/components/settings/GeneralPane.tsx index 2fa2a5ec817..c258da3a8d1 100644 --- a/src/renderer/src/components/settings/GeneralPane.tsx +++ b/src/renderer/src/components/settings/GeneralPane.tsx @@ -28,6 +28,8 @@ import { GENERAL_CACHE_TIMER_SEARCH_ENTRIES, GENERAL_CLI_SEARCH_ENTRIES, GENERAL_EDITOR_SEARCH_ENTRIES, + GENERAL_GEMINI_SEARCH_ENTRIES, + GENERAL_OPENCODE_SEARCH_ENTRIES, GENERAL_PANE_SEARCH_ENTRIES, GENERAL_SUPPORT_SEARCH_ENTRIES, GENERAL_UPDATE_SEARCH_ENTRIES, @@ -1014,6 +1016,131 @@ export function GeneralPane({ settings, updateSettings }: GeneralPaneProps): Rea ) : null, + matchesSettingsSearch(searchQuery, GENERAL_GEMINI_SEARCH_ENTRIES) ? ( +
+
+

Gemini

+

Configure Gemini provider settings.

+
+ + +
+ +

+ Extracts OAuth credentials from your local Gemini CLI installation to authenticate + with Google. This uses credentials issued to the Gemini CLI app, not Orca. May break + if Google updates the CLI. Use at your own risk. +

+
+ + +
+ ) : null, + matchesSettingsSearch(searchQuery, GENERAL_OPENCODE_SEARCH_ENTRIES) ? ( +
+
+

OpenCode Go

+

Configure OpenCode Go provider settings.

+
+ + + +
+ updateSettings({ opencodeSessionCookie: e.target.value })} + placeholder="Fe26.2**… token or auth=Fe26.2**… header" + spellCheck={false} + className="flex-1 text-xs" + /> + {settings.opencodeSessionCookie && ( + + )} +
+

+ Paste either the raw token value (e.g. Fe26.2**…) or + the full cookie header (e.g. auth=Fe26.2**…). Find it + in your browser's DevTools → Network → any opencode.ai request → Cookie header. +

+ + + + +
+ updateSettings({ opencodeWorkspaceId: e.target.value })} + placeholder="wrk_… (leave blank for automatic lookup)" + spellCheck={false} + className="flex-1 text-xs" + /> + {settings.opencodeWorkspaceId && ( + + )} +
+

+ Find this in the URL after logging into opencode.ai (e.g.{' '} + opencode.ai/workspace/wrk_…/go). +

+ +
+ ) : null, matchesSettingsSearch(searchQuery, GENERAL_UPDATE_SEARCH_ENTRIES) ? (
diff --git a/src/renderer/src/components/settings/general-search.ts b/src/renderer/src/components/settings/general-search.ts index 06a03927d33..45e58db24d3 100644 --- a/src/renderer/src/components/settings/general-search.ts +++ b/src/renderer/src/components/settings/general-search.ts @@ -86,6 +86,28 @@ export const GENERAL_CODEX_ACCOUNTS_SEARCH_ENTRIES: SettingsSearchEntry[] = [ } ] +export const GENERAL_GEMINI_SEARCH_ENTRIES: SettingsSearchEntry[] = [ + { + title: 'Use Gemini CLI credentials', + description: + 'Extracts OAuth credentials from your local Gemini CLI installation to authenticate with Google.', + keywords: ['gemini', 'cli', 'oauth', 'credentials', 'experimental', 'rate limit', 'status bar'] + } +] + +export const GENERAL_OPENCODE_SEARCH_ENTRIES: SettingsSearchEntry[] = [ + { + title: 'OpenCode Go Session Cookie', + description: 'Paste your opencode.ai session cookie for rate limit fetching.', + keywords: ['opencode', 'cookie', 'session', 'rate limit', 'status bar'] + }, + { + title: 'OpenCode Go Workspace ID', + description: 'Optional workspace ID override if the automatic lookup fails.', + keywords: ['opencode', 'workspace', 'id', 'wrk', 'rate limit', 'status bar'] + } +] + export const GENERAL_AGENT_SEARCH_ENTRIES: SettingsSearchEntry[] = [ { title: 'Default Agent', @@ -119,6 +141,8 @@ export const GENERAL_PANE_SEARCH_ENTRIES: SettingsSearchEntry[] = [ ...GENERAL_CACHE_TIMER_SEARCH_ENTRIES, ...GENERAL_CLAUDE_ACCOUNTS_SEARCH_ENTRIES, ...GENERAL_CODEX_ACCOUNTS_SEARCH_ENTRIES, + ...GENERAL_GEMINI_SEARCH_ENTRIES, + ...GENERAL_OPENCODE_SEARCH_ENTRIES, ...GENERAL_UPDATE_SEARCH_ENTRIES, ...GENERAL_SUPPORT_SEARCH_ENTRIES ] diff --git a/src/renderer/src/components/status-bar/StatusBar.tsx b/src/renderer/src/components/status-bar/StatusBar.tsx index 4d652f12e68..31e6096854e 100644 --- a/src/renderer/src/components/status-bar/StatusBar.tsx +++ b/src/renderer/src/components/status-bar/StatusBar.tsx @@ -28,7 +28,8 @@ import type { } from '../../../../shared/types' import type { ProviderRateLimits, RateLimitWindow } from '../../../../shared/rate-limit-types' import { ProviderIcon, ProviderPanel } from './tooltip' -import { ClaudeIcon, OpenAIIcon } from './icons' +import { ClaudeIcon, GeminiIcon, OpenAIIcon, OpenCodeGoIcon } from './icons' +import { formatWindowLabel } from '@/lib/window-label-formatter' import { markLiveCodexSessionsForRestart } from '@/lib/codex-session-restart' import { SshStatusSegment } from './SshStatusSegment' import { SessionsStatusSegment } from './SessionsStatusSegment' @@ -221,6 +222,10 @@ function WindowLabel({ w, label }: { w: RateLimitWindow; label: string }): React // Provider segment // --------------------------------------------------------------------------- +// Why: only Flash and the latest Pro are shown in the status bar — +// the rest (Flash Lite, experimental) are secondary and would clutter the bar. +const STATUS_BAR_BUCKET_NAMES = new Set(['Flash', 'Pro', '1.5 Pro']) + function ProviderSegment({ p, compact @@ -273,13 +278,40 @@ function ProviderSegment({ // Has data (ok, fetching with stale data, or error with stale data) const isStale = p.status === 'error' + + if (p.buckets && p.buckets.length > 0) { + const visibleBuckets = p.buckets.filter((b) => STATUS_BAR_BUCKET_NAMES.has(b.name)) + return ( + + + {visibleBuckets.map((bucket, i) => { + const left = Math.max(0, Math.round(100 - bucket.usedPercent)) + return ( + + {i > 0 && ·} + + {bucket.name} {left}% + + + ) + })} + {visibleBuckets.length === 0 && p.session && ( + + )} + {isStale && } + + ) + } + return ( {p.session && !compact && } - {p.session && } + {p.session && ( + + )} {p.session && p.weekly && ·} - {p.weekly && } + {p.weekly && } {isStale && } ) @@ -525,7 +557,13 @@ function ProviderDetailsMenu({ className={`inline-block h-2 w-2 rounded-full ${provider.session || provider.weekly ? 'bg-muted-foreground/60' : 'bg-muted-foreground/30'}`} /> - {provider.provider === 'claude' ? 'C' : 'X'} + {provider.provider === 'claude' + ? 'C' + : provider.provider === 'gemini' + ? 'G' + : provider.provider === 'opencode-go' + ? 'O' + : 'X'} ) : ( @@ -608,7 +646,7 @@ function StatusBarInner(): React.JSX.Element | null { return null } - const { claude, codex } = rateLimits + const { claude, codex, gemini, opencodeGo } = rateLimits // Why: hiding `unavailable` providers makes the status bar appear to lose a // provider at random after refreshes or wake/resume. Keeping the slot visible @@ -616,11 +654,20 @@ function StatusBarInner(): React.JSX.Element | null { // configured but currently unavailable. const showClaude = claude && statusBarItems.includes('claude') const showCodex = codex && statusBarItems.includes('codex') + // Why: hide only when the state hasn't loaded yet (null), not when unavailable. + // Gemini shows if credentials exist; OpenCode Go shows always so users can see + // the provider and know to configure the cookie in Settings. + const showGemini = gemini !== null && statusBarItems.includes('gemini') + const showOpencodeGo = opencodeGo !== null && statusBarItems.includes('opencode-go') const showSsh = statusBarItems.includes('ssh') const showSessions = statusBarItems.includes('sessions') const showMemory = statusBarItems.includes('memory') - const anyVisible = showClaude || showCodex - const anyFetching = claude?.status === 'fetching' || codex?.status === 'fetching' + const anyVisible = showClaude || showCodex || showGemini || showOpencodeGo || showMemory + const anyFetching = + claude?.status === 'fetching' || + codex?.status === 'fetching' || + gemini?.status === 'fetching' || + opencodeGo?.status === 'fetching' const compact = containerWidth < 900 const iconOnly = containerWidth < 500 @@ -646,6 +693,22 @@ function StatusBarInner(): React.JSX.Element | null {
{showClaude && } {showCodex && } + {showGemini && ( + + )} + {showOpencodeGo && ( + + )} {anyVisible && ( @@ -701,6 +764,20 @@ function StatusBarInner(): React.JSX.Element | null { Codex Usage + toggleStatusBarItem('gemini')} + > + + Gemini Usage + + toggleStatusBarItem('opencode-go')} + > + + OpenCode Go Usage + toggleStatusBarItem('ssh')} diff --git a/src/renderer/src/components/status-bar/icons.tsx b/src/renderer/src/components/status-bar/icons.tsx index 3f9bca4e80d..8e03bcd45a0 100644 --- a/src/renderer/src/components/status-bar/icons.tsx +++ b/src/renderer/src/components/status-bar/icons.tsx @@ -1,3 +1,5 @@ +import React from 'react' + export function OpenAIIcon({ size = 14 }: { size?: number }): React.JSX.Element { return ( `gi-${++_geminiIconCount}`) + const star = + 'M32.447 0c.68 0 1.273.465 1.439 1.125a38.904 38.904 0 001.999 5.905c2.152 5 5.105 9.376 8.854 13.125 3.751 3.75 8.126 6.703 13.125 8.855a38.98 38.98 0 005.906 1.999c.66.166 1.124.758 1.124 1.438 0 .68-.464 1.273-1.125 1.439a38.902 38.902 0 00-5.905 1.999c-5 2.152-9.375 5.105-13.125 8.854-3.749 3.751-6.702 8.126-8.854 13.125a38.973 38.973 0 00-2 5.906 1.485 1.485 0 01-1.438 1.124c-.68 0-1.272-.464-1.438-1.125a38.913 38.913 0 00-2-5.905c-2.151-5-5.103-9.375-8.854-13.125-3.75-3.749-8.125-6.702-13.125-8.854a38.973 38.973 0 00-5.905-2A1.485 1.485 0 010 32.448c0-.68.465-1.272 1.125-1.438a38.903 38.903 0 005.905-2c5-2.151 9.376-5.104 13.125-8.854 3.75-3.749 6.703-8.125 8.855-13.125a38.972 38.972 0 001.999-5.905A1.485 1.485 0 0132.447 0z' + return ( + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + ) +} + +export function OpenCodeGoIcon({ size = 14 }: { size?: number }): React.JSX.Element { + return ( + + + + + + + + + + + + + + + + + ) +} + export function ClaudeIcon({ size = 14 }: { size?: number }): React.JSX.Element { return ( diff --git a/src/renderer/src/components/status-bar/tooltip.test.ts b/src/renderer/src/components/status-bar/tooltip.test.ts new file mode 100644 index 00000000000..02c0a68fbeb --- /dev/null +++ b/src/renderer/src/components/status-bar/tooltip.test.ts @@ -0,0 +1,135 @@ +import { describe, expect, it } from 'vitest' +import type { ProviderRateLimits } from '../../../../shared/rate-limit-types' +import { getWindowSections } from './tooltip' + +describe('getWindowSections', () => { + it('returns buckets as sections when present', () => { + const p: ProviderRateLimits = { + provider: 'gemini', + session: { usedPercent: 80, windowMinutes: 300, resetsAt: null, resetDescription: null }, + weekly: null, + buckets: [ + { + name: 'Pro', + usedPercent: 30, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + }, + { + name: 'Flash', + usedPercent: 80, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + } + ], + updatedAt: Date.now(), + error: null, + status: 'ok' + } + const sections = getWindowSections(p) + expect(sections).toEqual([ + { label: 'Pro', window: p.buckets![0] }, + { label: 'Flash', window: p.buckets![1] }, + { label: 'Weekly', window: null } + ]) + }) + + it('returns session and weekly when buckets are absent', () => { + const p: ProviderRateLimits = { + provider: 'claude', + session: { usedPercent: 40, windowMinutes: 300, resetsAt: null, resetDescription: null }, + weekly: { usedPercent: 20, windowMinutes: 10080, resetsAt: null, resetDescription: null }, + updatedAt: Date.now(), + error: null, + status: 'ok' + } + const sections = getWindowSections(p) + expect(sections).toEqual([ + { label: 'Session', window: p.session }, + { label: 'Weekly', window: p.weekly } + ]) + }) + + it('returns session and weekly for empty buckets array', () => { + const p: ProviderRateLimits = { + provider: 'gemini', + session: { usedPercent: 50, windowMinutes: 300, resetsAt: null, resetDescription: null }, + weekly: null, + buckets: [], + updatedAt: Date.now(), + error: null, + status: 'ok' + } + const sections = getWindowSections(p) + expect(sections).toEqual([ + { label: 'Session', window: p.session }, + { label: 'Weekly', window: null } + ]) + }) + + it('does not expose bucket names via session window in compact rendering path', () => { + // Why: ProviderSegment (compact mode) reads only p.session — never p.buckets. + // This test locks the contract: getWindowSections returns buckets for detail + // views, while the plain session value remains independently available for + // compact rendering without bucket names bleeding through. + const p: ProviderRateLimits = { + provider: 'gemini', + session: { usedPercent: 80, windowMinutes: 300, resetsAt: null, resetDescription: null }, + weekly: null, + buckets: [ + { + name: 'Pro', + usedPercent: 30, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + }, + { + name: 'Flash', + usedPercent: 80, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + } + ], + updatedAt: Date.now(), + error: null, + status: 'ok' + } + // Compact path uses p.session directly — independent of getWindowSections. + expect(p.session?.usedPercent).toBe(80) + // getWindowSections (detail path) returns bucket rows, not session label. + const sections = getWindowSections(p) + const labels = sections.map((s) => s.label) + expect(labels).toContain('Pro') + expect(labels).toContain('Flash') + expect(labels).not.toContain('Session') + }) + + it('preserves reset metadata inside bucket windows', () => { + const p: ProviderRateLimits = { + provider: 'gemini', + session: null, + weekly: null, + buckets: [ + { + name: 'Pro', + usedPercent: 45, + windowMinutes: 300, + resetsAt: 18000000, + resetDescription: '5:00 PM' + } + ], + updatedAt: Date.now(), + error: null, + status: 'ok' + } + const sections = getWindowSections(p) + expect(sections).toHaveLength(2) + expect(sections[0].label).toBe('Pro') + expect(sections[0].window!.resetsAt).toBe(18000000) + expect(sections[0].window!.resetDescription).toBe('5:00 PM') + }) +}) diff --git a/src/renderer/src/components/status-bar/tooltip.tsx b/src/renderer/src/components/status-bar/tooltip.tsx index d1b1305bc6a..b7cf6f5e242 100644 --- a/src/renderer/src/components/status-bar/tooltip.tsx +++ b/src/renderer/src/components/status-bar/tooltip.tsx @@ -1,5 +1,5 @@ import type { ProviderRateLimits, RateLimitWindow } from '../../../../shared/rate-limit-types' -import { ClaudeIcon, OpenAIIcon } from './icons' +import { ClaudeIcon, GeminiIcon, OpenAIIcon, OpenCodeGoIcon } from './icons' // --------------------------------------------------------------------------- // Formatting helpers @@ -44,14 +44,23 @@ export function ProviderIcon({ provider }: { provider: string }): React.JSX.Elem if (provider === 'codex') { return } + if (provider === 'gemini') { + return + } + if (provider === 'opencode-go') { + return + } return } function ErrorMessage({ message, + stale = false, inverted = false }: { message: string + /** When true, prior data is still visible — show a softer "refresh failed" label. */ + stale?: boolean inverted?: boolean }): React.JSX.Element { const labelClass = inverted ? 'text-background/80' : 'text-foreground/85' @@ -59,12 +68,35 @@ function ErrorMessage({ return (
-
Usage unavailable
+
+ {stale ? 'Refresh failed — showing cached data' : 'Usage unavailable'} +
{message}
) } +// --------------------------------------------------------------------------- +// Window section derivation +// --------------------------------------------------------------------------- + +export function getWindowSections( + p: ProviderRateLimits +): { label: string; window: RateLimitWindow | null }[] { + if (p.buckets?.length) { + const bucketSections = p.buckets.map((b) => ({ label: b.name, window: b as RateLimitWindow })) + return [...bucketSections, { label: 'Weekly', window: p.weekly }] + } + const sections: { label: string; window: RateLimitWindow | null }[] = [ + { label: 'Session', window: p.session }, + { label: 'Weekly', window: p.weekly } + ] + if (p.monthly !== undefined && p.monthly !== null) { + sections.push({ label: 'Monthly', window: p.monthly }) + } + return sections +} + // --------------------------------------------------------------------------- // Tooltip — progress bar section for a single window // --------------------------------------------------------------------------- @@ -125,7 +157,16 @@ export function ProviderTooltip({ p }: { p: ProviderRateLimits | null }): React. return No data available } - const name = p.provider === 'claude' ? 'Claude' : 'Codex' + const name = + p.provider === 'claude' + ? 'Claude' + : p.provider === 'codex' + ? 'Codex' + : p.provider === 'gemini' + ? 'Gemini' + : p.provider === 'opencode-go' + ? 'OpenCode Go' + : p.provider if (p.status === 'unavailable') { return ( @@ -134,12 +175,12 @@ export function ProviderTooltip({ p }: { p: ProviderRateLimits | null }): React. {name}
-
{p.error ?? 'CLI not found'}
+
{p.error ?? 'Unavailable'}
) } - if (p.status === 'error' && !p.session && !p.weekly) { + if (p.status === 'error' && !p.session && !p.weekly && !p.monthly) { return (
@@ -167,14 +208,14 @@ export function ProviderTooltip({ p }: { p: ProviderRateLimits | null }): React. {/* Divider */}
- {/* Session window */} - + {getWindowSections(p).map((s) => ( + + ))} - {/* Weekly window */} - - - {/* Stale data warning */} - {p.error ? : null} + {/* Stale data warning — softer label when prior data is still shown */} + {p.error ? ( + + ) : null}
) } @@ -198,7 +239,16 @@ export function ProviderPanel({ return No data available } - const name = p.provider === 'claude' ? 'Claude' : 'Codex' + const name = + p.provider === 'claude' + ? 'Claude' + : p.provider === 'codex' + ? 'Codex' + : p.provider === 'gemini' + ? 'Gemini' + : p.provider === 'opencode-go' + ? 'OpenCode Go' + : p.provider if (p.status === 'unavailable') { return ( @@ -207,12 +257,12 @@ export function ProviderPanel({ {name}
-
{p.error ?? 'CLI not found'}
+
{p.error ?? 'Unavailable'}
) } - if (p.status === 'error' && !p.session && !p.weekly) { + if (p.status === 'error' && !p.session && !p.weekly && !p.monthly) { return (
@@ -270,10 +320,17 @@ export function ProviderPanel({
- - + {getWindowSections(p).map((s) => ( + + ))} - {p.error ? : null} + {p.error ? ( + + ) : null}
) } diff --git a/src/renderer/src/lib/window-label-formatter.test.ts b/src/renderer/src/lib/window-label-formatter.test.ts new file mode 100644 index 00000000000..730afbbd5ad --- /dev/null +++ b/src/renderer/src/lib/window-label-formatter.test.ts @@ -0,0 +1,40 @@ +import { describe, expect, it } from 'vitest' +import { formatWindowLabel } from './window-label-formatter' + +describe('formatWindowLabel', () => { + it('returns "5h" for 300 minutes', () => { + expect(formatWindowLabel(300)).toBe('5h') + }) + + it('returns "1h" for 60 minutes', () => { + expect(formatWindowLabel(60)).toBe('1h') + }) + + it('returns "wk" for 10080 minutes (7 days)', () => { + expect(formatWindowLabel(10080)).toBe('wk') + }) + + it('returns "1d" for 1440 minutes (1 day)', () => { + expect(formatWindowLabel(1440)).toBe('1d') + }) + + it('returns "2h" for 120 minutes', () => { + expect(formatWindowLabel(120)).toBe('2h') + }) + + it('returns "45m" for 45 minutes', () => { + expect(formatWindowLabel(45)).toBe('45m') + }) + + it('returns "2wk" for 20160 minutes (14 days)', () => { + expect(formatWindowLabel(20160)).toBe('2wk') + }) + + it('returns "30m" for 30 minutes', () => { + expect(formatWindowLabel(30)).toBe('30m') + }) + + it('returns "3d" for 4320 minutes (3 days)', () => { + expect(formatWindowLabel(4320)).toBe('3d') + }) +}) diff --git a/src/renderer/src/lib/window-label-formatter.ts b/src/renderer/src/lib/window-label-formatter.ts new file mode 100644 index 00000000000..1fbbcdb54f5 --- /dev/null +++ b/src/renderer/src/lib/window-label-formatter.ts @@ -0,0 +1,30 @@ +/** + * Returns a short human-readable label for a usage window duration. + * + * Why: 10080 minutes (7 days) is hard-coded as "wk" for backward + * compatibility with the original StatusBar implementation. + */ +export function formatWindowLabel(windowMinutes: number): string { + if (windowMinutes === 10080) { + return 'wk' + } + if (windowMinutes === 300) { + return '5h' + } + if (windowMinutes === 60) { + return '1h' + } + if (windowMinutes < 60) { + return `${windowMinutes}m` + } + if (windowMinutes % (60 * 24 * 7) === 0) { + return `${windowMinutes / (60 * 24 * 7)}wk` + } + if (windowMinutes % (60 * 24) === 0) { + return `${windowMinutes / (60 * 24)}d` + } + if (windowMinutes % 60 === 0) { + return `${windowMinutes / 60}h` + } + return `${windowMinutes}m` +} diff --git a/src/renderer/src/store/slices/rate-limits.ts b/src/renderer/src/store/slices/rate-limits.ts index f393ac92b7d..d3be60d9085 100644 --- a/src/renderer/src/store/slices/rate-limits.ts +++ b/src/renderer/src/store/slices/rate-limits.ts @@ -10,7 +10,7 @@ export type RateLimitSlice = { } export const createRateLimitSlice: StateCreator = (set) => ({ - rateLimits: { claude: null, codex: null }, + rateLimits: { claude: null, codex: null, gemini: null, opencodeGo: null }, fetchRateLimits: async () => { try { diff --git a/src/shared/constants.ts b/src/shared/constants.ts index 5b580ea9cff..a0ead7d4a5e 100644 --- a/src/shared/constants.ts +++ b/src/shared/constants.ts @@ -73,6 +73,8 @@ export const DEFAULT_WORKTREE_CARD_PROPERTIES: WorktreeCardProperty[] = [ export const DEFAULT_STATUS_BAR_ITEMS: StatusBarItem[] = [ 'claude', 'codex', + 'gemini', + 'opencode-go', 'ssh', 'sessions', 'memory' @@ -170,6 +172,9 @@ export function getDefaultSettings(homedir: string): GlobalSettings { defaultTaskSource: 'github', defaultRepoSelection: null, defaultLinearTeamSelection: null, + opencodeSessionCookie: '', + opencodeWorkspaceId: '', + geminiCliOAuthEnabled: false, agentCmdOverrides: {}, // Why: 'auto' runs a layout-aware probe at boot (see // src/renderer/src/lib/keyboard-layout/*) that picks 'true' for US and diff --git a/src/shared/rate-limit-types.ts b/src/shared/rate-limit-types.ts index 53f3b91ad82..7931cea4b21 100644 --- a/src/shared/rate-limit-types.ts +++ b/src/shared/rate-limit-types.ts @@ -11,12 +11,20 @@ export type RateLimitWindow = { export type ProviderRateLimitStatus = 'idle' | 'fetching' | 'ok' | 'error' | 'unavailable' +export type RateLimitBucket = RateLimitWindow & { + name: string +} + export type ProviderRateLimits = { - provider: 'claude' | 'codex' + provider: 'claude' | 'codex' | 'gemini' | 'opencode-go' /** 5-hour session window, null if not available. */ session: RateLimitWindow | null /** 7-day weekly window, null if not available. */ weekly: RateLimitWindow | null + /** 30-day monthly window (OpenCode Go only), null if not available. */ + monthly?: RateLimitWindow | null + /** Named per-model buckets (Gemini only). */ + buckets?: RateLimitBucket[] /** Unix ms timestamp of the last successful data update. */ updatedAt: number /** Human-readable error message, null when status is 'ok'. */ @@ -27,4 +35,6 @@ export type ProviderRateLimits = { export type RateLimitState = { claude: ProviderRateLimits | null codex: ProviderRateLimits | null + gemini: ProviderRateLimits | null + opencodeGo: ProviderRateLimits | null } diff --git a/src/shared/types.ts b/src/shared/types.ts index d5ff5d01271..96bef5728a6 100644 --- a/src/shared/types.ts +++ b/src/shared/types.ts @@ -984,6 +984,14 @@ export type GlobalSettings = { * Same nullable-array pattern as `defaultRepoSelection`: `null` = sticky-all, * `string[]` = frozen subset of team IDs. */ defaultLinearTeamSelection: string[] | null + /** Session cookie for OpenCode Go rate-limit fetching. Stored encrypted. */ + opencodeSessionCookie: string + /** Optional workspace ID override for OpenCode Go. When set, skips the + * workspaces lookup and fetches usage directly for this workspace. */ + opencodeWorkspaceId: string + /** Whether to extract OAuth credentials from the local Gemini CLI installation + * for rate-limit fetching. Disabled by default for explicit opt-in. */ + geminiCliOAuthEnabled: boolean /** Per-agent CLI command overrides. A missing key means use the catalog default binary name. */ agentCmdOverrides: Partial> /** Why: macOS terminals must choose between letting Option compose layout @@ -1054,8 +1062,14 @@ export type NotificationDispatchResult = { export type WorktreeCardProperty = 'status' | 'unread' | 'ci' | 'issue' | 'pr' | 'comment' -export type StatusBarItem = 'claude' | 'codex' | 'ssh' | 'sessions' | 'memory' - +export type StatusBarItem = + | 'claude' + | 'codex' + | 'gemini' + | 'opencode-go' + | 'ssh' + | 'sessions' + | 'memory' export type PersistedUIState = { lastActiveRepoId: string | null lastActiveWorktreeId: string | null