From 38ac287e4d64495f5e353d68aac07ebc53c6b636 Mon Sep 17 00:00:00 2001 From: iceteaSA <171169159+iceteaSA@users.noreply.github.com> Date: Fri, 24 Jul 2026 13:10:04 +0200 Subject: [PATCH 01/25] fix(commands): wire account dialog OAuth flow (cherry picked from commit 99fc9fbfceacabc4e06834cf8121ded69ed8ed02) --- .../src/plugin/account-command-oauth.test.ts | 112 ++++++++ .../src/plugin/account-command-oauth.ts | 142 ++++++++++ packages/opencode/src/plugin/command-data.ts | 20 +- packages/opencode/src/plugin/commands.test.ts | 144 ++++++++++ packages/opencode/src/plugin/commands.ts | 115 ++++++-- packages/opencode/src/plugin/index.ts | 12 + .../src/tui-compiled/tui/command-dialogs.tsx | 251 +++++++++++++----- .../opencode/src/tui/command-dialogs.test.tsx | 76 +++++- packages/opencode/src/tui/command-dialogs.tsx | 132 ++++++++- 9 files changed, 909 insertions(+), 95 deletions(-) create mode 100644 packages/opencode/src/plugin/account-command-oauth.test.ts create mode 100644 packages/opencode/src/plugin/account-command-oauth.ts diff --git a/packages/opencode/src/plugin/account-command-oauth.test.ts b/packages/opencode/src/plugin/account-command-oauth.test.ts new file mode 100644 index 0000000..7df18ce --- /dev/null +++ b/packages/opencode/src/plugin/account-command-oauth.test.ts @@ -0,0 +1,112 @@ +import { describe, expect, it, mock } from 'bun:test' +import { createAccountCommandOAuthService } from './account-command-oauth' + +const initialRows = [ + { + id: 'acct-0', + index: 0, + label: 'Primary account', + enabled: true, + current: true, + quota: [], + }, +] + +describe('createAccountCommandOAuthService', () => { + it('returns an OAuth URL and stores pending state for the session', async () => { + const authorize = mock(async () => ({ + url: 'https://accounts.google.test/authorize?state=signed-state&redirect_uri=http%3A%2F%2Flocalhost%3A51121%2Foauth-callback', + verifier: 'pkce-verifier', + projectId: '', + })) + const exchange = mock(async () => ({ + type: 'failed' as const, + error: 'unused', + })) + const service = createAccountCommandOAuthService({ + authorize, + exchange, + persist: mock(async () => {}), + listAccounts: mock(async () => initialRows), + }) + + const started = await service.start('session-1') + await service.finish('session-1', 'callback-code') + + expect(authorize).toHaveBeenCalledTimes(1) + expect(started.url).toContain('accounts.google.test') + expect(started.accounts).toEqual(initialRows) + expect(exchange).toHaveBeenCalledWith('callback-code', 'signed-state') + }) + + it('exchanges, persists, and returns updated label-only account rows', async () => { + const persisted = mock(async () => {}) + const updatedRows = [ + ...initialRows, + { + id: 'acct-1', + index: 1, + label: 'Work account', + enabled: true, + current: false, + quota: [], + }, + ] + const service = createAccountCommandOAuthService({ + authorize: async () => ({ + url: 'https://accounts.google.test/authorize?state=signed-state&redirect_uri=http%3A%2F%2Flocalhost%3A51121%2Foauth-callback', + verifier: 'pkce-verifier', + projectId: '', + }), + exchange: mock(async () => ({ + type: 'success' as const, + refresh: 'refresh-token|project', + access: 'access-token', + expires: 123, + email: 'private@example.test', + label: 'OAuth display name', + projectId: 'project', + })), + persist: persisted, + listAccounts: mock(async () => updatedRows), + }) + + await service.start('session-1') + const finished = await service.finish( + 'session-1', + 'callback-code', + 'Work account', + ) + + expect(persisted).toHaveBeenCalledWith( + expect.objectContaining({ label: 'Work account' }), + ) + expect(finished).toEqual({ + text: 'OAuth account added.', + accounts: updatedRows, + }) + }) + + it('returns a friendly error when the session has no pending OAuth flow', async () => { + const service = createAccountCommandOAuthService({ + authorize: async () => ({ + url: 'https://accounts.google.test/authorize?state=signed-state&redirect_uri=http%3A%2F%2Flocalhost%3A51121%2Foauth-callback', + verifier: 'pkce-verifier', + projectId: '', + }), + exchange: mock(async () => ({ + type: 'failed' as const, + error: 'unused', + })), + persist: mock(async () => {}), + listAccounts: mock(async () => initialRows), + }) + + const result = await service.finish('missing-session', 'callback-code') + + expect(result).toEqual({ + text: 'OAuth session expired. Please start again.', + accounts: initialRows, + }) + }) +}) diff --git a/packages/opencode/src/plugin/account-command-oauth.ts b/packages/opencode/src/plugin/account-command-oauth.ts new file mode 100644 index 0000000..5a9feef --- /dev/null +++ b/packages/opencode/src/plugin/account-command-oauth.ts @@ -0,0 +1,142 @@ +import type { + AntigravityAuthorization, + AntigravityTokenExchangeResult, +} from '../antigravity/oauth' +import type { CommandAccountRow } from './command-data' +import { parseOAuthCallbackInput } from './oauth-methods' + +type OAuthSuccess = Extract + +interface OAuthPendingEntry { + state: string + verifier: string + redirectUri: string + createdAt: number +} + +export interface AccountCommandOAuthServiceOptions { + authorize: () => Promise + exchange: ( + code: string, + state: string, + ) => Promise + persist: (result: OAuthSuccess) => Promise + listAccounts: () => Promise + now?: () => number +} + +export interface AccountCommandOAuthService { + start(sessionId: string): Promise<{ + url: string + accounts: CommandAccountRow[] + }> + finish( + sessionId: string, + callbackInput: string, + label?: string, + ): Promise<{ + text: string + accounts: CommandAccountRow[] + }> + dispose(): void +} + +const OAUTH_PENDING_TTL_MS = 10 * 60 * 1_000 +const OAUTH_PENDING_CAP = 50 + +export function createAccountCommandOAuthService( + options: AccountCommandOAuthServiceOptions, +): AccountCommandOAuthService { + const pendingBySession = new Map() + const now = options.now ?? (() => Date.now()) + + const cleanupExpired = (): void => { + const current = now() + for (const [sessionId, entry] of pendingBySession) { + if (current - entry.createdAt > OAUTH_PENDING_TTL_MS) { + pendingBySession.delete(sessionId) + } + } + } + + const takePending = (sessionId: string): OAuthPendingEntry | undefined => { + cleanupExpired() + const entry = pendingBySession.get(sessionId) + if (!entry || now() - entry.createdAt > OAUTH_PENDING_TTL_MS) { + pendingBySession.delete(sessionId) + return undefined + } + return entry + } + + return { + async start(sessionId) { + const authorization = await options.authorize() + const url = new URL(authorization.url) + const state = url.searchParams.get('state') + const redirectUri = url.searchParams.get('redirect_uri') + if (!state || !redirectUri) { + throw new Error('OAuth authorization URL is missing required state') + } + + cleanupExpired() + if (pendingBySession.size >= OAUTH_PENDING_CAP) { + const oldest = [...pendingBySession.entries()].reduce( + (previous, current) => + current[1].createdAt < previous[1].createdAt ? current : previous, + ) + pendingBySession.delete(oldest[0]) + } + pendingBySession.set(sessionId, { + state, + verifier: authorization.verifier, + redirectUri, + createdAt: now(), + }) + return { url: authorization.url, accounts: await options.listAccounts() } + }, + + async finish(sessionId, callbackInput, label) { + const pending = takePending(sessionId) + if (!pending) { + return { + text: 'OAuth session expired. Please start again.', + accounts: await options.listAccounts(), + } + } + + try { + const callback = parseOAuthCallbackInput(callbackInput, pending.state) + if ('error' in callback) { + return { + text: `OAuth authentication failed: ${callback.error}`, + accounts: await options.listAccounts(), + } + } + const result = await options.exchange(callback.code, callback.state) + if (result.type === 'failed') { + return { + text: 'OAuth authentication failed. Please check the code and try again.', + accounts: await options.listAccounts(), + } + } + await options.persist({ ...result, label: label || result.label }) + return { + text: 'OAuth account added.', + accounts: await options.listAccounts(), + } + } catch { + return { + text: 'OAuth exchange failed due to a network error. Please try again.', + accounts: await options.listAccounts(), + } + } finally { + pendingBySession.delete(sessionId) + } + }, + + dispose() { + pendingBySession.clear() + }, + } +} diff --git a/packages/opencode/src/plugin/command-data.ts b/packages/opencode/src/plugin/command-data.ts index 0c85a5d..e3c933f 100644 --- a/packages/opencode/src/plugin/command-data.ts +++ b/packages/opencode/src/plugin/command-data.ts @@ -105,7 +105,7 @@ type CommandDataAccountQuotaResult = { updatedAccount?: CommandDataAccountMetadata } -type CommandDataAccountStorage = { +export interface CommandDataAccountStorage { version: 4 activeIndex: number activeIndexByFamily?: { claude?: number; gemini?: number } @@ -212,6 +212,24 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { } } +export function projectCommandAccountRows( + storage: CommandDataAccountStorage | null | undefined, +): CommandAccountRow[] { + if (!storage) return [] + const activeIndex = storage.activeIndexByFamily?.claude ?? storage.activeIndex + return storage.accounts.map((entry, index) => + toCommandAccountRow({ + index, + refreshToken: entry.refreshToken, + label: entry.label, + enabled: entry.enabled !== false, + active: index === activeIndex, + cachedQuota: entry.cachedQuota, + cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, + }), + ) +} + /** * Live AccountManager view the data service needs. * diff --git a/packages/opencode/src/plugin/commands.test.ts b/packages/opencode/src/plugin/commands.test.ts index d4d2c47..51c8f5d 100644 --- a/packages/opencode/src/plugin/commands.test.ts +++ b/packages/opencode/src/plugin/commands.test.ts @@ -256,6 +256,52 @@ describe('createCommandExecuteBefore', () => { await settings.dispose() }) + it('includes cached accounts in the account OPEN notification payload', async () => { + const notifications: Array>> = + [] + const pushNotification = ( + payload: Awaited>, + ) => { + notifications.push(payload) + } + const settings = createOperatorSettingsController({ + projectConfigPath: join(dir, 'antigravity.json'), + userConfigPath: join(dir, 'user.json'), + }) + const handler = createCommandExecuteBefore( + { session: { promptAsync: mock(async () => {}) } } as never, + settings, + pushNotification, + { + listAccounts: async () => [ + { + index: 0, + label: 'Primary account', + enabled: true, + active: true, + quota: {}, + }, + ], + } as never, + { isTuiConnected: () => true }, + ) + + await expect( + handler?.( + { + command: 'antigravity-account', + arguments: '', + sessionID: 'session-1', + }, + { parts: [] }, + ), + ).rejects.toThrow('ANTIGRAVITY_COMMAND_HANDLED') + + const payload = notifications[0] + expect(payload?.knobs.accounts as unknown[] | undefined).toHaveLength(1) + await settings.dispose() + }) + it('suppresses the ignored fallback for every modal command when the tui is connected', async () => { const promptAsync = mock(async () => {}) const messages = mock(async () => ({ data: [] })) @@ -574,6 +620,104 @@ describe('applyCommand', () => { expect(refresh.knobs.timeoutMs).toBe(120_000) }) + it('returns an OAuth URL and current rows when add-oauth-start is applied', async () => { + const start = mock(async () => ({ + url: 'https://accounts.google.test/authorize', + accounts: [ + { + id: 'acct-0', + index: 0, + label: 'Primary account', + enabled: true, + current: true, + quota: [], + }, + ], + })) + + const result = await applyCommand( + { command: 'antigravity-account', arguments: 'add-oauth-start' }, + { + client: {} as never, + sessionID: 'session-1', + settings: ctx.settings, + accountOAuth: { start } as never, + }, + ) + + expect(start).toHaveBeenCalledWith('session-1') + expect(result.text).toContain('Open this URL') + expect(result.knobs.oauthUrl).toBe('https://accounts.google.test/authorize') + expect(result.knobs.accounts as unknown[]).toHaveLength(1) + expect(result.knobs.timeoutMs).toBe(120_000) + }) + + it('returns refreshed rows after add-oauth-finish', async () => { + const finish = mock(async () => ({ + text: 'OAuth account added.', + accounts: [ + { + id: 'acct-0', + index: 0, + label: 'Primary account', + enabled: true, + current: true, + quota: [], + }, + { + id: 'acct-1', + index: 1, + label: 'Work account', + enabled: true, + current: false, + quota: [], + }, + ], + })) + + const result = await applyCommand( + { + command: 'antigravity-account', + arguments: 'add-oauth-finish callback-code', + }, + { + client: {} as never, + sessionID: 'session-1', + settings: ctx.settings, + accountOAuth: { finish } as never, + }, + ) + + expect(finish).toHaveBeenCalledWith('session-1', 'callback-code') + expect(result.text).toBe('OAuth account added.') + expect(result.knobs.accounts as unknown[]).toHaveLength(2) + expect(result.knobs.timeoutMs).toBe(120_000) + }) + + it('returns the expired-pending result from add-oauth-finish', async () => { + const finish = mock(async () => ({ + text: 'OAuth session expired. Please start again.', + accounts: [], + })) + + const result = await applyCommand( + { + command: 'antigravity-account', + arguments: 'add-oauth-finish callback-code', + }, + { + client: {} as never, + sessionID: 'session-1', + settings: ctx.settings, + accountOAuth: { finish } as never, + }, + ) + + expect(finish).toHaveBeenCalledWith('session-1', 'callback-code') + expect(result.text).toBe('OAuth session expired. Please start again.') + expect(result.knobs.accounts as unknown[]).toHaveLength(0) + }) + it('account apply dispatches current/toggle/remove through the data service', async () => { const toggleSpy = mock(async () => [ { diff --git a/packages/opencode/src/plugin/commands.ts b/packages/opencode/src/plugin/commands.ts index 2b466c9..f22f5ad 100644 --- a/packages/opencode/src/plugin/commands.ts +++ b/packages/opencode/src/plugin/commands.ts @@ -117,7 +117,7 @@ interface CommandContext { * Commands that need to chain side effects (e.g. account add which * triggers an OAuth flow) hook their own refresh. */ - onApplied?: () => Promise | void + onApplied?: (accounts?: CommandAccountRow[]) => Promise | void /** * Privacy-safe data service that backs the data-first dialogs. The * production wiring injects one that reads from the live AccountManager @@ -126,6 +126,22 @@ interface CommandContext { * undefined and the dialog falls back to the legacy placeholder. */ commandData?: CommandDataService + accountOAuth?: AccountCommandOAuthService +} + +export interface AccountCommandOAuthService { + start(sessionId: string): Promise<{ + url: string + accounts: CommandAccountRow[] + }> + finish( + sessionId: string, + code: string, + label?: string, + ): Promise<{ + text: string + accounts: CommandAccountRow[] + }> } /** @@ -169,6 +185,9 @@ export async function buildDialogPayload( } case 'antigravity-account': { const action = argumentsText.trim().toLowerCase() + const accounts = context.commandData + ? await context.commandData.listAccounts() + : [] return { command, text: 'Antigravity accounts', @@ -180,6 +199,7 @@ export async function buildDialogPayload( action === 'list' ? action : 'list', + accounts, }, } } @@ -323,6 +343,8 @@ function parseLoggingLevel(input: string): OperatorSettings['log_level'] { type AccountAction = | { kind: 'add' } + | { kind: 'add-oauth-start' } + | { kind: 'add-oauth-finish'; code: string; label?: string } | { kind: 'refresh' } | { kind: 'current'; index: number } | { kind: 'toggle'; index: number } @@ -351,6 +373,22 @@ function parseAccountAction(input: string): AccountAction | undefined { const parts = trimmed.split(/\s+/).filter(Boolean) const head = parts[0]?.toLowerCase() if (head === 'add') return { kind: 'add' } + if (head === 'add-oauth-start' && parts.length === 1) { + return { kind: 'add-oauth-start' } + } + if (head === 'add-oauth-finish') { + const code = parts[1] + if (!code) return undefined + const labelAt = parts.indexOf('--label') + const label = + labelAt === -1 + ? undefined + : parts + .slice(labelAt + 1) + .join(' ') + .trim() + return { kind: 'add-oauth-finish', code, label: label || undefined } + } if (head === 'refresh') return { kind: 'refresh' } if (head === 'current' || head === 'toggle' || head === 'remove') { const raw = parts[1] @@ -396,7 +434,12 @@ export async function applyCommand( // sees a fresh `checkedAt`. The refresher is optional; tests and // read-only contexts leave it undefined and skip the write entirely. if (context.onApplied) { - await Promise.resolve(context.onApplied()).catch(() => { + const accounts = result.knobs.accounts + await Promise.resolve( + context.onApplied( + Array.isArray(accounts) ? (accounts as CommandAccountRow[]) : undefined, + ), + ).catch(() => { // Sidebar refresh must never break the command's apply response. }) } @@ -489,17 +532,6 @@ async function applyCommandInner( } } case 'antigravity-account': { - // Apply arguments follow the ` ` pattern the - // dialog uses to call back into the apply path: - // `current ` — pin the account as active for every family - // `toggle ` — flip the enabled flag - // `remove ` — drop the account (destructive, confirmed) - // `add` — backwards-compatible; opens the OAuth flow - // via the 120s timeout so the dialog can - // signal "Account add requested" (the OAuth - // implementation lands in a follow-up task). - // Bare `add` / `refresh` keep their 120s timeout so the OAuth - // handshake has room; everything else stays on the 2s default. const args = request.arguments.trim() const parsed = parseAccountAction(args) if (!parsed) { @@ -509,11 +541,46 @@ async function applyCommandInner( } } + if (parsed.kind === 'add-oauth-start') { + if (!context.accountOAuth) { + return { + text: 'OAuth account add is unavailable.', + knobs: { timeoutMs: 120_000, error: true }, + } + } + const result = await context.accountOAuth.start(context.sessionID) + return { + text: `Open this URL in your browser:\n${result.url}`, + knobs: { + oauthUrl: result.url, + accounts: result.accounts, + timeoutMs: 120_000, + }, + } + } + + if (parsed.kind === 'add-oauth-finish') { + if (!context.accountOAuth) { + return { + text: 'OAuth account add is unavailable.', + knobs: { timeoutMs: 120_000, error: true }, + } + } + const result = parsed.label + ? await context.accountOAuth.finish( + context.sessionID, + parsed.code, + parsed.label, + ) + : await context.accountOAuth.finish(context.sessionID, parsed.code) + return { + text: result.text, + knobs: { accounts: result.accounts, timeoutMs: 120_000 }, + } + } + // The data service drives the locked-storage mutator for the - // three CRUD mutations; `add` stays a no-op response so the - // dispatcher keeps its existing 120s RPC timeout contract. - // - // The locked-storage write is awaited FIRST (service-side) so a + // three CRUD mutations. The locked-storage write is awaited FIRST so a // failed disk write never leaves the runtime ahead of disk. When // that await throws — AccountStorageUnreadableError, lock // contention, I/O — we catch here and surface the message as the @@ -701,9 +768,17 @@ export function createSidebarRefresher( 'gemini-flash'?: { remainingFraction?: number; resetTime?: string } } }> | null, -): () => Promise { - return async () => { - const accounts = getAccounts() +): (accounts?: CommandAccountRow[]) => Promise { + return async (dialogAccounts) => { + const accounts = dialogAccounts + ? dialogAccounts.map((entry) => ({ + index: entry.index, + label: entry.label, + enabled: entry.enabled, + coolingDownUntil: undefined, + cachedQuota: undefined, + })) + : getAccounts() if (!accounts || accounts.length === 0) return try { await setSidebarMachineState( diff --git a/packages/opencode/src/plugin/index.ts b/packages/opencode/src/plugin/index.ts index bac9fb9..70cbf46 100644 --- a/packages/opencode/src/plugin/index.ts +++ b/packages/opencode/src/plugin/index.ts @@ -1,4 +1,5 @@ import { join } from 'node:path' +import { authorizeAntigravity, exchangeAntigravity } from '../antigravity/oauth' import { ANTIGRAVITY_PROVIDER_ID } from '../constants' import { createAutoUpdateCheckerHook } from '../hooks/auto-update-checker' import { drainNotifications, pushNotification } from '../rpc/notifications' @@ -10,6 +11,7 @@ import { promptAccountIndexForVerification, promptOpenVerificationUrl, } from './account-access' +import { createAccountCommandOAuthService } from './account-command-oauth' import { createAuthLoader } from './auth-loader' import { initDiskSignatureCache, shutdownDiskSignatureCache } from './cache' import { @@ -19,6 +21,7 @@ import { import { type CommandDataService, createCommandDataService, + projectCommandAccountRows, } from './command-data' import { applyCommand, @@ -285,6 +288,14 @@ export const createAntigravityPlugin = confirmOpenVerificationUrl: promptOpenVerificationUrl, }, }) + const accountOAuth = createAccountCommandOAuthService({ + authorize: () => authorizeAntigravity(), + exchange: exchangeAntigravity, + persist: (result) => accountAccess.persistAccountPool([result], false), + listAccounts: async () => + projectCommandAccountRows(await accountAccess.loadAccounts()), + }) + lifecycle.register({ dispose: () => accountOAuth.dispose() }) const oauthMethods = createOAuthMethods({ client, providerId, @@ -341,6 +352,7 @@ export const createAntigravityPlugin = settings: operatorSettings, onApplied: refreshSidebar, commandData, + accountOAuth, }) // /antigravity-logging mutates the log level — propagate // immediately so subsequent log calls in this session respect it. diff --git a/packages/opencode/src/tui-compiled/tui/command-dialogs.tsx b/packages/opencode/src/tui-compiled/tui/command-dialogs.tsx index dca8198..1ca7418 100644 --- a/packages/opencode/src/tui-compiled/tui/command-dialogs.tsx +++ b/packages/opencode/src/tui-compiled/tui/command-dialogs.tsx @@ -1,8 +1,8 @@ -import { createTextNode as _$createTextNode } from "opentui:runtime-module:%40opentui%2Fsolid"; -import { insertNode as _$insertNode } from "opentui:runtime-module:%40opentui%2Fsolid"; import { memo as _$memo } from "opentui:runtime-module:%40opentui%2Fsolid"; import { insert as _$insert } from "opentui:runtime-module:%40opentui%2Fsolid"; import { setProp as _$setProp } from "opentui:runtime-module:%40opentui%2Fsolid"; +import { createTextNode as _$createTextNode } from "opentui:runtime-module:%40opentui%2Fsolid"; +import { insertNode as _$insertNode } from "opentui:runtime-module:%40opentui%2Fsolid"; import { createElement as _$createElement } from "opentui:runtime-module:%40opentui%2Fsolid"; import { createComponent as _$createComponent } from "opentui:runtime-module:%40opentui%2Fsolid"; /** @jsxImportSource @opentui/solid */ @@ -228,27 +228,146 @@ function renderAccountDialog(api, initialPayload, apply) { renderMain(); return null; }; + const updateAccounts = knobs => { + const accounts = knobs.accounts; + if (!Array.isArray(accounts)) return; + payload.knobs = { + ...payload.knobs, + accounts: accounts + }; + }; + const openOAuthLabelPrompt = (code, oauthUrl) => { + const DialogPrompt = api.ui.DialogPrompt; + api.ui.dialog.setSize('xlarge'); + api.ui.dialog.replace(() => _$createComponent(DialogPrompt, { + title: "OAuth sign-in \u2014 label", + description: () => (() => { + var _el$ = _$createElement("text"); + _$insertNode(_el$, _$createTextNode(`A short name for this account (optional).`)); + return _el$; + })(), + placeholder: "e.g. work", + value: "", + onConfirm: value => { + const label = value.trim(); + const args = label ? `add-oauth-finish ${code} --label ${label}` : `add-oauth-finish ${code}`; + void apply('antigravity-account', args, { + timeoutMs: 120_000 + }).then(result => { + api.ui.toast({ + message: result.text + }); + updateAccounts(result.knobs); + renderMain(); + }).catch(() => { + api.ui.toast({ + message: 'OAuth account add failed' + }); + renderMain(); + }); + }, + onCancel: () => openOAuthCodePrompt(oauthUrl) + })); + }; + const openOAuthCodePrompt = oauthUrl => { + const DialogPrompt = api.ui.DialogPrompt; + api.ui.dialog.setSize('xlarge'); + api.ui.dialog.replace(() => _$createComponent(DialogPrompt, { + title: "OAuth sign-in \u2014 enter code", + description: () => (() => { + var _el$3 = _$createElement("text"); + _$insertNode(_el$3, _$createTextNode(`After signing in, paste the full callback URL or authorization code.`)); + return _el$3; + })(), + placeholder: "Paste callback URL or code here", + value: "", + onConfirm: value => { + const code = value.trim(); + if (!code) { + renderMain(); + return; + } + openOAuthLabelPrompt(code, oauthUrl); + }, + onCancel: () => openOAuthUrlScreen(oauthUrl) + })); + }; + const openOAuthUrlScreen = oauthUrl => { + const DialogSelect = api.ui.DialogSelect; + api.ui.dialog.setSize('xlarge'); + api.ui.dialog.replace(() => _$createComponent(DialogSelect, { + title: "OAuth sign-in", + options: [{ + title: 'Copy URL to clipboard', + value: 'copy', + description: oauthUrl + }, { + title: 'Enter sign-in code', + value: 'code', + description: 'Open the URL in your browser, sign in, then paste the callback URL or code.' + }, { + title: 'Cancel', + value: 'cancel' + }], + onSelect: option => { + if (option.value === 'cancel') { + renderMain(); + return; + } + if (option.value === 'copy') { + const copied = api.renderer.copyToClipboardOSC52(oauthUrl); + api.ui.toast({ + message: copied ? 'URL copied to clipboard' : 'Copy unavailable — select the URL text above to copy' + }); + openOAuthUrlScreen(oauthUrl); + return; + } + openOAuthCodePrompt(oauthUrl); + } + })); + }; + const openAddOAuthStart = () => { + void apply('antigravity-account', 'add-oauth-start', { + timeoutMs: 120_000 + }).then(result => { + updateAccounts(result.knobs); + const oauthUrl = result.knobs.oauthUrl; + if (typeof oauthUrl === 'string' && oauthUrl.length > 0) { + openOAuthUrlScreen(oauthUrl); + return; + } + api.ui.toast({ + message: result.text + }); + renderMain(); + }).catch(() => { + api.ui.toast({ + message: 'OAuth account add failed' + }); + renderMain(); + }); + }; const renderManageRow = row => { const DialogSelect = api.ui.DialogSelect; api.ui.dialog.setSize('xlarge'); api.ui.dialog.replace(() => (() => { - var _el$ = _$createElement("box"), - _el$2 = _$createElement("text"), - _el$3 = _$createElement("text"), - _el$4 = _$createElement("box"); - _$insertNode(_el$, _el$2); - _$insertNode(_el$, _el$3); - _$insertNode(_el$, _el$4); - _$setProp(_el$, "flexDirection", 'column'); - _$setProp(_el$, "padding", 1); - _$setProp(_el$, "width", '100%'); - _$insert(_el$2, () => row.label); - _$insert(_el$3, (() => { + var _el$5 = _$createElement("box"), + _el$6 = _$createElement("text"), + _el$7 = _$createElement("text"), + _el$8 = _$createElement("box"); + _$insertNode(_el$5, _el$6); + _$insertNode(_el$5, _el$7); + _$insertNode(_el$5, _el$8); + _$setProp(_el$5, "flexDirection", 'column'); + _$setProp(_el$5, "padding", 1); + _$setProp(_el$5, "width", '100%'); + _$insert(_el$6, () => row.label); + _$insert(_el$7, (() => { var _c$ = _$memo(() => row.quota.length === 0); return () => _c$() ? 'no cached quota' : row.quota.map(q => q.remainingPercent == null ? `${q.label} –%` : `${q.label} ${q.remainingPercent}%`).join(', '); })()); - _$setProp(_el$4, "marginTop", 1); - _$insert(_el$4, _$createComponent(DialogSelect, { + _$setProp(_el$8, "marginTop", 1); + _$insert(_el$8, _$createComponent(DialogSelect, { get title() { return `Manage ${row.label}`; }, @@ -297,7 +416,7 @@ function renderAccountDialog(api, initialPayload, apply) { }); } })); - return _el$; + return _el$5; })()); }; const promptRemove = index => { @@ -335,19 +454,19 @@ function renderAccountDialog(api, initialPayload, apply) { const rows = rowsFromKnobs(); const bodyLines = rows.length ? rows.map(formatQuotaRowLine) : ['No accounts configured. Add one via the menu below.']; api.ui.dialog.replace(() => (() => { - var _el$5 = _$createElement("box"), - _el$6 = _$createElement("box"); - _$insertNode(_el$5, _el$6); - _$setProp(_el$5, "flexDirection", 'column'); - _$setProp(_el$5, "padding", 1); - _$setProp(_el$5, "width", '100%'); - _$insert(_el$5, () => bodyLines.map(line => (() => { - var _el$7 = _$createElement("text"); - _$insert(_el$7, line); - return _el$7; - })()), _el$6); - _$setProp(_el$6, "marginTop", 1); - _$insert(_el$6, _$createComponent(DialogSelect, { + var _el$9 = _$createElement("box"), + _el$0 = _$createElement("box"); + _$insertNode(_el$9, _el$0); + _$setProp(_el$9, "flexDirection", 'column'); + _$setProp(_el$9, "padding", 1); + _$setProp(_el$9, "width", '100%'); + _$insert(_el$9, () => bodyLines.map(line => (() => { + var _el$1 = _$createElement("text"); + _$insert(_el$1, line); + return _el$1; + })()), _el$0); + _$setProp(_el$0, "marginTop", 1); + _$insert(_el$0, _$createComponent(DialogSelect, { title: "Antigravity accounts", get options() { return [{ @@ -370,13 +489,7 @@ function renderAccountDialog(api, initialPayload, apply) { return; } if (raw === 'add') { - void runApply('add', { - timeoutMs: 120_000 - }).catch(() => { - api.ui.toast({ - message: 'Account add failed' - }); - }); + openAddOAuthStart(); return; } if (raw.startsWith('__manage__ ')) { @@ -390,7 +503,7 @@ function renderAccountDialog(api, initialPayload, apply) { }); } })); - return _el$5; + return _el$9; })()); }; renderMain(); @@ -452,19 +565,19 @@ function renderRoutingDialog(api, initialPayload, apply) { } = readRouting(); const bodyLines = [`cli_first: ${cliFirst ? 'on' : 'off'}`, `quota_style_fallback: ${quotaFallback ? 'on' : 'off'}`]; api.ui.dialog.replace(() => (() => { - var _el$8 = _$createElement("box"), - _el$9 = _$createElement("box"); - _$insertNode(_el$8, _el$9); - _$setProp(_el$8, "flexDirection", 'column'); - _$setProp(_el$8, "padding", 1); - _$setProp(_el$8, "width", '100%'); - _$insert(_el$8, () => bodyLines.map(line => (() => { - var _el$0 = _$createElement("text"); - _$insert(_el$0, line); - return _el$0; - })()), _el$9); - _$setProp(_el$9, "marginTop", 1); - _$insert(_el$9, _$createComponent(DialogSelect, { + var _el$10 = _$createElement("box"), + _el$11 = _$createElement("box"); + _$insertNode(_el$10, _el$11); + _$setProp(_el$10, "flexDirection", 'column'); + _$setProp(_el$10, "padding", 1); + _$setProp(_el$10, "width", '100%'); + _$insert(_el$10, () => bodyLines.map(line => (() => { + var _el$12 = _$createElement("text"); + _$insert(_el$12, line); + return _el$12; + })()), _el$11); + _$setProp(_el$11, "marginTop", 1); + _$insert(_el$11, _$createComponent(DialogSelect, { title: "Antigravity routing", options: [{ title: `${cliFirst ? '●' : '○'} cli_first: ${cliFirst}`, @@ -498,7 +611,7 @@ function renderRoutingDialog(api, initialPayload, apply) { }); } })); - return _el$8; + return _el$10; })()); }; renderMain(); @@ -559,9 +672,9 @@ function renderKillswitchDialog(api, initialPayload, apply) { api.ui.dialog.replace(() => _$createComponent(DialogPrompt, { title: "Antigravity killswitch \u2014 set threshold", description: () => (() => { - var _el$1 = _$createElement("text"); - _$insertNode(_el$1, _$createTextNode(`Enter a new minimum_remaining_percent (0-100). Empty input cancels.`)); - return _el$1; + var _el$13 = _$createElement("text"); + _$insertNode(_el$13, _$createTextNode(`Enter a new minimum_remaining_percent (0-100). Empty input cancels.`)); + return _el$13; })(), placeholder: `${currentMinimum}`, value: `${currentMinimum}`, @@ -602,19 +715,19 @@ function renderKillswitchDialog(api, initialPayload, apply) { const accountKeys = Object.keys(accounts); const bodyLines = [`Killswitch: ${enabled ? 'enabled' : 'disabled'}`, `Minimum remaining percent: ${minimum}%`, ...(accountKeys.length > 0 ? ['Per-account overrides:', ...accountKeys.map(key => ` ${key}: ${accounts[key] ?? 0}%`)] : [])]; api.ui.dialog.replace(() => (() => { - var _el$11 = _$createElement("box"), - _el$12 = _$createElement("box"); - _$insertNode(_el$11, _el$12); - _$setProp(_el$11, "flexDirection", 'column'); - _$setProp(_el$11, "padding", 1); - _$setProp(_el$11, "width", '100%'); - _$insert(_el$11, () => bodyLines.map(line => (() => { - var _el$13 = _$createElement("text"); - _$insert(_el$13, line); - return _el$13; - })()), _el$12); - _$setProp(_el$12, "marginTop", 1); - _$insert(_el$12, _$createComponent(DialogSelect, { + var _el$15 = _$createElement("box"), + _el$16 = _$createElement("box"); + _$insertNode(_el$15, _el$16); + _$setProp(_el$15, "flexDirection", 'column'); + _$setProp(_el$15, "padding", 1); + _$setProp(_el$15, "width", '100%'); + _$insert(_el$15, () => bodyLines.map(line => (() => { + var _el$17 = _$createElement("text"); + _$insert(_el$17, line); + return _el$17; + })()), _el$16); + _$setProp(_el$16, "marginTop", 1); + _$insert(_el$16, _$createComponent(DialogSelect, { title: "Antigravity killswitch", options: [{ title: `${enabled ? '●' : '○'} Killswitch: ${enabled ? 'enabled' : 'disabled'}`, @@ -652,7 +765,7 @@ function renderKillswitchDialog(api, initialPayload, apply) { }); } })); - return _el$11; + return _el$15; })()); }; renderMain(); diff --git a/packages/opencode/src/tui/command-dialogs.test.tsx b/packages/opencode/src/tui/command-dialogs.test.tsx index ffbbdc9..6daa9ea 100644 --- a/packages/opencode/src/tui/command-dialogs.test.tsx +++ b/packages/opencode/src/tui/command-dialogs.test.tsx @@ -930,6 +930,78 @@ describe('openCommandDialog (imperative dispatcher)', () => { } }) + it('antigravity-account adds an OAuth account through URL, code, and label prompts', async () => { + const localFake = makeFakeApi() + applyMock.mockImplementationOnce(async (_cmd, _args, options) => ({ + text: 'Open this URL in your browser', + knobs: { + oauthUrl: 'https://accounts.google.test/authorize', + accounts: [], + timeoutMs: options?.timeoutMs ?? 120_000, + }, + })) + applyMock.mockImplementationOnce(async (_cmd, _args, options) => ({ + text: 'OAuth account added.', + knobs: { + accounts: [ + { + id: 'acct-0', + index: 0, + label: 'Work account', + enabled: true, + current: true, + quota: [], + }, + ], + timeoutMs: options?.timeoutMs ?? 120_000, + }, + })) + dispatcher.openCommandDialog( + localFake, + payloadFor('antigravity-account', { accounts: [] }), + applyFor('antigravity-account'), + ) + await renderDialog(localFake) + + const add = localFake.capturedSelectProps?.options?.find( + (option) => option.title === 'Add account…', + ) + localFake.capturedSelectProps?.onSelect?.({ + title: add?.title, + value: add?.value, + }) + for (let i = 0; i < 5; i += 1) { + await new Promise((resolve) => setImmediate(resolve)) + } + await renderDialog(localFake) + + expect(applyMock.mock.calls.at(0)?.[1]).toBe('add-oauth-start') + expect(applyMock.mock.calls.at(0)?.[2]?.timeoutMs).toBe(120_000) + expect(localFake.capturedSelectProps?.title).toBe('OAuth sign-in') + + const enterCode = localFake.capturedSelectProps?.options?.find( + (option) => option.title === 'Enter sign-in code', + ) + localFake.capturedSelectProps?.onSelect?.({ + title: enterCode?.title, + value: enterCode?.value, + }) + await renderDialog(localFake) + expect(localFake.capturedPromptProps?.title).toContain('enter code') + localFake.capturedPromptProps?.onConfirm?.('callback-code') + await renderDialog(localFake) + expect(localFake.capturedPromptProps?.title).toContain('label') + localFake.capturedPromptProps?.onConfirm?.('Work account') + for (let i = 0; i < 5; i += 1) { + await new Promise((resolve) => setImmediate(resolve)) + } + + expect(applyMock.mock.calls.at(1)?.[1]).toBe( + 'add-oauth-finish callback-code --label Work account', + ) + expect(applyMock.mock.calls.at(1)?.[2]?.timeoutMs).toBe(120_000) + }) + it('antigravity-account opens a row subdialog with toggle/current/remove/back', async () => { const localFake = makeFakeApi() dispatcher.openCommandDialog( @@ -1230,7 +1302,7 @@ describe('openCommandDialog (imperative dispatcher)', () => { expect(applyMock.mock.calls.length).toBe(applyCallsBefore) }) - it('antigravity-account add action calls apply with 120s timeout', async () => { + it('antigravity-account add action starts OAuth with a 120s timeout', async () => { const localFake = makeFakeApi() dispatcher.openCommandDialog( localFake, @@ -1262,7 +1334,7 @@ describe('openCommandDialog (imperative dispatcher)', () => { await renderDialog(localFake) const call = applyMock.mock.calls.at(-1) expect(call?.[0]).toBe('antigravity-account') - expect(call?.[1]).toBe('add') + expect(call?.[1]).toBe('add-oauth-start') expect(call?.[2]?.timeoutMs).toBe(120_000) }) diff --git a/packages/opencode/src/tui/command-dialogs.tsx b/packages/opencode/src/tui/command-dialogs.tsx index e127daf..9a2b1a8 100644 --- a/packages/opencode/src/tui/command-dialogs.tsx +++ b/packages/opencode/src/tui/command-dialogs.tsx @@ -252,6 +252,134 @@ function renderAccountDialog( return null } + const updateAccounts = (knobs: Record): void => { + const accounts = knobs.accounts + if (!Array.isArray(accounts)) return + payload.knobs = { + ...payload.knobs, + accounts: accounts as CommandAccountRow[], + } + } + + const openOAuthLabelPrompt = (code: string, oauthUrl: string): void => { + const DialogPrompt = api.ui.DialogPrompt + api.ui.dialog.setSize('xlarge') + api.ui.dialog.replace(() => ( + ( + A short name for this account (optional). + )} + placeholder='e.g. work' + value='' + onConfirm={(value: string) => { + const label = value.trim() + const args = label + ? `add-oauth-finish ${code} --label ${label}` + : `add-oauth-finish ${code}` + void apply('antigravity-account', args, { timeoutMs: 120_000 }) + .then((result) => { + api.ui.toast({ message: result.text }) + updateAccounts(result.knobs) + renderMain() + }) + .catch(() => { + api.ui.toast({ message: 'OAuth account add failed' }) + renderMain() + }) + }} + onCancel={() => openOAuthCodePrompt(oauthUrl)} + /> + )) + } + + const openOAuthCodePrompt = (oauthUrl: string): void => { + const DialogPrompt = api.ui.DialogPrompt + api.ui.dialog.setSize('xlarge') + api.ui.dialog.replace(() => ( + ( + + After signing in, paste the full callback URL or authorization code. + + )} + placeholder='Paste callback URL or code here' + value='' + onConfirm={(value: string) => { + const code = value.trim() + if (!code) { + renderMain() + return + } + openOAuthLabelPrompt(code, oauthUrl) + }} + onCancel={() => openOAuthUrlScreen(oauthUrl)} + /> + )) + } + + const openOAuthUrlScreen = (oauthUrl: string): void => { + const DialogSelect = api.ui.DialogSelect + api.ui.dialog.setSize('xlarge') + api.ui.dialog.replace(() => ( + { + if (option.value === 'cancel') { + renderMain() + return + } + if (option.value === 'copy') { + const copied = api.renderer.copyToClipboardOSC52(oauthUrl) + api.ui.toast({ + message: copied + ? 'URL copied to clipboard' + : 'Copy unavailable — select the URL text above to copy', + }) + openOAuthUrlScreen(oauthUrl) + return + } + openOAuthCodePrompt(oauthUrl) + }} + /> + )) + } + + const openAddOAuthStart = (): void => { + void apply('antigravity-account', 'add-oauth-start', { + timeoutMs: 120_000, + }) + .then((result) => { + updateAccounts(result.knobs) + const oauthUrl = result.knobs.oauthUrl + if (typeof oauthUrl === 'string' && oauthUrl.length > 0) { + openOAuthUrlScreen(oauthUrl) + return + } + api.ui.toast({ message: result.text }) + renderMain() + }) + .catch(() => { + api.ui.toast({ message: 'OAuth account add failed' }) + renderMain() + }) + } + const renderManageRow = (row: CommandAccountRow): void => { const DialogSelect = api.ui.DialogSelect api.ui.dialog.setSize('xlarge') @@ -395,9 +523,7 @@ function renderAccountDialog( return } if (raw === 'add') { - void runApply('add', { timeoutMs: 120_000 }).catch(() => { - api.ui.toast({ message: 'Account add failed' }) - }) + openAddOAuthStart() return } if (raw.startsWith('__manage__ ')) { From 0141af63320c994ec1f2b78458e32a146e4af27a Mon Sep 17 00:00:00 2001 From: iceteaSA <171169159+iceteaSA@users.noreply.github.com> Date: Fri, 24 Jul 2026 15:34:10 +0200 Subject: [PATCH 02/25] fix(quota): refresh sidebar state on modal open (cherry picked from commit d10b553882d37c2996627de447d933db65517da2) --- .../opencode/src/plugin/command-data.test.ts | 4 +- packages/opencode/src/plugin/command-data.ts | 4 +- packages/opencode/src/plugin/commands.test.ts | 88 +++ packages/opencode/src/plugin/commands.ts | 32 +- .../src/tui-compiled/plugin/command-data.ts | 712 ++++++++++++++++++ 5 files changed, 826 insertions(+), 14 deletions(-) create mode 100644 packages/opencode/src/tui-compiled/plugin/command-data.ts diff --git a/packages/opencode/src/plugin/command-data.test.ts b/packages/opencode/src/plugin/command-data.test.ts index e3e480a..e89c2ed 100644 --- a/packages/opencode/src/plugin/command-data.test.ts +++ b/packages/opencode/src/plugin/command-data.test.ts @@ -446,7 +446,7 @@ describe('createCommandDataService', () => { expect(harness.quotaCallLog).toEqual([]) }) - it('refreshQuota() refreshes through the shared quota manager and persists by refresh token', async () => { + it('refreshQuota() fetches every account, including an uncached new account, through the shared quota manager', async () => { const baseAccount = ( refreshToken: string, groups: Record, @@ -506,8 +506,6 @@ describe('createCommandDataService', () => { makeAccountFixture({ refreshToken: 'refresh-b', label: 'Beta', - cachedQuota: { 'gemini-flash': { remainingFraction: 0.05 } }, - cachedQuotaUpdatedAt: 1, }), ], refreshResults, diff --git a/packages/opencode/src/plugin/command-data.ts b/packages/opencode/src/plugin/command-data.ts index e3c933f..08330d9 100644 --- a/packages/opencode/src/plugin/command-data.ts +++ b/packages/opencode/src/plugin/command-data.ts @@ -430,8 +430,8 @@ export function createCommandDataService( buildSidebarMachineStateFromAccounts(accounts, { checkedAt: now() }), { stateFile: sidebarStateFile }, ).catch(() => { - // Lock contention is the only realistic failure mode; the next - // periodic writer will catch up. + // Sidebar writes are best-effort; the next command or quota refresh + // will publish the current snapshot. }) } diff --git a/packages/opencode/src/plugin/commands.test.ts b/packages/opencode/src/plugin/commands.test.ts index 51c8f5d..6d0bce3 100644 --- a/packages/opencode/src/plugin/commands.test.ts +++ b/packages/opencode/src/plugin/commands.test.ts @@ -3,12 +3,15 @@ import { mkdtempSync, rmSync } from 'node:fs' import { tmpdir } from 'node:os' import { join } from 'node:path' import { AccountStorageUnreadableError } from '@cortexkit/antigravity-auth-core' + import type { CommandModalName } from '../rpc/protocol' +import { readSidebarState } from '../sidebar-state' import { registerAntigravityCommands } from './catalog' import { applyCommand, buildDialogPayload, createCommandExecuteBefore, + createSidebarRefresher, MODAL_COMMANDS, } from './commands' import { GEMINI_DUMP_COMMAND_NAME } from './gemini-dump' @@ -240,6 +243,7 @@ describe('createCommandExecuteBefore', () => { quota: {}, }, ], + refreshQuota: async () => [], } as never, { isTuiConnected: () => true }, ) @@ -694,6 +698,90 @@ describe('applyCommand', () => { expect(result.knobs.timeoutMs).toBe(120_000) }) + it('preserves an existing cached quota in the sidebar after add-oauth-finish', async () => { + const sidebarFile = join(dir, 'sidebar-state.json') + const previousSidebarFile = process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = sidebarFile + try { + const finish = mock(async () => ({ + text: 'OAuth account added.', + accounts: [ + { + id: 'acct-0', + index: 0, + label: 'Primary account', + enabled: true, + current: true, + quota: [ + { + key: 'claude' as const, + label: 'Claude', + remainingPercent: 50, + }, + ], + }, + { + id: 'acct-1', + index: 1, + label: 'New account', + enabled: true, + current: false, + quota: [], + }, + ], + })) + + await applyCommand( + { + command: 'antigravity-account', + arguments: 'add-oauth-finish callback-code', + }, + { + client: {} as never, + sessionID: 'session-1', + settings: ctx.settings, + accountOAuth: { finish } as never, + onApplied: createSidebarRefresher(() => []), + }, + ) + + const state = readSidebarState(sidebarFile) + expect(state.accounts).toHaveLength(2) + expect(state.accounts[0]?.quota.claude?.remainingPercent).toBe(50) + expect(state.accounts[1]?.quota).toEqual({}) + } finally { + if (previousSidebarFile === undefined) { + delete process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + } else { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + + it('starts a quota refresh when the quota dialog opens without delaying its cached payload', async () => { + const listAccounts = mock(async () => [ + { + id: 'acct-0', + index: 0, + label: 'Primary account', + enabled: true, + current: true, + quota: [], + }, + ]) + const refreshQuota = mock(async () => []) + + const payload = await buildDialogPayload('antigravity-quota', '', { + client: {} as never, + sessionID: 'session-1', + settings: ctx.settings, + commandData: { listAccounts, refreshQuota } as never, + }) + + expect(payload.knobs.accounts).toHaveLength(1) + expect(refreshQuota).toHaveBeenCalledTimes(1) + }) + it('returns the expired-pending result from add-oauth-finish', async () => { const finish = mock(async () => ({ text: 'OAuth session expired. Please start again.', diff --git a/packages/opencode/src/plugin/commands.ts b/packages/opencode/src/plugin/commands.ts index f22f5ad..e0989fa 100644 --- a/packages/opencode/src/plugin/commands.ts +++ b/packages/opencode/src/plugin/commands.ts @@ -165,15 +165,14 @@ export async function buildDialogPayload( switch (command) { case 'antigravity-quota': { const action = argumentsText.trim().toLowerCase() - // Opening the quota dialog is cache-only — we render the - // privacy-safe rows straight from the data service so the user - // sees the last-known percentages the moment the dialog mounts, - // with zero network I/O. The Refresh action (handled by - // `applyCommandInner`) is the only path that performs a live - // fetch. const accounts = context.commandData ? await context.commandData.listAccounts() : [] + // Render the cached snapshot immediately while the shared manager + // refreshes all accounts; the mounted panel polls the fenced state file. + if (context.commandData) { + void context.commandData.refreshQuota().catch(() => {}) + } return { command, text: 'Antigravity quota', @@ -776,7 +775,22 @@ export function createSidebarRefresher( label: entry.label, enabled: entry.enabled, coolingDownUntil: undefined, - cachedQuota: undefined, + cachedQuota: Object.fromEntries( + entry.quota.flatMap((group) => { + if (group.remainingPercent == null) return [] + return [ + [ + group.key, + { + remainingFraction: group.remainingPercent / 100, + ...(group.resetAt === undefined + ? {} + : { resetTime: new Date(group.resetAt).toISOString() }), + }, + ], + ] + }), + ), })) : getAccounts() if (!accounts || accounts.length === 0) return @@ -793,8 +807,8 @@ export function createSidebarRefresher( ), ) } catch { - // Lock contention is the only realistic failure — sidebar refresh - // is best-effort and the next periodic writer will catch up. + // Sidebar refresh is best-effort; the next command or quota refresh + // will publish the current account snapshot. } } } diff --git a/packages/opencode/src/tui-compiled/plugin/command-data.ts b/packages/opencode/src/tui-compiled/plugin/command-data.ts new file mode 100644 index 0000000..e7f76e6 --- /dev/null +++ b/packages/opencode/src/tui-compiled/plugin/command-data.ts @@ -0,0 +1,712 @@ +/** + * Privacy-safe data service for the data-first slash-command dialogs. + * + * `/antigravity-quota` (this task) and the future `/antigravity-account` + * dialogs (Tasks 10-11) read from a single shared service so they + * never touch raw account storage, never see the email PII field, and + * never run quota network I/O during the dialog's open path — the + * Refresh action is the only path that performs a live fetch. + * + * Why this is a separate module: + * + * - `commands.ts` owns the slash-command orchestration; mixing the row + * projection + quota refresh logic in there would balloon the surface + * for what is fundamentally a small read/refresh service. + * - The row type (`CommandAccountRow`) is the projection that crosses + * the PII firewall into the dialog payload. Defining it here (next + * to the code that constructs it) keeps the firewall reviewable. + * - Tests can pin the read-vs-refresh boundary, the email redaction, and + * the refresh-token-keyed persistence in one file rather than chasing + * them across the dispatcher + RPC + apply layers. + * + * Cache-only opening contract (Task 9 operator requirement): + * + * `listAccounts()` is a pure read of the live AccountManager view + * the auth-loader materialized at session start. It performs + * ZERO quota manager calls — opening the dialog must be instant even + * when the network is unreachable, and quota refresh must remain an + * explicit user-driven action so the cached percentages never quietly + * rewrite themselves behind the user's back. + * + * Refresh-token-keyed persistence (Task 9 plan trap): + * + * `refreshQuota()` runs the shared quota manager across every enabled + * account and folds the results back into the live AccountManager + + * storage. Concurrent OAuth can renumber the flat `accounts[]` array + * between read and write, so we re-read under the lock and key the + * update by `refreshToken` (the canonical identity) rather than the + * array index. This way a successful OAuth add that lands between + * the read and the write cannot cause the refresh to overwrite the + * wrong account. + */ + +import { + buildSidebarMachineStateFromAccounts, + type SidebarAccountRedactionInput, + setSidebarMachineState, +} from '../sidebar-state' + +/** + * Local copies of the few core types the data service touches. + * + * The shipped TUI tree cannot depend on the `@cortexkit/antigravity-auth-core` + * barrel (only subpath imports like `./file-lock` are allowed) — the + * compiled `command-data.ts` is copied verbatim into `tui-compiled/`, + * and the type-import would otherwise leak the full core surface into + * the dialog render path. Declaring the structural shape locally keeps + * the compiled tree free of the barrel while preserving duck-typed + * compatibility with the production quota manager. + */ +type CommandDataAccountMetadata = { + email?: string + refreshToken: string + projectId?: string + managedProjectId?: string + addedAt: number + lastUsed: number + enabled?: boolean + label?: string + cachedQuota?: Partial< + Record + > + cachedQuotaUpdatedAt?: number +} + +type CommandDataQuotaGroup = + | 'claude' + | 'gemini-pro' + | 'gemini-flash' + | 'gpt-oss' + +type CommandDataQuotaGroupSummary = { + remainingFraction?: number + resetTime?: string + modelCount: number +} + +type CommandDataAccountQuotaResult = { + index: number + status: 'ok' | 'disabled' | 'error' + error?: string + disabled?: boolean + quota?: { + groups: Partial> + perModel?: Array<{ + modelId: string + displayName?: string + group: CommandDataQuotaGroup | null + remainingFraction: number + resetTime?: string + }> + modelCount: number + error?: string + } + updatedAccount?: CommandDataAccountMetadata +} + +export interface CommandDataAccountStorage { + version: 4 + activeIndex: number + activeIndexByFamily?: { claude?: number; gemini?: number } + accounts: CommandDataAccountMetadata[] +} + +/** + * Privacy-safe per-account row shown in `/antigravity-*` data-first + * dialogs. Carries the cached quota percentages and the labels the + * dialog needs, but NEVER the email — the email stays in private + * account storage and is invisible to the sidebar and the dialog. + * + * `index` is the position in the live account array at the time of + * projection. It is informational only and is NOT a stable identity + * across refreshes — concurrent OAuth can renumber the array between + * the dialog opening and a refresh. + */ +export interface CommandAccountRow { + /** Stable identity for dialog keying (`acct-`). */ + id: string + /** Position in the live account array at projection time. */ + index: number + /** Display label (PII-free OAuth `name`, falling back to `Account N`). */ + label: string + enabled: boolean + /** `true` when this row matches the harness-active account. */ + current: boolean + quota: Array<{ + key: 'claude' | 'gemini-pro' | 'gemini-flash' + label: string + remainingPercent: number | null + resetAt?: number + }> +} + +/** + * Quota display label for each supported quota group. Kept here so the + * dialog and the quota manager agree on the same vocabulary. + */ +const QUOTA_GROUP_LABELS: Record< + CommandAccountRow['quota'][number]['key'], + string +> = { + claude: 'Claude', + 'gemini-pro': 'Gemini Pro', + 'gemini-flash': 'Gemini Flash', +} + +const SUPPORTED_QUOTA_KEYS = [ + 'claude', + 'gemini-pro', + 'gemini-flash', +] as const satisfies readonly CommandAccountRow['quota'][number]['key'][] + +interface LiveAccountSnapshot { + index: number + refreshToken: string + label?: string + enabled: boolean + active: boolean + cachedQuota?: Partial< + Record + > + cachedQuotaUpdatedAt?: number +} + +function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { + const cached = entry.cachedQuota + const quota: CommandAccountRow['quota'] = [] + for (const key of SUPPORTED_QUOTA_KEYS) { + const cachedEntry = cached?.[key] + if (!cachedEntry) continue + const fraction = cachedEntry.remainingFraction + const remainingPercent = + typeof fraction === 'number' && Number.isFinite(fraction) + ? Math.round(fraction * 100) + : null + let resetAt: number | undefined + if ( + typeof cachedEntry.resetTime === 'string' && + cachedEntry.resetTime.length > 0 + ) { + const parsed = Date.parse(cachedEntry.resetTime) + if (Number.isFinite(parsed)) resetAt = parsed + } + quota.push({ + key, + label: QUOTA_GROUP_LABELS[key], + remainingPercent, + resetAt, + }) + } + const label = entry.label ?? `Account ${entry.index + 1}` + return { + id: `acct-${entry.index}`, + index: entry.index, + label, + enabled: entry.enabled, + current: entry.active, + quota, + } +} + +export function projectCommandAccountRows( + storage: CommandDataAccountStorage | null | undefined, +): CommandAccountRow[] { + if (!storage) return [] + const activeIndex = storage.activeIndexByFamily?.claude ?? storage.activeIndex + return storage.accounts.map((entry, index) => + toCommandAccountRow({ + index, + refreshToken: entry.refreshToken, + label: entry.label, + enabled: entry.enabled !== false, + active: index === activeIndex, + cachedQuota: entry.cachedQuota, + cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, + }), + ) +} + +/** + * Live AccountManager view the data service needs. + * + * - `getAccounts()` returns the in-memory snapshot — used for cache-only + * reads during dialog open and for re-reading the post-mutation state. + * `getAccountsForQuotaCheck()` returns the freshest `AccountMetadataV3` + * (the canonical shape the quota manager expects). + * - `updateQuotaCache(index, groups)` + `requestSaveToDisk()` fold the + * refreshed quota back into the live view and persist it. The service + * calls them together after the network fetch resolves. + * - `setAccountEnabled`, `setAccountCurrent`, `removeAccountByIndex` + * mirror the CLI menu's mutation primitives — the dialog actions + * (Task 10) reuse them so the TUI never invents its own mutation + * logic. Each method returns `true` when the live view changed. + * - `getRefreshTokenAt(index)` lets the service identify the canonical + * account identity before it reaches the locked storage mutator; a + * concurrent OAuth could renumber the flat array between the dialog + * opening and the apply, so keying by refresh token is mandatory. + * - `flushSaveToDisk()` drains the AccountManager's debounced save so a + * dialog-triggered mutation lands on disk before the dialog's response + * returns. Without it, the dialog could toast "Account removed" while + * the file still carries the old pool. + * - `activeIndex()` returns the position the harness considers active; + * the service uses it to mark the matching row as `current: true`. + */ +export interface CommandDataAccountManagerView { + getAccounts(): LiveAccountSnapshot[] + getAccountsForQuotaCheck(): CommandDataAccountMetadata[] + updateQuotaCache( + index: number, + groups: Partial< + Record + >, + ): void + requestSaveToDisk(): void + flushSaveToDisk(): Promise + activeIndex(): number + /** Enable/disable the account at `index`. Returns true when it changed. */ + setAccountEnabled(index: number, enabled: boolean): boolean + /** Pin `index` as the active account for every family the dialog cares about. */ + setAccountCurrent(index: number): boolean + /** Remove the account at `index` from the live view. Returns true when it changed. */ + removeAccountByIndex(index: number): boolean + /** Canonical refresh token for the account at `index`, or undefined. */ + getRefreshTokenAt(index: number): string | undefined +} + +/** + * Storage adapter the data service uses for re-read-under-lock writes. + * `mutate` runs the supplied callback against the latest snapshot under + * the file lock — the callback receives the current storage and may + * return a replacement; returning the input unchanged is a no-op. + * + * The promise resolves with the (possibly mutated) storage snapshot so + * the data service can await it. A rejected promise signals a failed + * write — `AccountStorageUnreadableError`, lock contention, or any + * other I/O failure — which the data service surfaces to the dialog + * as a friendly error toast. + */ +export interface CommandDataStorage { + mutate( + mutator: ( + current: CommandDataAccountStorage, + ) => + | CommandDataAccountStorage + | undefined + | Promise, + ): Promise | undefined +} + +/** + * Options for `createCommandDataService`. Each field is required so + * production wiring is explicit — a missing dependency is a startup + * error, not a silent no-op at dialog-open time. + */ +export interface CommandDataServiceOptions { + accountManagerView: CommandDataAccountManagerView + quotaManager: { + refreshAccounts( + accounts: CommandDataAccountMetadata[], + options: { + indexFor?: (account: CommandDataAccountMetadata) => number + force?: boolean + }, + ): Promise + } + /** Path to the label-only sidebar state file. */ + sidebarStateFile: string + /** + * Optional storage adapter. When provided, the service persists the + * refreshed quota to disk via the lock-held mutator (best-effort — + * a lock contention never breaks the dialog response). When omitted, + * the service still folds results into the live AccountManager view; + * the auth-loader is responsible for the on-disk write. + */ + storage?: CommandDataStorage + /** Clock for `cachedQuotaUpdatedAt`. Tests inject a fixed clock. */ + now?: () => number +} + +/** + * Public surface of the command-data service. Each method is the + * smallest possible projection so the dialog layer never has to know + * how the underlying quota manager or storage adapter is wired. + */ +export interface CommandDataService { + /** + * Cache-only snapshot. Performs zero quota manager calls — safe to + * call as part of the dialog's open path. + */ + listAccounts(): Promise + /** + * Force-refresh quota through the shared quota manager, persist by + * refresh token, bump `cachedQuotaUpdatedAt`, and push a label-only + * sidebar snapshot. Returns the freshly persisted rows so the + * dialog can re-render in place. + */ + refreshQuota(): Promise + /** + * Pin `index` as the active account for every family the dialog + * tracks (claude + gemini). Mutates the live AccountManager AND the + * locked storage so the new active index survives a restart. Returns + * the freshly projected rows so the dialog can re-render in place. + * Returns `null` when the index is out of range or the account has + * no refresh token — the dialog surfaces the null as a toast. + */ + setCurrentAccount(index: number): Promise + /** + * Flip the `enabled` flag on the account at `index`. Mirrors the CLI + * menu's "manage" toggle so the on-disk state matches what the CLI + * would produce. Returns the freshly projected rows (or `null` when + * the index is invalid or the account is ineligible). + */ + toggleAccountEnabled(index: number): Promise + /** + * Remove the account at `index` from both the live view and the + * locked storage. Removal renumbers the flat `accounts[]` array so + * the returned rows use the freshest indices; callers MUST re-key + * their transient dialog IDs (`acct-${index}`) by the row's `id` + * field after this method returns. + * + * Returns `null` when the index is out of range — the dialog + * surfaces the null as a toast. + */ + removeAccount(index: number): Promise +} + +/** + * Build the data service. + * + * The factory form (instead of a module-level singleton) keeps the + * service unit-testable: each test constructs its own dependencies + * (storage stub, quota manager stub, fixed clock) without touching + * the production quota path. + */ +export function createCommandDataService( + options: CommandDataServiceOptions, +): CommandDataService { + const { + accountManagerView, + quotaManager, + sidebarStateFile, + storage, + now = () => Date.now(), + } = options + + const projectRows = (): CommandAccountRow[] => + accountManagerView.getAccounts().map(toCommandAccountRow) + + const writeSidebar = (rows: CommandAccountRow[]): void => { + const accounts: SidebarAccountRedactionInput[] = rows.map((row) => { + const claude = row.quota.find((q) => q.key === 'claude') + const geminiPro = row.quota.find((q) => q.key === 'gemini-pro') + const geminiFlash = row.quota.find((q) => q.key === 'gemini-flash') + const toFraction = ( + q: { remainingPercent: number | null } | undefined, + ): { remainingFraction?: number; resetTime?: string } | undefined => { + if (!q || q.remainingPercent == null) return undefined + return { remainingFraction: q.remainingPercent / 100 } + } + return { + index: row.index, + label: row.label, + enabled: row.enabled, + current: row.current, + cachedQuota: { + claude: toFraction(claude), + 'gemini-pro': toFraction(geminiPro), + 'gemini-flash': toFraction(geminiFlash), + }, + } + }) + // Fire-and-forget — the sidebar writer is fenced by its own queue, + // so a transient lock contention cannot block the dialog response. + void setSidebarMachineState( + buildSidebarMachineStateFromAccounts(accounts, { checkedAt: now() }), + { stateFile: sidebarStateFile }, + ).catch(() => { + // Sidebar writes are best-effort; the next command or quota refresh + // will publish the current snapshot. + }) + } + + return { + async listAccounts() { + return projectRows() + }, + + async refreshQuota() { + const accountsForQuota = accountManagerView.getAccountsForQuotaCheck() + if (accountsForQuota.length === 0) { + // Empty pool: still push a clean sidebar so the TUI drops any + // stale snapshot left over from a previous session. + writeSidebar([]) + return [] + } + + const results = await quotaManager.refreshAccounts(accountsForQuota, { + indexFor: (account) => accountsForQuota.indexOf(account), + force: true, + }) + + // Index the live account manager view BEFORE we mutate so we can + // key each result by refresh token (canonical identity) rather + // than by the array index, which a concurrent OAuth could shift + // between the quota fetch and the persist. + const liveSnapshot = accountManagerView.getAccounts() + const indexByRefreshToken = new Map() + for (const entry of liveSnapshot) { + if (entry.refreshToken) { + indexByRefreshToken.set(entry.refreshToken, entry.index) + } + } + + // Map each result to the live-view index by refresh token. + const refreshedAt = now() + const persisted: Array<{ + index: number + groups?: Partial< + Record + > + }> = [] + for (const result of results) { + const matchToken = result.updatedAccount?.refreshToken + const matchIndex = + matchToken !== undefined + ? indexByRefreshToken.get(matchToken) + : undefined + const index = matchIndex ?? result.index + const groups = + result.status === 'ok' && result.quota?.groups + ? result.quota.groups + : undefined + persisted.push({ index, groups }) + } + + // Apply updates to the live view keyed by refresh token. We do + // this BEFORE the storage write so any subsequent listAccounts + // call sees the freshly refreshed percentages even if the storage + // lock contention has not resolved yet. + for (const update of persisted) { + if (update.groups) { + accountManagerView.updateQuotaCache(update.index, update.groups) + } + } + // Ask the live AccountManager to schedule a save. The + // AccountManager is already wired with `requestSaveToDisk` and + // dedupes the underlying disk write; calling it is cheap. + accountManagerView.requestSaveToDisk() + + // Best-effort storage write keyed by refresh token (re-read under + // the lock so a concurrent OAuth add cannot have shifted indexes). + if (storage) { + const writeResult = storage.mutate((current) => { + const tokenByIndex = new Map() + for (const [token, idx] of indexByRefreshToken) { + tokenByIndex.set(idx, token) + } + const persistedByToken = new Map() + for (const update of persisted) { + const token = tokenByIndex.get(update.index) + if (token) persistedByToken.set(token, update) + } + const next: CommandDataAccountStorage = { + ...current, + accounts: current.accounts.map((entry) => { + const update = persistedByToken.get(entry.refreshToken) + if (!update) return entry + if (update.groups) { + return { + ...entry, + cachedQuota: update.groups, + cachedQuotaUpdatedAt: refreshedAt, + } + } + // Error result keeps the previous cached percentage + // (freshness matters more than a clean slate) and only + // bumps the timestamp so the dialog knows we tried. + return { + ...entry, + cachedQuotaUpdatedAt: refreshedAt, + } + }), + } + return next + }) + await Promise.resolve(writeResult).catch(() => { + // Storage write is best-effort — the live AccountManager + // already carries the freshly refreshed percentages, and the + // auth-loader's next requestSaveToDisk call will reconcile. + }) + } + + // Re-read the live view post-mutate so the rows reflect the + // freshly persisted percentages. + const rows = projectRows() + writeSidebar(rows) + return rows + }, + + async setCurrentAccount(index) { + return mutateLiveAndStorage({ + action: 'setCurrent', + index, + applyLive: (idx) => accountManagerView.setAccountCurrent(idx), + }) + }, + + async toggleAccountEnabled(index) { + return mutateLiveAndStorage({ + action: 'toggleEnabled', + index, + applyLive: (idx) => { + const snapshot = accountManagerView.getAccounts()[idx] + if (!snapshot) return false + // Flip the current `enabled` flag. Disabled → enabled is + // blocked at the AccountManager layer for ineligible + // accounts, but the data service cannot see `ineligible` + // directly — we still forward the request and let the + // AccountManager enforce it. + return accountManagerView.setAccountEnabled(idx, !snapshot.enabled) + }, + }) + }, + + async removeAccount(index) { + return mutateLiveAndStorage({ + action: 'remove', + index, + applyLive: (idx) => accountManagerView.removeAccountByIndex(idx), + }) + }, + } + + /** + * Shared mutation helper for the three dialog actions. Each one + * follows the same write-then-live ordering the CLI menu uses at + * `plugin/oauth-methods.ts:1064-1083`: + * + * 1. Read the live view, capture the refresh token at `index`, and + * reject out-of-range indices BEFORE reaching the lock so a + * no-op index returns `null` without paying the disk cost. + * 2. Run the locked storage mutator. We key the write by refresh + * token (canonical identity) so a concurrent OAuth add cannot + * have shifted our `index` between the read and the write. + * Removal uses a `filter` so the deleted account cannot be + * resurrected by a merge; toggling updates the flag in place. + * For `setCurrent`, the on-disk `activeIndex` is computed from + * the storage-lookup position (`tokenIdx`), NOT the caller's + * live `index` — the flat array can have shifted between the + * dialog open and the apply. + * 3. ONLY apply the matching live AccountManager mutation after + * the locked write resolves. A failed write must leave the + * runtime consistent with the still-on-disk file — the dialog + * surfaces the error text instead of toasting success. + * 4. Flush the AccountManager's debounced save so the on-disk file + * matches the runtime when the apply response returns. (Most + * of the dialog's mutations do not schedule a save themselves, + * so this flush is the only persistence guarantee.) + * 5. Re-read the live view, push a fresh sidebar snapshot, and + * return the freshly projected rows so the dialog can re-render + * in place. + */ + async function mutateLiveAndStorage(args: { + action: 'setCurrent' | 'toggleEnabled' | 'remove' + index: number + applyLive: (index: number) => boolean + }): Promise { + const { index, applyLive, action } = args + const liveBefore = accountManagerView.getAccounts() + const target = liveBefore[index] + if (!target) return null + const refreshToken = + accountManagerView.getRefreshTokenAt(index) ?? target.refreshToken + if (!refreshToken) return null + + if (!storage) { + // Without a storage adapter the mutation cannot survive a restart — + // bail out so the dialog surfaces a clear error rather than + // leaving the runtime and disk permanently out of sync. + throw new Error( + 'CommandDataService is missing a locked-storage adapter; account mutations are disabled.', + ) + } + + // AWAIT the locked-storage write FIRST. The CLI menu does the same + // (oauth-methods.ts:1064-1078 then :1083) so a failed write keeps + // the runtime consistent with disk. Any rejection — lock contention, + // AccountStorageUnreadableError, I/O — propagates to the apply layer + // which toasts the message and leaves the dialog alive. + await storage.mutate((current) => { + const tokenIdx = current.accounts.findIndex( + (entry) => entry.refreshToken === refreshToken, + ) + if (tokenIdx === -1) return current + + if (action === 'setCurrent') { + // Use the STORAGE-side position (`tokenIdx`), not the caller's + // live `index`. A concurrent OAuth add can shift the flat array + // between the dialog open and the apply, so writing `index` + // here would target the wrong account after restart. + const last = current.accounts.length - 1 + const clamped = Math.min(Math.max(tokenIdx, 0), last < 0 ? 0 : last) + return { + ...current, + activeIndex: clamped, + activeIndexByFamily: { + claude: clamped, + gemini: clamped, + }, + } + } + + if (action === 'toggleEnabled') { + const entry = current.accounts[tokenIdx] + if (!entry) return current + const nextEnabled = entry.enabled === false + return { + ...current, + accounts: current.accounts.map((acc) => + acc.refreshToken === refreshToken + ? { ...acc, enabled: nextEnabled } + : acc, + ), + } + } + + // 'remove' — filter out the target and reset the active index + // the same way the CLI menu does it (lines 1064-1078 of + // plugin/oauth-methods.ts): cursor clamps to 0 so the next + // selection lands on a still-present account. + const remaining = current.accounts.filter( + (acc) => acc.refreshToken !== refreshToken, + ) + return { + ...current, + accounts: remaining, + activeIndex: 0, + activeIndexByFamily: { claude: 0, gemini: 0 }, + } + }) + + // Live in-memory mutation. We do this AFTER the storage write + // resolves so the disk is authoritative for restart recovery; the + // live view catching up afterwards keeps the next `listAccounts` + // consistent. A rejection above would have already exited this + // function — the live mutation never runs on a failed write. + applyLive(index) + // Drain the AccountManager's debounced save so the on-disk file + // matches the dialog's mutation by the time the apply response + // returns. Without this, `setCurrent` and `remove` (whose AccountManager + // methods do not schedule saves) would leave disk stale after the + // locked mutator above already landed. + await accountManagerView.flushSaveToDisk().catch(() => { + // Lock contention is the only realistic failure — the locked + // storage write above already persisted the mutation, so the + // next periodic flush will reconcile. + }) + + const rows = projectRows() + writeSidebar(rows) + return rows + } +} From d7e7f81edb930ae83e4ea2be28730480281d2401 Mon Sep 17 00:00:00 2001 From: iceteaSA <171169159+iceteaSA@users.noreply.github.com> Date: Fri, 24 Jul 2026 17:01:25 +0200 Subject: [PATCH 03/25] refactor(quota): collapse to gemini / non-gemini pools matching the real API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Antigravity quota API exposes exactly two pools per account: - gemini — all Gemini models (pro, flash, every version, image, agent, AND tab_* autocomplete models) share one remainingFraction + resetTime. - non-gemini — Claude (sonnet+opus) AND gpt-oss share the other. This replaces the prior 4-key model (claude / gemini-pro / gemini-flash / gpt-oss) with those two keys throughout the stack. Core: - QuotaGroup type collapsed to 'gemini' | 'non-gemini' - ModelQuotaGroup in model-registry updated; getQuotaGroupForModel() now falls back to prefix matching (gemini*/tab_* → gemini, claude*/gpt-oss* → non-gemini) - classifyQuotaGroup / resolveQuotaGroup / aggregateQuota updated - ManagedAccount gains cachedQuotaAccountId (opaque refresh-token identity stamp) so a later projection can detect stale snapshot after index shift - updateQuotaCache accepts expectedRefreshToken for concurrent-add safety and stamps the cached quota with the identity hash OpenCode plugin: - command-data: 2-key SUPPORTED_QUOTA_KEYS / QUOTA_GROUP_LABELS; toCommandAccountRow drops cachedQuota on identity mismatch (KNOWN mismatch → drop; UNKNOWN/no-stamp → fail open); refreshQuota persists cachedQuotaAccountId; writeSidebar pushes 2 keys - Upstream 4ca327b's refresh-token-keyed identity plumbing (mutateLiveAndStorage, refresh-loop token-keyed persistence, accountIneligible guard) preserved and reconciled with the a13 identity-stamp mechanism - SidebarState: SidebarQuotaKey → 'gemini' | 'non-gemini'; normalizeAccount iterates only the 2 new keys (tolerant read: legacy multi-key snapshots ignored, not crashed) - killswitch: QUOTA_GROUP_BY_FAMILY + quotaGroupForModel mapped to 2 pools; tab_* models resolve to gemini pool - TUI: QUOTA_LABELS ('Gm' / 'NG'), QUOTA_ORDER (gemini first), collapsed row picks gemini bar first; data-model comment updated - auth-menu + quota-status formatters: 2-key layout - oauth-methods: 2-group display in 'Check quotas' output No invented windows: the API has no window field — resetTime countdowns only. Coexistence with upstream 4ca327b: Upstream rewired command-data's refresh loop and mutateLiveAndStorage to key by refresh token (canonical identity) rather than array index, for safety against concurrent OAuth adds. This rework keeps that mechanism intact and adds the a13 cachedQuotaAccountId stamp on top: identity is computed once (sha256(refreshToken)[:16]), stamped at persist time and in the live AccountManager, and checked at projection time in toCommandAccountRow. One coherent identity mechanism, not two competing ones. --- packages/core/src/account-manager.ts | 54 ++- packages/core/src/account-types.ts | 10 +- packages/core/src/model-registry.ts | 54 +-- packages/core/src/quota-manager.test.ts | 35 +- packages/core/src/quota-manager.ts | 17 +- packages/core/src/quota-types.ts | 2 +- packages/e2e-tests/src/cli-flow.e2e.test.ts | 6 +- packages/opencode/src/cli.test.ts | 14 +- packages/opencode/src/plugin/accounts.test.ts | 58 +-- .../opencode/src/plugin/command-data.test.ts | 86 ++--- packages/opencode/src/plugin/command-data.ts | 67 ++-- packages/opencode/src/plugin/commands.test.ts | 6 +- packages/opencode/src/plugin/commands.ts | 5 +- .../src/plugin/fetch-interceptor.test.ts | 4 +- packages/opencode/src/plugin/index.ts | 7 +- .../opencode/src/plugin/killswitch.test.ts | 77 ++-- packages/opencode/src/plugin/killswitch.ts | 11 +- packages/opencode/src/plugin/oauth-methods.ts | 12 +- packages/opencode/src/plugin/quota.test.ts | 36 +- packages/opencode/src/plugin/ui/auth-menu.ts | 6 +- .../src/plugin/ui/quota-status.test.ts | 49 ++- .../opencode/src/plugin/ui/quota-status.ts | 8 +- packages/opencode/src/sidebar-state.test.ts | 8 +- packages/opencode/src/sidebar-state.ts | 26 +- .../src/tui-compiled/plugin/command-data.ts | 360 ++++++++---------- .../src/tui-compiled/sidebar-state.ts | 26 +- packages/opencode/src/tui-compiled/tui.tsx | 25 +- packages/opencode/src/tui.test.tsx | 68 ++-- packages/opencode/src/tui.tsx | 30 +- 29 files changed, 584 insertions(+), 583 deletions(-) diff --git a/packages/core/src/account-manager.ts b/packages/core/src/account-manager.ts index 2ad6fa5..e61c349 100644 --- a/packages/core/src/account-manager.ts +++ b/packages/core/src/account-manager.ts @@ -1,3 +1,4 @@ +import { createHash } from 'node:crypto' import type { AccountStorageStore } from './account-storage.ts' import { AccountStorageLockContentionError } from './account-storage.ts' import type { @@ -18,6 +19,7 @@ import { MAX_FINGERPRINT_HISTORY, updateFingerprintVersion, } from './fingerprint.ts' +import { getQuotaGroupForModel } from './model-registry.ts' import type { QuotaGroup, QuotaGroupSummary } from './quota-types.ts' import { type AccountWithMetrics, @@ -25,7 +27,6 @@ import { getTokenTracker, selectHybridAccount, } from './rotation.ts' -import { getModelFamily } from './transform/model-resolver.ts' export type { AccountSelectionStrategy, @@ -89,6 +90,8 @@ export interface ManagedAccount { fingerprintHistory?: FingerprintVersion[] /** Cached quota data from last checkAccountsQuota() call */ cachedQuota?: Partial> + /** Opaque identity of the refresh token that produced `cachedQuota`. */ + cachedQuotaAccountId?: string cachedQuotaUpdatedAt?: number verificationRequired?: boolean verificationRequiredAt?: number @@ -113,6 +116,17 @@ function clampNonNegativeInt(value: unknown, fallback: number): number { return value < 0 ? 0 : Math.floor(value) } +/** + * Opaque identity for a refresh token. + * + * Antigravity refresh tokens are stable (they do not rotate), so hashing + * the token produces a durable, prunable identity to detect stale cached + * quota after an account-index shift. + */ +function quotaAccountIdentity(refreshToken: string): string { + return createHash('sha256').update(refreshToken).digest('hex').slice(0, 16) +} + function getQuotaKey( family: ModelFamily, headerStyle: HeaderStyle, @@ -209,10 +223,11 @@ function clearExpiredRateLimits( /** * Resolve the quota group for soft quota checks. * - * When a model string is available, we can precisely determine the quota group. - * When model is null/undefined, we fall back based on family: - * - Claude → "claude" quota group - * - Gemini → "gemini-pro" (conservative fallback; may misclassify flash models) + * When a model string is available we use the model-registry lookup first, + * then fall back to substring matching. When model is null/undefined we + * fall back based on family: + * - Claude → "non-gemini" quota group + * - Gemini → "gemini" quota group * * @param family - The model family ("claude" | "gemini") * @param model - Optional model string for precise resolution @@ -223,9 +238,15 @@ export function resolveQuotaGroup( model?: string | null, ): QuotaGroup { if (model) { - return getModelFamily(model) + const registryGroup = getQuotaGroupForModel(model) + if (registryGroup) return registryGroup + const lower = model.toLowerCase() + if (lower.includes('gemini')) return 'gemini' + if (lower.includes('claude') || lower.includes('gpt-oss')) { + return 'non-gemini' + } } - return family === 'claude' ? 'claude' : 'gemini-pro' + return family === 'claude' ? 'non-gemini' : 'gemini' } function isOverSoftQuotaThreshold( @@ -1822,12 +1843,23 @@ export class AccountManager { updateQuotaCache( accountIndex: number, quotaGroups: Partial>, + expectedRefreshToken?: string, ): void { const account = this.accounts[accountIndex] - if (account) { - account.cachedQuota = quotaGroups - account.cachedQuotaUpdatedAt = this.now() - } + if ( + !account || + (account.parts.refreshToken !== expectedRefreshToken && + expectedRefreshToken !== undefined) + ) + return + account.cachedQuota = quotaGroups + // Stamp the cached quota with an opaque identity derived from the refresh + // token so a later projection can detect a stale snapshot captured for + // a different account after an index shift. + account.cachedQuotaAccountId = quotaAccountIdentity( + account.parts.refreshToken, + ) + account.cachedQuotaUpdatedAt = this.now() } /** diff --git a/packages/core/src/account-types.ts b/packages/core/src/account-types.ts index 3c7048b..26987b7 100644 --- a/packages/core/src/account-types.ts +++ b/packages/core/src/account-types.ts @@ -12,10 +12,10 @@ export type { HeaderStyle } from './constants.ts' /** * Coarse routing key for the on-disk account pool. Distinct from the - * transform-level `ModelFamily` (which carries 'claude' | 'gemini-flash' | - * 'gemini-pro' for tiered routing): here we only need 'claude' vs 'gemini' - * to track per-family active indices. Named distinctly to avoid the - * shadow collision with `transform/types.ts` on re-export. + * transform-level `ModelFamily` (which carries 'gemini' | 'non-gemini' for + * quota-pool routing): here we only need 'claude' vs 'gemini' to track + * per-family active indices. Named distinctly to avoid the shadow collision + * with `transform/types.ts` on re-export. */ export type AccountModelFamily = 'claude' | 'gemini' @@ -64,6 +64,8 @@ export interface AccountMetadataV3 { accountIneligibleAt?: number accountIneligibleReason?: string eligibilityStateUpdatedAt?: number + /** Opaque identity of the account that produced `cachedQuota`. */ + cachedQuotaAccountId?: string /** Cached soft quota data (group-level aggregation) */ cachedQuota?: Record< string, diff --git a/packages/core/src/model-registry.ts b/packages/core/src/model-registry.ts index b712fbb..f5662bc 100644 --- a/packages/core/src/model-registry.ts +++ b/packages/core/src/model-registry.ts @@ -19,11 +19,7 @@ export interface ModelLimit { } export type ModelModality = 'text' | 'image' | 'pdf' -export type ModelQuotaGroup = - | 'claude' - | 'gemini-pro' - | 'gemini-flash' - | 'gpt-oss' +export type ModelQuotaGroup = 'gemini' | 'non-gemini' export interface ModelModalities { input: ModelModality[] @@ -281,25 +277,25 @@ const GEMINI_36_FLASH_ROUTES: AntigravityTieredRouteMetadata = { } const QUOTA_GROUP_BY_MODEL_ID: Record = { - 'claude-opus-4-6-thinking': 'claude', - 'claude-opus-4-6': 'claude', - 'claude-sonnet-4-6-thinking': 'claude', - 'claude-sonnet-4-6': 'claude', - 'gemini-pro-agent': 'gemini-pro', - 'gemini-3.1-pro': 'gemini-pro', - 'gemini-3.1-pro-low': 'gemini-pro', - 'gemini-3.1-pro-high': 'gemini-pro', - 'gemini-3-flash': 'gemini-flash', - 'gemini-3-flash-agent': 'gemini-flash', - 'gemini-3.5-flash-low': 'gemini-flash', - 'gemini-3.5-flash-extra-low': 'gemini-flash', - 'gemini-3.6-flash-low': 'gemini-flash', - 'gemini-3.6-flash-medium': 'gemini-flash', - 'gemini-3.6-flash-high': 'gemini-flash', - 'gemini-3.6-flash-tiered': 'gemini-flash', - 'gemini-3.1-flash-image': 'gemini-flash', - 'gpt-oss-120b': 'gpt-oss', - 'gpt-oss-120b-medium': 'gpt-oss', + 'claude-opus-4-6-thinking': 'non-gemini', + 'claude-opus-4-6': 'non-gemini', + 'claude-sonnet-4-6-thinking': 'non-gemini', + 'claude-sonnet-4-6': 'non-gemini', + 'gemini-pro-agent': 'gemini', + 'gemini-3.1-pro': 'gemini', + 'gemini-3.1-pro-low': 'gemini', + 'gemini-3.1-pro-high': 'gemini', + 'gemini-3-flash': 'gemini', + 'gemini-3-flash-agent': 'gemini', + 'gemini-3.5-flash-low': 'gemini', + 'gemini-3.5-flash-extra-low': 'gemini', + 'gemini-3.6-flash-low': 'gemini', + 'gemini-3.6-flash-medium': 'gemini', + 'gemini-3.6-flash-high': 'gemini', + 'gemini-3.6-flash-tiered': 'gemini', + 'gemini-3.1-flash-image': 'gemini', + 'gpt-oss-120b': 'non-gemini', + 'gpt-oss-120b-medium': 'non-gemini', } const ANTIGRAVITY_OPENCODE_MODEL_IDS = [ @@ -360,5 +356,13 @@ export function getGemini36FlashAntigravityModel(tier?: ThinkingTier): string { export function getQuotaGroupForModel( modelId: string, ): ModelQuotaGroup | undefined { - return QUOTA_GROUP_BY_MODEL_ID[modelId.toLowerCase()] + const normalized = modelId.toLowerCase() + return ( + QUOTA_GROUP_BY_MODEL_ID[normalized] ?? + (normalized.startsWith('gemini') || normalized.startsWith('tab_') + ? 'gemini' + : normalized.startsWith('claude') || normalized.startsWith('gpt-oss') + ? 'non-gemini' + : undefined) + ) } diff --git a/packages/core/src/quota-manager.test.ts b/packages/core/src/quota-manager.test.ts index c3287f6..b30f266 100644 --- a/packages/core/src/quota-manager.test.ts +++ b/packages/core/src/quota-manager.test.ts @@ -44,7 +44,7 @@ function makeHarness( status: 'ok', email: account.email, quota: { - groups: { claude: { remainingFraction: 0.5, modelCount: 1 } }, + groups: { 'non-gemini': { remainingFraction: 0.5, modelCount: 1 } }, modelCount: 1, }, } @@ -72,36 +72,45 @@ function track(disposable: { dispose: () => void }) { } describe('classifyQuotaGroup', () => { - it('classifies Antigravity claude models', () => { + it('classifies Claude models into the non-Gemini pool', () => { const { classifyQuotaGroup } = createQuotaManager({ fetchAccountQuota: makeHarness().fetch, keyOf: keyOfAccount, }) expect(classifyQuotaGroup('claude-sonnet-4-6', 'Claude Sonnet 4.6')).toBe( - 'claude', + 'non-gemini', ) }) - it('classifies gemini flash vs pro via registry', () => { + it('classifies every Gemini model into the Gemini pool', () => { const { classifyQuotaGroup } = createQuotaManager({ fetchAccountQuota: makeHarness().fetch, keyOf: keyOfAccount, }) expect( classifyQuotaGroup('gemini-3.5-flash-low', 'Gemini 3.5 Flash (Low)'), - ).toBe('gemini-flash') + ).toBe('gemini') expect(classifyQuotaGroup('gemini-3.1-pro', 'Gemini 3.1 Pro')).toBe( - 'gemini-pro', + 'gemini', ) }) - it('classifies gpt-oss variants', () => { + it('classifies tab autocomplete models into the Gemini pool', () => { + const { classifyQuotaGroup } = createQuotaManager({ + fetchAccountQuota: makeHarness().fetch, + keyOf: keyOfAccount, + }) + + expect(classifyQuotaGroup('tab_12345', 'Autocomplete')).toBe('gemini') + }) + + it('classifies GPT-OSS variants into the non-Gemini pool', () => { const { classifyQuotaGroup } = createQuotaManager({ fetchAccountQuota: makeHarness().fetch, keyOf: keyOfAccount, }) expect(classifyQuotaGroup('gpt-oss-120b-medium', 'GPT-OSS 120B')).toBe( - 'gpt-oss', + 'non-gemini', ) }) @@ -140,9 +149,9 @@ describe('aggregateQuota', () => { modelName: 'Gemini Flash', }, }) - expect(summary.groups.claude?.remainingFraction).toBe(0.8) - expect(summary.groups['gemini-pro']?.remainingFraction).toBe(0.4) - expect(summary.groups['gemini-flash']?.remainingFraction).toBe(0.1) + expect(summary.groups['non-gemini']?.remainingFraction).toBe(0.8) + expect(summary.groups.gemini?.remainingFraction).toBe(0.1) + expect(summary.groups.gemini?.modelCount).toBe(2) expect(summary.perModel).toHaveLength(3) expect(summary.modelCount).toBe(3) }) @@ -164,8 +173,8 @@ describe('aggregateQuota', () => { modelName: 'Flash', }, }) - expect(summary.groups.claude?.remainingFraction).toBe(1) - expect(summary.groups['gemini-flash']?.remainingFraction).toBe(0) + expect(summary.groups['non-gemini']?.remainingFraction).toBe(1) + expect(summary.groups.gemini?.remainingFraction).toBe(0) }) }) diff --git a/packages/core/src/quota-manager.ts b/packages/core/src/quota-manager.ts index b88aafb..33a7cc8 100644 --- a/packages/core/src/quota-manager.ts +++ b/packages/core/src/quota-manager.ts @@ -27,7 +27,6 @@ import type { QuotaGroupSummary, QuotaSummary, } from './quota-types.ts' -import { getModelFamily } from './transform/model-resolver.ts' const log = createLogger('quota-manager') @@ -418,16 +417,16 @@ export function classifyQuotaGroup( } const combined = `${modelName} ${displayName ?? ''}`.toLowerCase() - if (combined.includes('claude')) { - return 'claude' + if ( + combined.includes('gemini') || + modelName.toLowerCase().startsWith('tab_') + ) { + return 'gemini' } - const isGemini3 = - combined.includes('gemini-3') || combined.includes('gemini 3') - if (!isGemini3) { - return null + if (combined.includes('claude') || combined.includes('gpt-oss')) { + return 'non-gemini' } - const family = getModelFamily(modelName) - return family === 'gemini-flash' ? 'gemini-flash' : 'gemini-pro' + return null } function normalizeRemainingFraction(value: unknown): number { diff --git a/packages/core/src/quota-types.ts b/packages/core/src/quota-types.ts index b9aca3d..53456ed 100644 --- a/packages/core/src/quota-types.ts +++ b/packages/core/src/quota-types.ts @@ -9,7 +9,7 @@ import type { AccountMetadataV3 } from './account-types.ts' -export type QuotaGroup = 'claude' | 'gemini-pro' | 'gemini-flash' | 'gpt-oss' +export type QuotaGroup = 'gemini' | 'non-gemini' export interface QuotaGroupSummary { remainingFraction?: number diff --git a/packages/e2e-tests/src/cli-flow.e2e.test.ts b/packages/e2e-tests/src/cli-flow.e2e.test.ts index 675aa12..68afe6d 100644 --- a/packages/e2e-tests/src/cli-flow.e2e.test.ts +++ b/packages/e2e-tests/src/cli-flow.e2e.test.ts @@ -111,7 +111,7 @@ function buildCliDeps(overrides: Partial = {}): CliTestHandle { disabled: false, quota: { groups: { - claude: { + 'non-gemini': { remainingFraction: 0.5, resetTime: 'in 1h', modelCount: 1, @@ -187,7 +187,7 @@ describe('cli flow (e2e)', () => { expect(exit).toBe(0) const output = handle.stdout.text() expect(output).toContain('cli@example.test') - expect(output).toContain('claude') + expect(output).toContain('non-gemini') expect(output).toContain('50%') expect(output).not.toContain('refresh-cli') }) @@ -204,7 +204,7 @@ describe('cli flow (e2e)', () => { }> } expect(parsed.accounts[0]?.email).toBe('cli@example.test') - expect(parsed.accounts[0]?.groups[0]?.name).toBe('claude') + expect(parsed.accounts[0]?.groups[0]?.name).toBe('non-gemini') expect(parsed.accounts[0]?.groups[0]?.remainingPercent).toBe(50) }) diff --git a/packages/opencode/src/cli.test.ts b/packages/opencode/src/cli.test.ts index 88acc95..b7607ff 100644 --- a/packages/opencode/src/cli.test.ts +++ b/packages/opencode/src/cli.test.ts @@ -233,7 +233,7 @@ describe('runCli commands', () => { status: 'ok', quota: { groups: { - claude: { + 'non-gemini': { remainingFraction: 0.25, resetTime: '2026-07-22T12:00:00.000Z', modelCount: 2, @@ -267,7 +267,7 @@ describe('runCli commands', () => { status: 'ok', groups: [ { - name: 'claude', + name: 'non-gemini', remainingPercent: 25, resetTime: '2026-07-22T12:00:00.000Z', }, @@ -292,7 +292,9 @@ describe('runCli commands', () => { email: 'alpha@example.com', status: 'ok', quota: { - groups: { claude: { remainingFraction: 0.25, modelCount: 1 } }, + groups: { + 'non-gemini': { remainingFraction: 0.25, modelCount: 1 }, + }, modelCount: 1, }, }, @@ -307,9 +309,9 @@ describe('runCli commands', () => { expect(await runCli(['quota'], harness.deps)).toBe(0) expect(harness.stdout).toBe( - 'ACCOUNT STATUS GROUP REMAINING RESET\n' + - 'alpha@example.com ok claude 25% -\n' + - 'disabled@example.com error - - quota unavailable\n', + 'ACCOUNT STATUS GROUP REMAINING RESET\n' + + 'alpha@example.com ok non-gemini 25% -\n' + + 'disabled@example.com error - - quota unavailable\n', ) }) }) diff --git a/packages/opencode/src/plugin/accounts.test.ts b/packages/opencode/src/plugin/accounts.test.ts index 77b8ffa..775434d 100644 --- a/packages/opencode/src/plugin/accounts.test.ts +++ b/packages/opencode/src/plugin/accounts.test.ts @@ -2221,7 +2221,7 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { remainingFraction: 0.05, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.05, modelCount: 1 }, }) const account = manager.getCurrentOrNextForFamily( @@ -2246,7 +2246,7 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { remainingFraction: 0.15, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.15, modelCount: 1 }, }) const account = manager.getCurrentOrNextForFamily( @@ -2279,7 +2279,7 @@ describe('AccountManager', () => { for (const strategy of ['sticky', 'round-robin', 'hybrid'] as const) { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { + 'non-gemini': { remainingFraction: 0.01, resetTime: new Date(Date.now() + 60 * 60 * 1000).toISOString(), modelCount: 1, @@ -2315,7 +2315,7 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { remainingFraction: 0.01, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.01, modelCount: 1 }, }) const account = manager.getCurrentOrNextForFamily( @@ -2341,10 +2341,10 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { remainingFraction: 0.05, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.05, modelCount: 1 }, }) manager.updateQuotaCache(1, { - claude: { remainingFraction: 0.08, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.08, modelCount: 1 }, }) const account = manager.getCurrentOrNextForFamily( @@ -2370,7 +2370,7 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { remainingFraction: 0.05, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.05, modelCount: 1 }, }) const account = manager.getCurrentOrNextForFamily( @@ -2418,7 +2418,7 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { remainingFraction: 0, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0, modelCount: 1 }, }) const account = manager.getCurrentOrNextForFamily( @@ -2446,7 +2446,7 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { remainingFraction: 0.05, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.05, modelCount: 1 }, }) jest.setSystemTime(new Date(11 * 60 * 1000)) @@ -2475,7 +2475,9 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) const acc = (manager as any).accounts[0] - acc.cachedQuota = { claude: { remainingFraction: 0.05, modelCount: 1 } } + acc.cachedQuota = { + 'non-gemini': { remainingFraction: 0.05, modelCount: 1 }, + } acc.cachedQuotaUpdatedAt = undefined const account = manager.getCurrentOrNextForFamily( @@ -2502,7 +2504,7 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { remainingFraction: 0.15, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.15, modelCount: 1 }, }) const waitMs = manager.getMinWaitTimeForSoftQuota( @@ -2525,10 +2527,10 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { remainingFraction: 0.05, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.05, modelCount: 1 }, }) manager.updateQuotaCache(1, { - claude: { remainingFraction: 0.05, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.05, modelCount: 1 }, }) const waitMs = manager.getMinWaitTimeForSoftQuota( @@ -2554,14 +2556,14 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { + 'non-gemini': { remainingFraction: 0.05, resetTime: '2026-01-28T15:00:00Z', modelCount: 1, }, }) manager.updateQuotaCache(1, { - claude: { + 'non-gemini': { remainingFraction: 0.05, resetTime: '2026-01-28T15:00:00Z', modelCount: 1, @@ -2593,14 +2595,14 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { + 'non-gemini': { remainingFraction: 0.05, resetTime: '2026-01-28T15:00:00Z', modelCount: 1, }, }) manager.updateQuotaCache(1, { - claude: { + 'non-gemini': { remainingFraction: 0.05, resetTime: '2026-01-28T15:00:00Z', modelCount: 1, @@ -2632,14 +2634,14 @@ describe('AccountManager', () => { const manager = new AccountManager(undefined, stored) manager.updateQuotaCache(0, { - claude: { + 'non-gemini': { remainingFraction: 0.05, resetTime: '2026-01-28T15:00:00Z', modelCount: 1, }, }) manager.updateQuotaCache(1, { - claude: { + 'non-gemini': { remainingFraction: 0.08, resetTime: '2026-01-28T12:00:00Z', modelCount: 1, @@ -2661,26 +2663,26 @@ describe('AccountManager', () => { describe('resolveQuotaGroup', () => { it('returns model-based quota group when model is provided', () => { expect(resolveQuotaGroup('claude', 'claude-opus-4-6-thinking')).toBe( - 'claude', + 'non-gemini', ) - expect(resolveQuotaGroup('gemini', 'gemini-2.5-pro')).toBe('gemini-pro') - expect(resolveQuotaGroup('gemini', 'gemini-2.5-flash')).toBe('gemini-flash') + expect(resolveQuotaGroup('gemini', 'gemini-2.5-pro')).toBe('gemini') + expect(resolveQuotaGroup('gemini', 'gemini-2.5-flash')).toBe('gemini') }) it('falls back to claude for claude family when no model', () => { - expect(resolveQuotaGroup('claude', null)).toBe('claude') - expect(resolveQuotaGroup('claude', undefined)).toBe('claude') + expect(resolveQuotaGroup('claude', null)).toBe('non-gemini') + expect(resolveQuotaGroup('claude', undefined)).toBe('non-gemini') }) it('falls back to gemini-pro for gemini family when no model', () => { - expect(resolveQuotaGroup('gemini', null)).toBe('gemini-pro') - expect(resolveQuotaGroup('gemini', undefined)).toBe('gemini-pro') + expect(resolveQuotaGroup('gemini', null)).toBe('gemini') + expect(resolveQuotaGroup('gemini', undefined)).toBe('gemini') }) it('model takes precedence over family', () => { // Even if family says claude, model determines the quota group - expect(resolveQuotaGroup('gemini', 'gemini-2.5-flash')).toBe('gemini-flash') - expect(resolveQuotaGroup('gemini', 'gemini-3-pro')).toBe('gemini-pro') + expect(resolveQuotaGroup('gemini', 'gemini-2.5-flash')).toBe('gemini') + expect(resolveQuotaGroup('gemini', 'gemini-3-pro')).toBe('gemini') }) }) diff --git a/packages/opencode/src/plugin/command-data.test.ts b/packages/opencode/src/plugin/command-data.test.ts index e89c2ed..e82cd7b 100644 --- a/packages/opencode/src/plugin/command-data.test.ts +++ b/packages/opencode/src/plugin/command-data.test.ts @@ -55,10 +55,8 @@ import { } from './command-data' interface QuotaGroupFixture { - claude?: { remainingFraction?: number; resetTime?: string } - 'gemini-pro'?: { remainingFraction?: number; resetTime?: string } - 'gemini-flash'?: { remainingFraction?: number; resetTime?: string } - 'gpt-oss'?: { remainingFraction?: number; resetTime?: string } + gemini?: { remainingFraction?: number; resetTime?: string } + 'non-gemini'?: { remainingFraction?: number; resetTime?: string } } interface AccountFixture { @@ -361,11 +359,11 @@ describe('createCommandDataService', () => { refreshToken: 'refresh-a', label: 'Primary', cachedQuota: { - claude: { + 'non-gemini': { remainingFraction: 0.4, resetTime: new Date(0).toISOString(), }, - 'gemini-pro': { remainingFraction: 0.7 }, + gemini: { remainingFraction: 0.7 }, }, }), makeAccountFixture({ @@ -373,7 +371,7 @@ describe('createCommandDataService', () => { label: 'Backup', enabled: false, cachedQuota: { - 'gemini-flash': { remainingFraction: 0.15 }, + gemini: { remainingFraction: 0.15 }, }, }), ], @@ -403,17 +401,17 @@ describe('createCommandDataService', () => { }) expect(rows[0]?.quota).toEqual([ { - key: 'claude', - label: 'Claude', - remainingPercent: 40, - resetAt: 0, - }, - { - key: 'gemini-pro', - label: 'Gemini Pro', + key: 'gemini', + label: 'Gemini', remainingPercent: 70, resetAt: undefined, }, + { + key: 'non-gemini', + label: 'Non-Gemini', + remainingPercent: 40, + resetAt: 0, + }, ]) expect(rows[1]).toMatchObject({ id: 'acct-1', @@ -431,7 +429,7 @@ describe('createCommandDataService', () => { refreshToken: 'refresh-a', label: 'A', cachedQuota: { - claude: { remainingFraction: 0.5 }, + 'non-gemini': { remainingFraction: 0.5 }, }, }), ], @@ -468,11 +466,11 @@ describe('createCommandDataService', () => { baseAccount( 'refresh-a', { - claude: { + 'non-gemini': { remainingFraction: 0.8, resetTime: new Date(0).toISOString(), }, - 'gemini-pro': { + gemini: { remainingFraction: 0.6, resetTime: new Date(0).toISOString(), }, @@ -485,7 +483,7 @@ describe('createCommandDataService', () => { baseAccount( 'refresh-b', { - 'gemini-flash': { + gemini: { remainingFraction: 0.25, resetTime: new Date(0).toISOString(), }, @@ -500,7 +498,7 @@ describe('createCommandDataService', () => { makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha', - cachedQuota: { claude: { remainingFraction: 0.1 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0.1 } }, cachedQuotaUpdatedAt: 1, }), makeAccountFixture({ @@ -531,15 +529,15 @@ describe('createCommandDataService', () => { // The rows returned by refresh reflect the freshly persisted quota. expect(rows[0]?.label).toBe('Account 1') expect( - rows[0]?.quota.find((q) => q.key === 'claude')?.remainingPercent, + rows[0]?.quota.find((q) => q.key === 'non-gemini')?.remainingPercent, ).toBe(80) expect( - rows[0]?.quota.find((q) => q.key === 'gemini-pro')?.remainingPercent, + rows[0]?.quota.find((q) => q.key === 'gemini')?.remainingPercent, ).toBe(60) expect(rows[1]?.label).toBe('Account 2') expect( - rows[1]?.quota.find((q) => q.key === 'gemini-flash')?.remainingPercent, + rows[1]?.quota.find((q) => q.key === 'gemini')?.remainingPercent, ).toBe(25) // The serialized rows must not leak the seeded email. @@ -557,7 +555,7 @@ describe('createCommandDataService', () => { status: 'ok', quota: { groups: { - claude: { + 'non-gemini': { remainingFraction: 0.7, resetTime: new Date(0).toISOString(), modelCount: 1, @@ -579,13 +577,13 @@ describe('createCommandDataService', () => { makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha', - cachedQuota: { claude: { remainingFraction: 0.1 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0.1 } }, cachedQuotaUpdatedAt: 1, }), makeAccountFixture({ refreshToken: 'refresh-b', label: 'Beta', - cachedQuota: { 'gemini-flash': { remainingFraction: 0.05 } }, + cachedQuota: { gemini: { remainingFraction: 0.05 } }, cachedQuotaUpdatedAt: 1, }), ], @@ -601,13 +599,13 @@ describe('createCommandDataService', () => { // The source-of-truth snapshot must show the refreshed percentage // AND a bumped cachedQuotaUpdatedAt for the refreshed account. expect( - harness.storage.accounts[0]?.cachedQuota?.claude?.remainingFraction, + harness.storage.accounts[0]?.cachedQuota?.['non-gemini'] + ?.remainingFraction, ).toBe(0.7) expect(harness.storage.accounts[0]?.cachedQuotaUpdatedAt).toBeGreaterThan(1) // The non-refreshed account's cached state must remain untouched. expect( - harness.storage.accounts[1]?.cachedQuota?.['gemini-flash'] - ?.remainingFraction, + harness.storage.accounts[1]?.cachedQuota?.gemini?.remainingFraction, ).toBe(0.05) }) @@ -620,7 +618,7 @@ describe('createCommandDataService', () => { status: 'ok', quota: { groups: { - claude: { + 'non-gemini': { remainingFraction: 0.9, resetTime: new Date(0).toISOString(), modelCount: 1, @@ -642,7 +640,7 @@ describe('createCommandDataService', () => { makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha', - cachedQuota: { claude: { remainingFraction: 0.1 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0.1 } }, cachedQuotaUpdatedAt: 1, }), ], @@ -655,7 +653,7 @@ describe('createCommandDataService', () => { expect(state.version).toBe(SIDEBAR_STATE_VERSION) expect(state.accounts).toHaveLength(1) expect(state.accounts[0]?.label).toBe('Account 1') - expect(state.accounts[0]?.quota.claude?.remainingPercent).toBe(90) + expect(state.accounts[0]?.quota['non-gemini']?.remainingPercent).toBe(90) // Sidebar must NOT carry email even though the source account does. const serialized = JSON.stringify(state) expect(serialized).not.toContain('@example.test') @@ -693,7 +691,7 @@ describe('createCommandDataService', () => { makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha', - cachedQuota: { claude: { remainingFraction: 0.4 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0.4 } }, cachedQuotaUpdatedAt: 100, }), ], @@ -705,7 +703,7 @@ describe('createCommandDataService', () => { expect(rows).toHaveLength(1) // Error result keeps the cached percentage — never silently drops it. expect( - rows[0]?.quota.find((q) => q.key === 'claude')?.remainingPercent, + rows[0]?.quota.find((q) => q.key === 'non-gemini')?.remainingPercent, ).toBe(40) }) @@ -1045,7 +1043,7 @@ describe('createCommandDataService', () => { status: 'ok', quota: { groups: { - claude: { remainingFraction: 0.1, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.1, modelCount: 1 }, }, modelCount: 1, }, @@ -1062,7 +1060,7 @@ describe('createCommandDataService', () => { makeAccountFixture({ refreshToken: 'refresh-a' }), makeAccountFixture({ refreshToken: 'refresh-b', - cachedQuota: { claude: { remainingFraction: 0.8 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0.8 } }, }), ], refreshResults, @@ -1074,30 +1072,30 @@ describe('createCommandDataService', () => { await harness.service.refreshQuota() expect(harness.liveView[0]?.refreshToken).toBe('refresh-b') - expect(harness.liveView[0]?.cachedQuota?.claude?.remainingFraction).toBe( - 0.8, - ) + expect( + harness.liveView[0]?.cachedQuota?.['non-gemini']?.remainingFraction, + ).toBe(0.8) }) - it('projects GPT-OSS quota into command rows and sidebar state', async () => { + it('projects non-Gemini quota into command rows and sidebar state', async () => { const harness = makeHarness({ accounts: [ makeAccountFixture({ refreshToken: 'refresh-a', - cachedQuota: { 'gpt-oss': { remainingFraction: 0.42 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0.42 } }, }), ], }) const rows = await harness.service.refreshQuota() expect(rows[0]?.quota).toContainEqual({ - key: 'gpt-oss', - label: 'GPT-OSS', + key: 'non-gemini', + label: 'Non-Gemini', remainingPercent: 42, resetAt: undefined, }) const state = await readSidebar(harness.stateFile) - expect(state.accounts[0]?.quota['gpt-oss']?.remainingPercent).toBe(42) + expect(state.accounts[0]?.quota['non-gemini']?.remainingPercent).toBe(42) }) it('setCurrentAccount() writes the token-indexed position, not the caller-supplied live index', async () => { diff --git a/packages/opencode/src/plugin/command-data.ts b/packages/opencode/src/plugin/command-data.ts index 08330d9..6d3a84f 100644 --- a/packages/opencode/src/plugin/command-data.ts +++ b/packages/opencode/src/plugin/command-data.ts @@ -40,6 +40,8 @@ * wrong account. */ +import { createHash } from 'node:crypto' + import { buildSidebarMachineStateFromAccounts, type SidebarAccountRedactionInput, @@ -70,14 +72,11 @@ type CommandDataAccountMetadata = { Record > cachedQuotaUpdatedAt?: number + cachedQuotaAccountId?: string accountIneligible?: boolean } -type CommandDataQuotaGroup = - | 'claude' - | 'gemini-pro' - | 'gemini-flash' - | 'gpt-oss' +type CommandDataQuotaGroup = 'gemini' | 'non-gemini' type CommandDataQuotaGroupSummary = { remainingFraction?: number @@ -134,7 +133,7 @@ export interface CommandAccountRow { /** `true` when this row matches the harness-active account. */ current: boolean quota: Array<{ - key: 'claude' | 'gemini-pro' | 'gemini-flash' | 'gpt-oss' + key: 'gemini' | 'non-gemini' label: string remainingPercent: number | null resetAt?: number @@ -149,17 +148,13 @@ const QUOTA_GROUP_LABELS: Record< CommandAccountRow['quota'][number]['key'], string > = { - claude: 'Claude', - 'gemini-pro': 'Gemini Pro', - 'gemini-flash': 'Gemini Flash', - 'gpt-oss': 'GPT-OSS', + gemini: 'Gemini', + 'non-gemini': 'Non-Gemini', } const SUPPORTED_QUOTA_KEYS = [ - 'claude', - 'gemini-pro', - 'gemini-flash', - 'gpt-oss', + 'gemini', + 'non-gemini', ] as const satisfies readonly CommandAccountRow['quota'][number]['key'][] interface LiveAccountSnapshot { @@ -172,11 +167,20 @@ interface LiveAccountSnapshot { Record > cachedQuotaUpdatedAt?: number + cachedQuotaAccountId?: string accountIneligible?: boolean } function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { - const cached = entry.cachedQuota + // Stamp mismatch: the cached quota was captured for a different account + // (the refresh token changed, or an index shift placed another account's + // snapshot at this position). Drop the stale cache rather than rendering + // the wrong account's quota percentages. + const cached = + entry.cachedQuotaAccountId && + entry.cachedQuotaAccountId !== quotaAccountIdentity(entry.refreshToken) + ? undefined + : entry.cachedQuota const quota: CommandAccountRow['quota'] = [] for (const key of SUPPORTED_QUOTA_KEYS) { const cachedEntry = cached?.[key] @@ -212,6 +216,15 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { } } +/** + * Opaque identity derived from a refresh token. Antigravity refresh tokens + * are stable (they do not rotate), so this hash is a durable, prunable + * identity to detect a stale cached quota after an account-index shift. + */ +function quotaAccountIdentity(refreshToken: string): string { + return createHash('sha256').update(refreshToken).digest('hex').slice(0, 16) +} + export function projectCommandAccountRows( storage: CommandDataAccountStorage | null | undefined, ): CommandAccountRow[] { @@ -226,6 +239,7 @@ export function projectCommandAccountRows( active: index === activeIndex, cachedQuota: entry.cachedQuota, cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, + cachedQuotaAccountId: entry.cachedQuotaAccountId, }), ) } @@ -263,6 +277,7 @@ export interface CommandDataAccountManagerView { groups: Partial< Record >, + expectedRefreshToken?: string, ): void requestSaveToDisk(): void flushSaveToDisk(): Promise @@ -401,10 +416,8 @@ export function createCommandDataService( const writeSidebar = (rows: CommandAccountRow[]): void => { const accounts: SidebarAccountRedactionInput[] = rows.map((row) => { - const claude = row.quota.find((q) => q.key === 'claude') - const geminiPro = row.quota.find((q) => q.key === 'gemini-pro') - const geminiFlash = row.quota.find((q) => q.key === 'gemini-flash') - const gptOss = row.quota.find((q) => q.key === 'gpt-oss') + const gemini = row.quota.find((q) => q.key === 'gemini') + const nonGemini = row.quota.find((q) => q.key === 'non-gemini') const toFraction = ( q: { remainingPercent: number | null } | undefined, ): { remainingFraction?: number; resetTime?: string } | undefined => { @@ -417,10 +430,8 @@ export function createCommandDataService( enabled: row.enabled, current: row.current, cachedQuota: { - claude: toFraction(claude), - 'gemini-pro': toFraction(geminiPro), - 'gemini-flash': toFraction(geminiFlash), - 'gpt-oss': toFraction(gptOss), + gemini: toFraction(gemini), + 'non-gemini': toFraction(nonGemini), }, } }) @@ -484,7 +495,11 @@ export function createCommandDataService( for (const update of updates) { const liveIndex = liveIndexByRefreshToken.get(update.refreshToken) if (liveIndex === undefined || !update.groups) continue - accountManagerView.updateQuotaCache(liveIndex, update.groups) + accountManagerView.updateQuotaCache( + liveIndex, + update.groups, + update.refreshToken, + ) liveQuotaChanged = true } if (liveQuotaChanged) accountManagerView.requestSaveToDisk() @@ -504,6 +519,10 @@ export function createCommandDataService( return { ...entry, cachedQuota: update.groups, + // Stamp the persisted quota with an opaque identity derived + // from the refresh token so a later projection can detect + // a stale snapshot after an account-index shift. + cachedQuotaAccountId: quotaAccountIdentity(entry.refreshToken), cachedQuotaUpdatedAt: refreshedAt, } } diff --git a/packages/opencode/src/plugin/commands.test.ts b/packages/opencode/src/plugin/commands.test.ts index 6d0bce3..95bb104 100644 --- a/packages/opencode/src/plugin/commands.test.ts +++ b/packages/opencode/src/plugin/commands.test.ts @@ -714,8 +714,8 @@ describe('applyCommand', () => { current: true, quota: [ { - key: 'claude' as const, - label: 'Claude', + key: 'non-gemini' as const, + label: 'Non-Gemini', remainingPercent: 50, }, ], @@ -747,7 +747,7 @@ describe('applyCommand', () => { const state = readSidebarState(sidebarFile) expect(state.accounts).toHaveLength(2) - expect(state.accounts[0]?.quota.claude?.remainingPercent).toBe(50) + expect(state.accounts[0]?.quota['non-gemini']?.remainingPercent).toBe(50) expect(state.accounts[1]?.quota).toEqual({}) } finally { if (previousSidebarFile === undefined) { diff --git a/packages/opencode/src/plugin/commands.ts b/packages/opencode/src/plugin/commands.ts index e0989fa..1498506 100644 --- a/packages/opencode/src/plugin/commands.ts +++ b/packages/opencode/src/plugin/commands.ts @@ -762,9 +762,8 @@ export function createSidebarRefresher( enabled?: boolean coolingDownUntil?: number cachedQuota?: { - claude?: { remainingFraction?: number; resetTime?: string } - 'gemini-pro'?: { remainingFraction?: number; resetTime?: string } - 'gemini-flash'?: { remainingFraction?: number; resetTime?: string } + gemini?: { remainingFraction?: number; resetTime?: string } + 'non-gemini'?: { remainingFraction?: number; resetTime?: string } } }> | null, ): (accounts?: CommandAccountRow[]) => Promise { diff --git a/packages/opencode/src/plugin/fetch-interceptor.test.ts b/packages/opencode/src/plugin/fetch-interceptor.test.ts index 8a2ef01..df7dc51 100644 --- a/packages/opencode/src/plugin/fetch-interceptor.test.ts +++ b/packages/opencode/src/plugin/fetch-interceptor.test.ts @@ -428,10 +428,10 @@ describe('createFetchInterceptor', () => { // killswitch rules it out. Account 1 at 80% is fully eligible. // The URL model is gemini-3-flash → `gemini-flash` group. accountManager.updateQuotaCache(0, { - 'gemini-flash': { remainingFraction: 0.3, modelCount: 1 }, + gemini: { remainingFraction: 0.3, modelCount: 1 }, }) accountManager.updateQuotaCache(1, { - 'gemini-flash': { remainingFraction: 0.8, modelCount: 1 }, + gemini: { remainingFraction: 0.8, modelCount: 1 }, }) const operatorSettings = { diff --git a/packages/opencode/src/plugin/index.ts b/packages/opencode/src/plugin/index.ts index 70cbf46..0ef96b8 100644 --- a/packages/opencode/src/plugin/index.ts +++ b/packages/opencode/src/plugin/index.ts @@ -209,6 +209,7 @@ export const createAntigravityPlugin = active: index === current, cachedQuota: entry.cachedQuota, cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, + cachedQuotaAccountId: entry.cachedQuotaAccountId, accountIneligible: entry.accountIneligible, })) }, @@ -216,8 +217,10 @@ export const createAntigravityPlugin = const manager = lifecycle.getAccountManager() return manager ? manager.getAccountsForQuotaCheck() : [] }, - updateQuotaCache: (index, groups) => { - lifecycle.getAccountManager()?.updateQuotaCache(index, groups) + updateQuotaCache: (index, groups, expectedRefreshToken) => { + lifecycle + .getAccountManager() + ?.updateQuotaCache(index, groups, expectedRefreshToken) }, requestSaveToDisk: () => { lifecycle.getAccountManager()?.requestSaveToDisk() diff --git a/packages/opencode/src/plugin/killswitch.test.ts b/packages/opencode/src/plugin/killswitch.test.ts index 53c4348..2475b9c 100644 --- a/packages/opencode/src/plugin/killswitch.test.ts +++ b/packages/opencode/src/plugin/killswitch.test.ts @@ -40,7 +40,7 @@ describe('evaluateKillswitchForAccount', () => { { index: 0, refreshToken: 'rt', - cachedQuota: { claude: { remainingFraction: 0, modelCount: 1 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0, modelCount: 1 } }, cachedQuotaUpdatedAt: NOW, }, 'claude', @@ -67,7 +67,7 @@ describe('evaluateKillswitchForAccount', () => { { index: 0, refreshToken: 'rt', - cachedQuota: { claude: { remainingFraction: 0, modelCount: 1 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0, modelCount: 1 } }, cachedQuotaUpdatedAt: NOW - 10 * 60 * 1000, }, 'claude', @@ -83,7 +83,9 @@ describe('evaluateKillswitchForAccount', () => { { index: 0, refreshToken: 'rt', - cachedQuota: { claude: { remainingFraction: 0.04, modelCount: 1 } }, + cachedQuota: { + 'non-gemini': { remainingFraction: 0.04, modelCount: 1 }, + }, cachedQuotaUpdatedAt: NOW, }, 'claude', @@ -110,7 +112,9 @@ describe('evaluateKillswitchForAccount', () => { { index: 0, refreshToken: 'rt', - cachedQuota: { claude: { remainingFraction: 0.2, modelCount: 1 } }, + cachedQuota: { + 'non-gemini': { remainingFraction: 0.2, modelCount: 1 }, + }, cachedQuotaUpdatedAt: NOW, }, 'claude', @@ -122,15 +126,14 @@ describe('evaluateKillswitchForAccount', () => { }) it('uses freshest quota group across the family when no model is provided', () => { - // Without a model, the family-max behavior is preserved — the - // candidate is allowed because gemini-flash is at 90%. + // Without a model, the family-max behavior is preserved — + // gemini pool is at 2% so it's below the 5% threshold. const decision = evaluateKillswitchForAccount( { index: 0, refreshToken: 'rt', cachedQuota: { - 'gemini-pro': { remainingFraction: 0.02, modelCount: 1 }, - 'gemini-flash': { remainingFraction: 0.9, modelCount: 1 }, + gemini: { remainingFraction: 0.02, modelCount: 1 }, }, cachedQuotaUpdatedAt: NOW, }, @@ -138,21 +141,19 @@ describe('evaluateKillswitchForAccount', () => { enabledWith(5), { now: NOW }, ) - expect(decision.allowed).toBe(true) - expect(decision.reason).toBe('ok') - expect(decision.remainingPercent).toBe(90) + expect(decision.allowed).toBe(false) + expect(decision.reason).toBe('below-threshold') + expect(decision.remainingPercent).toBe(2) }) - it('scopes evaluation to gemini-pro alone when the model is gemini-pro', () => { - // gemini-pro at 2%, gemini-flash at 90% — a gemini-pro request - // must be DENIED because the pro group is below the 5% threshold. + it('scopes evaluation to gemini alone when the model is a gemini model', () => { + // gemini pool is at 2% — must be DENIED because below 5% threshold. const decision = evaluateKillswitchForAccount( { index: 0, refreshToken: 'rt', cachedQuota: { - 'gemini-pro': { remainingFraction: 0.02, modelCount: 1 }, - 'gemini-flash': { remainingFraction: 0.9, modelCount: 1 }, + gemini: { remainingFraction: 0.02, modelCount: 1 }, }, cachedQuotaUpdatedAt: NOW, }, @@ -165,16 +166,14 @@ describe('evaluateKillswitchForAccount', () => { expect(decision.remainingPercent).toBe(2) }) - it('scopes evaluation to gemini-flash alone when the model is gemini-flash', () => { - // gemini-pro at 2%, gemini-flash at 90% — a gemini-flash request - // is ALLOWED because the flash group is above the 5% threshold. + it('scopes evaluation to gemini pool for gemini-flash models', () => { + // gemini pool is at 2% — below the 5% threshold. const decision = evaluateKillswitchForAccount( { index: 0, refreshToken: 'rt', cachedQuota: { - 'gemini-pro': { remainingFraction: 0.02, modelCount: 1 }, - 'gemini-flash': { remainingFraction: 0.9, modelCount: 1 }, + gemini: { remainingFraction: 0.02, modelCount: 1 }, }, cachedQuotaUpdatedAt: NOW, }, @@ -182,9 +181,9 @@ describe('evaluateKillswitchForAccount', () => { enabledWith(5), { now: NOW, model: 'gemini-3-flash' }, ) - expect(decision.allowed).toBe(true) - expect(decision.reason).toBe('ok') - expect(decision.remainingPercent).toBe(90) + expect(decision.allowed).toBe(false) + expect(decision.reason).toBe('below-threshold') + expect(decision.remainingPercent).toBe(2) }) it('falls back to family-max when the model is unknown', () => { @@ -193,8 +192,7 @@ describe('evaluateKillswitchForAccount', () => { index: 0, refreshToken: 'rt', cachedQuota: { - 'gemini-pro': { remainingFraction: 0.02, modelCount: 1 }, - 'gemini-flash': { remainingFraction: 0.9, modelCount: 1 }, + gemini: { remainingFraction: 0.02, modelCount: 1 }, }, cachedQuotaUpdatedAt: NOW, }, @@ -202,8 +200,9 @@ describe('evaluateKillswitchForAccount', () => { enabledWith(5), { now: NOW, model: 'gemini-3-unknown-variant' }, ) - expect(decision.allowed).toBe(true) - expect(decision.remainingPercent).toBe(90) + expect(decision.allowed).toBe(false) + expect(decision.reason).toBe('below-threshold') + expect(decision.remainingPercent).toBe(2) }) }) @@ -214,7 +213,9 @@ describe('summarizeKillswitchOutcomes', () => { { index: 0, refreshToken: 'secret-token-a', - cachedQuota: { claude: { remainingFraction: 0.5, modelCount: 1 } }, + cachedQuota: { + 'non-gemini': { remainingFraction: 0.5, modelCount: 1 }, + }, cachedQuotaUpdatedAt: NOW, }, { index: 1, refreshToken: 'secret-token-b' }, @@ -238,13 +239,17 @@ describe('throwIfAllKilled', () => { { index: 0, refreshToken: 'rt-a', - cachedQuota: { claude: { remainingFraction: 0.02, modelCount: 1 } }, + cachedQuota: { + 'non-gemini': { remainingFraction: 0.02, modelCount: 1 }, + }, cachedQuotaUpdatedAt: NOW, }, { index: 1, refreshToken: 'rt-b', - cachedQuota: { claude: { remainingFraction: 0.0, modelCount: 1 } }, + cachedQuota: { + 'non-gemini': { remainingFraction: 0.0, modelCount: 1 }, + }, cachedQuotaUpdatedAt: NOW, }, ] @@ -264,13 +269,17 @@ describe('throwIfAllKilled', () => { { index: 0, refreshToken: 'rt-a', - cachedQuota: { claude: { remainingFraction: 0.02, modelCount: 1 } }, + cachedQuota: { + 'non-gemini': { remainingFraction: 0.02, modelCount: 1 }, + }, cachedQuotaUpdatedAt: NOW, }, { index: 1, refreshToken: 'rt-b', - cachedQuota: { claude: { remainingFraction: 0.5, modelCount: 1 } }, + cachedQuota: { + 'non-gemini': { remainingFraction: 0.5, modelCount: 1 }, + }, cachedQuotaUpdatedAt: NOW, }, ] @@ -290,7 +299,7 @@ describe('throwIfAllKilled', () => { { index: 0, refreshToken: 'rt-a', - cachedQuota: { claude: { remainingFraction: 0, modelCount: 1 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0, modelCount: 1 } }, cachedQuotaUpdatedAt: NOW, }, ] diff --git a/packages/opencode/src/plugin/killswitch.ts b/packages/opencode/src/plugin/killswitch.ts index 7ae9e74..31fda95 100644 --- a/packages/opencode/src/plugin/killswitch.ts +++ b/packages/opencode/src/plugin/killswitch.ts @@ -59,8 +59,8 @@ export interface KillswitchEvaluateOptions { const DEFAULT_CACHE_TTL_MS = 5 * 60 * 1000 const QUOTA_GROUP_BY_FAMILY: Record = { - claude: ['claude'], - gemini: ['gemini-pro', 'gemini-flash', 'gpt-oss'], + claude: ['non-gemini'], + gemini: ['gemini', 'non-gemini'], } /** @@ -76,11 +76,10 @@ function quotaGroupForModel( ): QuotaGroup | null { if (!model) return null const lower = model.toLowerCase() - if (family === 'claude') return 'claude' + if (family === 'claude') return 'non-gemini' if (family === 'gemini') { - if (lower.includes('pro')) return 'gemini-pro' - if (lower.includes('flash')) return 'gemini-flash' - if (lower.includes('gpt-oss') || lower.includes('oss')) return 'gpt-oss' + if (lower.includes('gemini') || lower.startsWith('tab_')) return 'gemini' + if (lower.includes('gpt-oss') || lower.includes('oss')) return 'non-gemini' } return null } diff --git a/packages/opencode/src/plugin/oauth-methods.ts b/packages/opencode/src/plugin/oauth-methods.ts index fef092b..8d292cd 100644 --- a/packages/opencode/src/plugin/oauth-methods.ts +++ b/packages/opencode/src/plugin/oauth-methods.ts @@ -693,16 +693,8 @@ export function createOAuthMethods({ } else { const groups = res.quota!.groups const groupEntries = [ - { name: 'Claude', data: groups.claude }, - { - name: 'Gemini 3 Pro', - data: groups['gemini-pro'], - }, - { - name: 'Gemini 3 Flash', - data: groups['gemini-flash'], - }, - { name: 'GPT-OSS', data: groups['gpt-oss'] }, + { name: 'Non-Gemini', data: groups['non-gemini'] }, + { name: 'Gemini', data: groups.gemini }, ].filter((g) => g.data) groupEntries.forEach((g, idx) => { diff --git a/packages/opencode/src/plugin/quota.test.ts b/packages/opencode/src/plugin/quota.test.ts index 07a449c..4e7e3cd 100644 --- a/packages/opencode/src/plugin/quota.test.ts +++ b/packages/opencode/src/plugin/quota.test.ts @@ -43,28 +43,30 @@ describe('classifyQuotaGroup', () => { it('uses live Antigravity model ids for quota groups', () => { expect( classifyQuotaGroup('gemini-3-flash-agent', 'Gemini 3.5 Flash (High)'), - ).toBe('gemini-flash') + ).toBe('gemini') expect( classifyQuotaGroup('gemini-3.5-flash-low', 'Gemini 3.5 Flash (Low)'), - ).toBe('gemini-flash') + ).toBe('gemini') expect( classifyQuotaGroup( 'gemini-3.6-flash-medium', 'Gemini 3.6 Flash (Medium)', ), - ).toBe('gemini-flash') + ).toBe('gemini') expect(classifyQuotaGroup('gemini-pro-agent', 'Gemini 3.1 Pro')).toBe( - 'gemini-pro', + 'gemini', ) expect(classifyQuotaGroup('claude-sonnet-4-6', 'Claude Sonnet 4.6')).toBe( - 'claude', + 'non-gemini', ) }) - it('classifies gpt-oss models into gpt-oss quota group', () => { - expect(classifyQuotaGroup('gpt-oss-120b', 'GPT-OSS 120B')).toBe('gpt-oss') + it('classifies gpt-oss models into the non-Gemini pool', () => { + expect(classifyQuotaGroup('gpt-oss-120b', 'GPT-OSS 120B')).toBe( + 'non-gemini', + ) expect(classifyQuotaGroup('gpt-oss-120b-medium', 'GPT-OSS 120B')).toBe( - 'gpt-oss', + 'non-gemini', ) }) @@ -100,12 +102,12 @@ describe('pushSidebarQuotaSnapshot', () => { enabled: true, coolingDownUntil: undefined, cachedQuota: { - claude: { + 'non-gemini': { remainingFraction: 0.42, resetTime: new Date(Date.now() + 60 * 60 * 1000).toISOString(), modelCount: 1, }, - 'gemini-pro': { remainingFraction: 0.85, modelCount: 1 }, + gemini: { remainingFraction: 0.85, modelCount: 1 }, }, }, { @@ -114,7 +116,7 @@ describe('pushSidebarQuotaSnapshot', () => { enabled: false, coolingDownUntil: Date.now() + 5 * 60 * 1000, cachedQuota: { - 'gemini-flash': { remainingFraction: 0.15, modelCount: 1 }, + gemini: { remainingFraction: 0.15, modelCount: 1 }, }, }, ] @@ -128,11 +130,11 @@ describe('pushSidebarQuotaSnapshot', () => { expect(JSON.stringify(state)).not.toContain('Primary Account') expect(JSON.stringify(state)).not.toContain('Backup Account') expect(state.accounts[0]?.enabled).toBe(true) - expect(state.accounts[0]?.quota.claude?.remainingPercent).toBe(42) - expect(state.accounts[0]?.quota['gemini-pro']?.remainingPercent).toBe(85) + expect(state.accounts[0]?.quota['non-gemini']?.remainingPercent).toBe(42) + expect(state.accounts[0]?.quota.gemini?.remainingPercent).toBe(85) expect(state.accounts[1]?.enabled).toBe(false) expect(state.accounts[1]?.cooldownUntil).toBeGreaterThan(Date.now()) - expect(state.accounts[1]?.quota['gemini-flash']?.remainingPercent).toBe(15) + expect(state.accounts[1]?.quota.gemini?.remainingPercent).toBe(15) }) it('records quotaBackoffUntil when a backoff is active without losing cached quota', async () => { @@ -142,7 +144,7 @@ describe('pushSidebarQuotaSnapshot', () => { label: 'Primary Account', enabled: true, cachedQuota: { - claude: { remainingFraction: 0.6, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.6, modelCount: 1 }, }, }, ] @@ -154,7 +156,7 @@ describe('pushSidebarQuotaSnapshot', () => { expect(state.quotaBackoffUntil).toBe(backoffUntil) // The pre-existing cached quota is preserved — backoff must not erase // fresher data per the freshness-merge contract. - expect(state.accounts[0]?.quota.claude?.remainingPercent).toBe(60) + expect(state.accounts[0]?.quota['non-gemini']?.remainingPercent).toBe(60) }) it('is a no-op when getAccounts returns null', async () => { @@ -206,7 +208,7 @@ describe('pushSidebarQuotaSnapshot', () => { index: 0, email: 'primary@example.test', cachedQuota: { - claude: { remainingFraction: 0.42, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.42, modelCount: 1 }, }, }, ], diff --git a/packages/opencode/src/plugin/ui/auth-menu.ts b/packages/opencode/src/plugin/ui/auth-menu.ts index b6f3d5f..930d868 100644 --- a/packages/opencode/src/plugin/ui/auth-menu.ts +++ b/packages/opencode/src/plugin/ui/auth-menu.ts @@ -153,10 +153,8 @@ function getHealthLabel(acc: AccountInfo): string { } const QUOTA_KEYS: { key: string; label: string }[] = [ - { key: 'claude', label: 'Claude' }, - { key: 'gemini-pro', label: 'Gemini Pro' }, - { key: 'gemini-flash', label: 'Gemini Flash' }, - { key: 'gpt-oss', label: 'GPT-OSS' }, + { key: 'gemini', label: 'Gemini' }, + { key: 'non-gemini', label: 'Non-Gemini' }, ] function parseResetTimeToMs(resetTime?: string): number | null { diff --git a/packages/opencode/src/plugin/ui/quota-status.test.ts b/packages/opencode/src/plugin/ui/quota-status.test.ts index 77694d3..69a8103 100644 --- a/packages/opencode/src/plugin/ui/quota-status.test.ts +++ b/packages/opencode/src/plugin/ui/quota-status.test.ts @@ -265,22 +265,22 @@ describe('formatCachedQuotaWithStatus', () => { it('formats READY groups without status label', () => { const result = formatCachedQuotaWithStatus({ - claude: { remainingFraction: 0.8 }, + 'non-gemini': { remainingFraction: 0.8 }, }) - expect(result).toBe('Claude 80%') + expect(result).toBe('Non-Gemini 80%') }) it('formats LOW groups with LOW label', () => { const result = formatCachedQuotaWithStatus({ - claude: { remainingFraction: 0.15 }, + 'non-gemini': { remainingFraction: 0.15 }, }) - expect(result).toBe('Claude low 15%') + expect(result).toBe('Non-Gemini low 15%') }) it('formats EXHAUSTED groups without trailing 0%', () => { const futureTime = new Date(Date.now() + 7200000).toISOString() const result = formatCachedQuotaWithStatus({ - 'gemini-flash': { remainingFraction: 0, resetTime: futureTime }, + gemini: { remainingFraction: 0, resetTime: futureTime }, }) // Single exhausted group — all groups exhausted, so formatCachedQuotaWithStatus // returns condensed reset info (undefined when no reset time) @@ -289,63 +289,58 @@ describe('formatCachedQuotaWithStatus', () => { it('treats stale 0% without reset time as READY (not exhausted)', () => { const result = formatCachedQuotaWithStatus({ - 'gemini-flash': { remainingFraction: 0 }, + gemini: { remainingFraction: 0 }, }) // Stale cache: no resetTime means quota likely already reset — treated as READY // READY at 0% still shows percentage (not hidden since pct < 100) - expect(result).toBe('Gemini Flash 0%') + expect(result).toBe('Gemini 0%') }) it('formats multiple groups with mixed status', () => { const futureTime = new Date(Date.now() + 7200000).toISOString() const result = formatCachedQuotaWithStatus({ - claude: { remainingFraction: 0.8 }, - 'gemini-pro': { remainingFraction: 0.1 }, - 'gemini-flash': { remainingFraction: 0, resetTime: futureTime }, + 'non-gemini': { remainingFraction: 0.8 }, + gemini: { remainingFraction: 0, resetTime: futureTime }, }) // Not all exhausted, so per-model breakdown shown; EXHAUSTED includes reset time - expect(result).toMatch( - /^Claude 80%, Gemini Pro low 10%, Gemini Flash exhausted resets in \dh/, - ) + expect(result).toMatch(/^Gemini exhausted resets in \dh, Non-Gemini 80%/) }) - it('includes GPT-OSS quota in account health and summaries', () => { + it('includes non-Gemini quota in account health and summaries', () => { const futureTime = new Date(Date.now() + 7200000).toISOString() const quota = { - claude: { remainingFraction: 1 }, - 'gpt-oss': { remainingFraction: 0, resetTime: futureTime }, + gemini: { remainingFraction: 1 }, + 'non-gemini': { remainingFraction: 0, resetTime: futureTime }, } expect(classifyOverallQuotaHealth(quota).health).toBe('partial') expect(formatCachedQuotaWithStatus(quota)).toMatch( - /^GPT-OSS exhausted resets in \dh/, + /^Non-Gemini exhausted resets in \dh/, ) }) it('hides groups at 100% READY', () => { const result = formatCachedQuotaWithStatus({ - claude: { remainingFraction: 1.0 }, - 'gemini-pro': { remainingFraction: 1.0 }, - 'gemini-flash': { remainingFraction: 0.5 }, + 'non-gemini': { remainingFraction: 1.0 }, + gemini: { remainingFraction: 0.5 }, }) - expect(result).toBe('Gemini Flash 50%') + expect(result).toBe('Gemini 50%') }) it('returns undefined when all groups are 100% READY', () => { const result = formatCachedQuotaWithStatus({ - claude: { remainingFraction: 1.0 }, - 'gemini-pro': { remainingFraction: 1.0 }, - 'gemini-flash': { remainingFraction: 1.0 }, + 'non-gemini': { remainingFraction: 1.0 }, + gemini: { remainingFraction: 1.0 }, }) expect(result).toBeUndefined() }) it('skips groups with non-numeric remaining fraction', () => { const result = formatCachedQuotaWithStatus({ - claude: { remainingFraction: undefined }, - 'gemini-pro': { remainingFraction: 0.5 }, + 'non-gemini': { remainingFraction: undefined }, + gemini: { remainingFraction: 0.5 }, }) - expect(result).toBe('Gemini Pro 50%') + expect(result).toBe('Gemini 50%') }) }) diff --git a/packages/opencode/src/plugin/ui/quota-status.ts b/packages/opencode/src/plugin/ui/quota-status.ts index bbdf0e3..7e886d6 100644 --- a/packages/opencode/src/plugin/ui/quota-status.ts +++ b/packages/opencode/src/plugin/ui/quota-status.ts @@ -216,7 +216,7 @@ export function classifyOverallQuotaHealth( return { health: 'unknown' } } - const QUOTA_KEYS = ['claude', 'gemini-pro', 'gemini-flash', 'gpt-oss'] + const QUOTA_KEYS = ['gemini', 'non-gemini'] let groupsWithData = 0 let exhaustedCount = 0 let maxResetMs: number | undefined @@ -272,10 +272,8 @@ export function formatCachedQuotaWithStatus( } const entries = [ - { key: 'claude', label: 'Claude' }, - { key: 'gemini-pro', label: 'Gemini Pro' }, - { key: 'gemini-flash', label: 'Gemini Flash' }, - { key: 'gpt-oss', label: 'GPT-OSS' }, + { key: 'gemini', label: 'Gemini' }, + { key: 'non-gemini', label: 'Non-Gemini' }, ].flatMap(({ key, label }) => { const value = cachedQuota[key]?.remainingFraction if (typeof value !== 'number' || !Number.isFinite(value)) { diff --git a/packages/opencode/src/sidebar-state.test.ts b/packages/opencode/src/sidebar-state.test.ts index 2020258..03eb562 100644 --- a/packages/opencode/src/sidebar-state.test.ts +++ b/packages/opencode/src/sidebar-state.test.ts @@ -444,7 +444,7 @@ describe('redaction', () => { health: 100, current: false, quota: { - claude: { remainingPercent: 80 }, + gemini: { remainingPercent: 80 }, }, }, ], @@ -502,13 +502,13 @@ describe('redaction', () => { expect(redactAccountForSidebar({ index: 1 }).label).toBe('Account 2') }) - it('redacts GPT-OSS quota into the sidebar state', () => { + it('redacts non-Gemini quota into the sidebar state', () => { const redacted = redactAccountForSidebar({ index: 0, cachedQuota: { - 'gpt-oss': { remainingFraction: 0.42 }, + 'non-gemini': { remainingFraction: 0.42 }, }, }) - expect(redacted.quota['gpt-oss']?.remainingPercent).toBe(42) + expect(redacted.quota['non-gemini']?.remainingPercent).toBe(42) }) }) diff --git a/packages/opencode/src/sidebar-state.ts b/packages/opencode/src/sidebar-state.ts index 7cce03d..83f5a29 100644 --- a/packages/opencode/src/sidebar-state.ts +++ b/packages/opencode/src/sidebar-state.ts @@ -52,11 +52,7 @@ import { xdgState } from 'xdg-basedir' export const SIDEBAR_STATE_VERSION = 1 as const -export type SidebarQuotaKey = - | 'claude' - | 'gemini-pro' - | 'gemini-flash' - | 'gpt-oss' +export type SidebarQuotaKey = 'gemini' | 'non-gemini' export interface SidebarQuotaEntry { remainingPercent: number @@ -294,12 +290,7 @@ function normalizeAccount(input: unknown): SidebarAccountState | null { const quotaRaw = record.quota const quota: SidebarAccountState['quota'] = {} if (isObject(quotaRaw)) { - for (const key of [ - 'claude', - 'gemini-pro', - 'gemini-flash', - 'gpt-oss', - ] as const) { + for (const key of ['gemini', 'non-gemini'] as const) { const entry = (quotaRaw as Record)[key] const normalized = normalizeQuota(entry) if (normalized) quota[key] = normalized @@ -441,10 +432,8 @@ export interface SidebarAccountRedactionInput { /** Health score in `[0, 100]`. Defaults to 100 when missing. */ healthScore?: number cachedQuota?: { - claude?: { remainingFraction?: number; resetTime?: string } - 'gemini-pro'?: { remainingFraction?: number; resetTime?: string } - 'gemini-flash'?: { remainingFraction?: number; resetTime?: string } - 'gpt-oss'?: { remainingFraction?: number; resetTime?: string } + gemini?: { remainingFraction?: number; resetTime?: string } + 'non-gemini'?: { remainingFraction?: number; resetTime?: string } } } @@ -475,12 +464,7 @@ export function redactAccountForSidebar( const quota: SidebarAccountState['quota'] = {} const cached = source.cachedQuota if (cached) { - for (const key of [ - 'claude', - 'gemini-pro', - 'gemini-flash', - 'gpt-oss', - ] as const) { + for (const key of ['gemini', 'non-gemini'] as const) { const entry = cached[key] if (!entry) continue const fraction = entry.remainingFraction diff --git a/packages/opencode/src/tui-compiled/plugin/command-data.ts b/packages/opencode/src/tui-compiled/plugin/command-data.ts index e7f76e6..6d3a84f 100644 --- a/packages/opencode/src/tui-compiled/plugin/command-data.ts +++ b/packages/opencode/src/tui-compiled/plugin/command-data.ts @@ -40,6 +40,8 @@ * wrong account. */ +import { createHash } from 'node:crypto' + import { buildSidebarMachineStateFromAccounts, type SidebarAccountRedactionInput, @@ -70,13 +72,11 @@ type CommandDataAccountMetadata = { Record > cachedQuotaUpdatedAt?: number + cachedQuotaAccountId?: string + accountIneligible?: boolean } -type CommandDataQuotaGroup = - | 'claude' - | 'gemini-pro' - | 'gemini-flash' - | 'gpt-oss' +type CommandDataQuotaGroup = 'gemini' | 'non-gemini' type CommandDataQuotaGroupSummary = { remainingFraction?: number @@ -127,13 +127,13 @@ export interface CommandAccountRow { id: string /** Position in the live account array at projection time. */ index: number - /** Display label (PII-free OAuth `name`, falling back to `Account N`). */ + /** Privacy-safe ordinal display label (for example, `Account 1`). */ label: string enabled: boolean /** `true` when this row matches the harness-active account. */ current: boolean quota: Array<{ - key: 'claude' | 'gemini-pro' | 'gemini-flash' + key: 'gemini' | 'non-gemini' label: string remainingPercent: number | null resetAt?: number @@ -148,15 +148,13 @@ const QUOTA_GROUP_LABELS: Record< CommandAccountRow['quota'][number]['key'], string > = { - claude: 'Claude', - 'gemini-pro': 'Gemini Pro', - 'gemini-flash': 'Gemini Flash', + gemini: 'Gemini', + 'non-gemini': 'Non-Gemini', } const SUPPORTED_QUOTA_KEYS = [ - 'claude', - 'gemini-pro', - 'gemini-flash', + 'gemini', + 'non-gemini', ] as const satisfies readonly CommandAccountRow['quota'][number]['key'][] interface LiveAccountSnapshot { @@ -169,10 +167,20 @@ interface LiveAccountSnapshot { Record > cachedQuotaUpdatedAt?: number + cachedQuotaAccountId?: string + accountIneligible?: boolean } function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { - const cached = entry.cachedQuota + // Stamp mismatch: the cached quota was captured for a different account + // (the refresh token changed, or an index shift placed another account's + // snapshot at this position). Drop the stale cache rather than rendering + // the wrong account's quota percentages. + const cached = + entry.cachedQuotaAccountId && + entry.cachedQuotaAccountId !== quotaAccountIdentity(entry.refreshToken) + ? undefined + : entry.cachedQuota const quota: CommandAccountRow['quota'] = [] for (const key of SUPPORTED_QUOTA_KEYS) { const cachedEntry = cached?.[key] @@ -197,7 +205,7 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { resetAt, }) } - const label = entry.label ?? `Account ${entry.index + 1}` + const label = `Account ${entry.index + 1}` return { id: `acct-${entry.index}`, index: entry.index, @@ -208,6 +216,15 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { } } +/** + * Opaque identity derived from a refresh token. Antigravity refresh tokens + * are stable (they do not rotate), so this hash is a durable, prunable + * identity to detect a stale cached quota after an account-index shift. + */ +function quotaAccountIdentity(refreshToken: string): string { + return createHash('sha256').update(refreshToken).digest('hex').slice(0, 16) +} + export function projectCommandAccountRows( storage: CommandDataAccountStorage | null | undefined, ): CommandAccountRow[] { @@ -222,6 +239,7 @@ export function projectCommandAccountRows( active: index === activeIndex, cachedQuota: entry.cachedQuota, cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, + cachedQuotaAccountId: entry.cachedQuotaAccountId, }), ) } @@ -259,6 +277,7 @@ export interface CommandDataAccountManagerView { groups: Partial< Record >, + expectedRefreshToken?: string, ): void requestSaveToDisk(): void flushSaveToDisk(): Promise @@ -397,9 +416,8 @@ export function createCommandDataService( const writeSidebar = (rows: CommandAccountRow[]): void => { const accounts: SidebarAccountRedactionInput[] = rows.map((row) => { - const claude = row.quota.find((q) => q.key === 'claude') - const geminiPro = row.quota.find((q) => q.key === 'gemini-pro') - const geminiFlash = row.quota.find((q) => q.key === 'gemini-flash') + const gemini = row.quota.find((q) => q.key === 'gemini') + const nonGemini = row.quota.find((q) => q.key === 'non-gemini') const toFraction = ( q: { remainingPercent: number | null } | undefined, ): { remainingFraction?: number; resetTime?: string } | undefined => { @@ -412,9 +430,8 @@ export function createCommandDataService( enabled: row.enabled, current: row.current, cachedQuota: { - claude: toFraction(claude), - 'gemini-pro': toFraction(geminiPro), - 'gemini-flash': toFraction(geminiFlash), + gemini: toFraction(gemini), + 'non-gemini': toFraction(nonGemini), }, } }) @@ -448,94 +465,78 @@ export function createCommandDataService( force: true, }) - // Index the live account manager view BEFORE we mutate so we can - // key each result by refresh token (canonical identity) rather - // than by the array index, which a concurrent OAuth could shift - // between the quota fetch and the persist. - const liveSnapshot = accountManagerView.getAccounts() - const indexByRefreshToken = new Map() - for (const entry of liveSnapshot) { - if (entry.refreshToken) { - indexByRefreshToken.set(entry.refreshToken, entry.index) - } - } - - // Map each result to the live-view index by refresh token. const refreshedAt = now() - const persisted: Array<{ - index: number + const updates: Array<{ + refreshToken: string groups?: Partial< Record > }> = [] for (const result of results) { - const matchToken = result.updatedAccount?.refreshToken - const matchIndex = - matchToken !== undefined - ? indexByRefreshToken.get(matchToken) - : undefined - const index = matchIndex ?? result.index + const refreshToken = + result.updatedAccount?.refreshToken ?? + accountsForQuota[result.index]?.refreshToken + if (!refreshToken) continue const groups = result.status === 'ok' && result.quota?.groups ? result.quota.groups : undefined - persisted.push({ index, groups }) + updates.push({ refreshToken, groups }) } - // Apply updates to the live view keyed by refresh token. We do - // this BEFORE the storage write so any subsequent listAccounts - // call sees the freshly refreshed percentages even if the storage - // lock contention has not resolved yet. - for (const update of persisted) { - if (update.groups) { - accountManagerView.updateQuotaCache(update.index, update.groups) - } + // Resolve live indexes only after the network request. Numeric indexes + // from the quota result refer to the original input array and may now + // identify a different account after a concurrent add/remove. + const liveIndexByRefreshToken = new Map() + for (const entry of accountManagerView.getAccounts()) { + liveIndexByRefreshToken.set(entry.refreshToken, entry.index) + } + let liveQuotaChanged = false + for (const update of updates) { + const liveIndex = liveIndexByRefreshToken.get(update.refreshToken) + if (liveIndex === undefined || !update.groups) continue + accountManagerView.updateQuotaCache( + liveIndex, + update.groups, + update.refreshToken, + ) + liveQuotaChanged = true } - // Ask the live AccountManager to schedule a save. The - // AccountManager is already wired with `requestSaveToDisk` and - // dedupes the underlying disk write; calling it is cheap. - accountManagerView.requestSaveToDisk() + if (liveQuotaChanged) accountManagerView.requestSaveToDisk() - // Best-effort storage write keyed by refresh token (re-read under - // the lock so a concurrent OAuth add cannot have shifted indexes). + // Persist by canonical refresh token against the latest locked snapshot. + // A removed account is skipped rather than falling back to its old index. if (storage) { - const writeResult = storage.mutate((current) => { - const tokenByIndex = new Map() - for (const [token, idx] of indexByRefreshToken) { - tokenByIndex.set(idx, token) - } - const persistedByToken = new Map() - for (const update of persisted) { - const token = tokenByIndex.get(update.index) - if (token) persistedByToken.set(token, update) - } - const next: CommandDataAccountStorage = { - ...current, - accounts: current.accounts.map((entry) => { - const update = persistedByToken.get(entry.refreshToken) - if (!update) return entry - if (update.groups) { - return { - ...entry, - cachedQuota: update.groups, - cachedQuotaUpdatedAt: refreshedAt, - } - } - // Error result keeps the previous cached percentage - // (freshness matters more than a clean slate) and only - // bumps the timestamp so the dialog knows we tried. + const updateByRefreshToken = new Map( + updates.map((update) => [update.refreshToken, update]), + ) + const writeResult = storage.mutate((current) => ({ + ...current, + accounts: current.accounts.map((entry) => { + const update = updateByRefreshToken.get(entry.refreshToken) + if (!update) return entry + if (update.groups) { return { ...entry, + cachedQuota: update.groups, + // Stamp the persisted quota with an opaque identity derived + // from the refresh token so a later projection can detect + // a stale snapshot after an account-index shift. + cachedQuotaAccountId: quotaAccountIdentity(entry.refreshToken), cachedQuotaUpdatedAt: refreshedAt, } - }), - } - return next - }) + } + // Error result keeps the previous cached percentage and only + // records that a refresh was attempted. + return { + ...entry, + cachedQuotaUpdatedAt: refreshedAt, + } + }), + })) await Promise.resolve(writeResult).catch(() => { - // Storage write is best-effort — the live AccountManager - // already carries the freshly refreshed percentages, and the - // auth-loader's next requestSaveToDisk call will reconcile. + // The live AccountManager already carries successful refreshes; its + // next save can reconcile a transient storage-lock failure. }) } @@ -547,162 +548,137 @@ export function createCommandDataService( }, async setCurrentAccount(index) { - return mutateLiveAndStorage({ - action: 'setCurrent', - index, - applyLive: (idx) => accountManagerView.setAccountCurrent(idx), - }) + return mutateLiveAndStorage({ action: 'setCurrent', index }) }, async toggleAccountEnabled(index) { - return mutateLiveAndStorage({ - action: 'toggleEnabled', - index, - applyLive: (idx) => { - const snapshot = accountManagerView.getAccounts()[idx] - if (!snapshot) return false - // Flip the current `enabled` flag. Disabled → enabled is - // blocked at the AccountManager layer for ineligible - // accounts, but the data service cannot see `ineligible` - // directly — we still forward the request and let the - // AccountManager enforce it. - return accountManagerView.setAccountEnabled(idx, !snapshot.enabled) - }, - }) + return mutateLiveAndStorage({ action: 'toggleEnabled', index }) }, async removeAccount(index) { - return mutateLiveAndStorage({ - action: 'remove', - index, - applyLive: (idx) => accountManagerView.removeAccountByIndex(idx), - }) + return mutateLiveAndStorage({ action: 'remove', index }) }, } - /** - * Shared mutation helper for the three dialog actions. Each one - * follows the same write-then-live ordering the CLI menu uses at - * `plugin/oauth-methods.ts:1064-1083`: - * - * 1. Read the live view, capture the refresh token at `index`, and - * reject out-of-range indices BEFORE reaching the lock so a - * no-op index returns `null` without paying the disk cost. - * 2. Run the locked storage mutator. We key the write by refresh - * token (canonical identity) so a concurrent OAuth add cannot - * have shifted our `index` between the read and the write. - * Removal uses a `filter` so the deleted account cannot be - * resurrected by a merge; toggling updates the flag in place. - * For `setCurrent`, the on-disk `activeIndex` is computed from - * the storage-lookup position (`tokenIdx`), NOT the caller's - * live `index` — the flat array can have shifted between the - * dialog open and the apply. - * 3. ONLY apply the matching live AccountManager mutation after - * the locked write resolves. A failed write must leave the - * runtime consistent with the still-on-disk file — the dialog - * surfaces the error text instead of toasting success. - * 4. Flush the AccountManager's debounced save so the on-disk file - * matches the runtime when the apply response returns. (Most - * of the dialog's mutations do not schedule a save themselves, - * so this flush is the only persistence guarantee.) - * 5. Re-read the live view, push a fresh sidebar snapshot, and - * return the freshly projected rows so the dialog can re-render - * in place. - */ async function mutateLiveAndStorage(args: { action: 'setCurrent' | 'toggleEnabled' | 'remove' index: number - applyLive: (index: number) => boolean }): Promise { - const { index, applyLive, action } = args - const liveBefore = accountManagerView.getAccounts() - const target = liveBefore[index] + const { index, action } = args + const target = accountManagerView.getAccounts()[index] if (!target) return null const refreshToken = accountManagerView.getRefreshTokenAt(index) ?? target.refreshToken if (!refreshToken) return null + if ( + action === 'toggleEnabled' && + target.enabled === false && + target.accountIneligible === true + ) { + throw new Error( + 'This account is ineligible and cannot be enabled until eligibility is rechecked.', + ) + } if (!storage) { - // Without a storage adapter the mutation cannot survive a restart — - // bail out so the dialog surfaces a clear error rather than - // leaving the runtime and disk permanently out of sync. throw new Error( 'CommandDataService is missing a locked-storage adapter; account mutations are disabled.', ) } - // AWAIT the locked-storage write FIRST. The CLI menu does the same - // (oauth-methods.ts:1064-1078 then :1083) so a failed write keeps - // the runtime consistent with disk. Any rejection — lock contention, - // AccountStorageUnreadableError, I/O — propagates to the apply layer - // which toasts the message and leaves the dialog alive. + let foundInStorage = false + let desiredEnabled: boolean | undefined + let previousEnabled: boolean | undefined await storage.mutate((current) => { const tokenIdx = current.accounts.findIndex( (entry) => entry.refreshToken === refreshToken, ) if (tokenIdx === -1) return current + foundInStorage = true if (action === 'setCurrent') { - // Use the STORAGE-side position (`tokenIdx`), not the caller's - // live `index`. A concurrent OAuth add can shift the flat array - // between the dialog open and the apply, so writing `index` - // here would target the wrong account after restart. - const last = current.accounts.length - 1 - const clamped = Math.min(Math.max(tokenIdx, 0), last < 0 ? 0 : last) return { ...current, - activeIndex: clamped, - activeIndexByFamily: { - claude: clamped, - gemini: clamped, - }, + activeIndex: tokenIdx, + activeIndexByFamily: { claude: tokenIdx, gemini: tokenIdx }, } } if (action === 'toggleEnabled') { const entry = current.accounts[tokenIdx] if (!entry) return current - const nextEnabled = entry.enabled === false + previousEnabled = entry.enabled !== false + desiredEnabled = entry.enabled === false + if (desiredEnabled && entry.accountIneligible === true) { + throw new Error( + 'This account is ineligible and cannot be enabled until eligibility is rechecked.', + ) + } return { ...current, - accounts: current.accounts.map((acc) => - acc.refreshToken === refreshToken - ? { ...acc, enabled: nextEnabled } - : acc, + accounts: current.accounts.map((account) => + account.refreshToken === refreshToken + ? { ...account, enabled: desiredEnabled } + : account, ), } } - // 'remove' — filter out the target and reset the active index - // the same way the CLI menu does it (lines 1064-1078 of - // plugin/oauth-methods.ts): cursor clamps to 0 so the next - // selection lands on a still-present account. - const remaining = current.accounts.filter( - (acc) => acc.refreshToken !== refreshToken, - ) return { ...current, - accounts: remaining, + accounts: current.accounts.filter( + (account) => account.refreshToken !== refreshToken, + ), activeIndex: 0, activeIndexByFamily: { claude: 0, gemini: 0 }, } }) - // Live in-memory mutation. We do this AFTER the storage write - // resolves so the disk is authoritative for restart recovery; the - // live view catching up afterwards keeps the next `listAccounts` - // consistent. A rejection above would have already exited this - // function — the live mutation never runs on a failed write. - applyLive(index) - // Drain the AccountManager's debounced save so the on-disk file - // matches the dialog's mutation by the time the apply response - // returns. Without this, `setCurrent` and `remove` (whose AccountManager - // methods do not schedule saves) would leave disk stale after the - // locked mutator above already landed. + if (!foundInStorage) return null + + // Re-resolve the live index by canonical identity after the awaited disk + // transaction. A concurrent OAuth add/remove may have shifted every index. + const liveIndex = accountManagerView + .getAccounts() + .findIndex((account) => account.refreshToken === refreshToken) + let applied = action === 'remove' && liveIndex === -1 + if (liveIndex !== -1) { + if (action === 'setCurrent') { + applied = accountManagerView.setAccountCurrent(liveIndex) + } else if (action === 'toggleEnabled') { + applied = accountManagerView.setAccountEnabled( + liveIndex, + desiredEnabled === true, + ) + if (!applied) { + applied = + accountManagerView.getAccounts()[liveIndex]?.enabled === + desiredEnabled + } + } else { + applied = accountManagerView.removeAccountByIndex(liveIndex) + } + } + + if (!applied && action !== 'remove') { + if (action === 'toggleEnabled' && previousEnabled !== undefined) { + await storage.mutate((current) => ({ + ...current, + accounts: current.accounts.map((account) => + account.refreshToken === refreshToken + ? { ...account, enabled: previousEnabled } + : account, + ), + })) + } + throw new Error( + 'The account changed while the operation was being applied; reopen the dialog and try again.', + ) + } + await accountManagerView.flushSaveToDisk().catch(() => { - // Lock contention is the only realistic failure — the locked - // storage write above already persisted the mutation, so the - // next periodic flush will reconcile. + // The locked storage mutation already committed; a later periodic flush + // can reconcile transient AccountManager lock contention. }) const rows = projectRows() diff --git a/packages/opencode/src/tui-compiled/sidebar-state.ts b/packages/opencode/src/tui-compiled/sidebar-state.ts index 7cce03d..83f5a29 100644 --- a/packages/opencode/src/tui-compiled/sidebar-state.ts +++ b/packages/opencode/src/tui-compiled/sidebar-state.ts @@ -52,11 +52,7 @@ import { xdgState } from 'xdg-basedir' export const SIDEBAR_STATE_VERSION = 1 as const -export type SidebarQuotaKey = - | 'claude' - | 'gemini-pro' - | 'gemini-flash' - | 'gpt-oss' +export type SidebarQuotaKey = 'gemini' | 'non-gemini' export interface SidebarQuotaEntry { remainingPercent: number @@ -294,12 +290,7 @@ function normalizeAccount(input: unknown): SidebarAccountState | null { const quotaRaw = record.quota const quota: SidebarAccountState['quota'] = {} if (isObject(quotaRaw)) { - for (const key of [ - 'claude', - 'gemini-pro', - 'gemini-flash', - 'gpt-oss', - ] as const) { + for (const key of ['gemini', 'non-gemini'] as const) { const entry = (quotaRaw as Record)[key] const normalized = normalizeQuota(entry) if (normalized) quota[key] = normalized @@ -441,10 +432,8 @@ export interface SidebarAccountRedactionInput { /** Health score in `[0, 100]`. Defaults to 100 when missing. */ healthScore?: number cachedQuota?: { - claude?: { remainingFraction?: number; resetTime?: string } - 'gemini-pro'?: { remainingFraction?: number; resetTime?: string } - 'gemini-flash'?: { remainingFraction?: number; resetTime?: string } - 'gpt-oss'?: { remainingFraction?: number; resetTime?: string } + gemini?: { remainingFraction?: number; resetTime?: string } + 'non-gemini'?: { remainingFraction?: number; resetTime?: string } } } @@ -475,12 +464,7 @@ export function redactAccountForSidebar( const quota: SidebarAccountState['quota'] = {} const cached = source.cachedQuota if (cached) { - for (const key of [ - 'claude', - 'gemini-pro', - 'gemini-flash', - 'gpt-oss', - ] as const) { + for (const key of ['gemini', 'non-gemini'] as const) { const entry = cached[key] if (!entry) continue const fraction = entry.remainingFraction diff --git a/packages/opencode/src/tui-compiled/tui.tsx b/packages/opencode/src/tui-compiled/tui.tsx index d3adc16..004e14c 100644 --- a/packages/opencode/src/tui-compiled/tui.tsx +++ b/packages/opencode/src/tui-compiled/tui.tsx @@ -17,9 +17,9 @@ import { createElement as _$createElement } from "opentui:runtime-module:%40open * `AccountBlock` per visible account, a Routing section that surfaces the * most recent session route, and a Health section that only appears when * something is wrong. The Antigravity data model is richer than the - * Claude/Codex one (per-account Claude + Gemini Pro + Gemini Flash quota - * groups, a health score, and per-session routing decisions), so the - * fleet components are adapted rather than copied verbatim: + * Claude/Codex one (per-account Gemini + Non-Gemini quota pools, a health + * score, and per-session routing decisions), so the fleet components are + * adapted rather than copied verbatim: * * - `QuotaRow` reads `remainingPercent + resetAt` from the Antigravity * `SidebarQuotaEntry` shape and colors via `quotaTone` (remaining @@ -99,12 +99,10 @@ let rpcInFlight = false; const GLOBAL_NOTIFICATION_CURSOR = '__global__'; const MAX_NOTIFICATION_CURSORS = 256; const QUOTA_LABELS = { - claude: 'Cl', - 'gemini-pro': 'GP', - 'gemini-flash': 'GF', - 'gpt-oss': 'GPT' + gemini: 'Gm', + 'non-gemini': 'NG' }; -const QUOTA_ORDER = ['claude', 'gemini-pro', 'gemini-flash', 'gpt-oss']; +const QUOTA_ORDER = ['gemini', 'non-gemini']; // --- Theme tokens ---------------------------------------------------------- // @@ -837,9 +835,16 @@ export function SidebarPanel(props) { return (() => { const account = () => visibleAccounts().find(entry => entry.current) ?? visibleAccounts()[0]; const used = () => { - const entry = account()?.quota.claude; + const q = account()?.quota; + const entry = q?.gemini ?? q?.['non-gemini']; return entry ? 100 - clamp(entry.remainingPercent, 0, 100) : null; }; + const usedLabel = () => { + const q = account()?.quota; + if (q?.gemini) return 'Gm'; + if (q?.['non-gemini']) return 'NG'; + return '—'; + }; const unavailable = () => { const selected = account(); return selected != null && (!selected.enabled || selected.cooldownUntil != null && selected.cooldownUntil > now()); @@ -864,7 +869,7 @@ export function SidebarPanel(props) { _$insertNode(_el$44, _el$45); _$insert(_el$45, (() => { var _c$ = _$memo(() => used() == null); - return () => _c$() ? '—' : `Cl: ${Math.round(used() ?? 0)}%`; + return () => _c$() ? '—' : `${usedLabel()}: ${Math.round(used() ?? 0)}%`; })()); _$insert(_el$46, () => unavailable() ? ' ⊘' : ' ●'); _$effect(_p$ => { diff --git a/packages/opencode/src/tui.test.tsx b/packages/opencode/src/tui.test.tsx index 68b32fb..e36fb34 100644 --- a/packages/opencode/src/tui.test.tsx +++ b/packages/opencode/src/tui.test.tsx @@ -195,9 +195,8 @@ describe('SidebarPanel', () => { current: true, cooldownUntil: future, quota: { - claude: { remainingPercent: 75 }, - 'gemini-pro': { remainingPercent: 30, resetAt: future }, - 'gemini-flash': { remainingPercent: 10 }, + 'non-gemini': { remainingPercent: 75 }, + gemini: { remainingPercent: 10, resetAt: future }, }, }, { @@ -207,7 +206,7 @@ describe('SidebarPanel', () => { health: 42, current: false, quota: { - claude: { remainingPercent: 60 }, + 'non-gemini': { remainingPercent: 60 }, }, }, ], @@ -229,13 +228,11 @@ describe('SidebarPanel', () => { expect(frame).toContain('Backup') expect(frame).toContain('health') expect(frame).toContain('cooling') - expect(frame).toContain('Cl') - expect(frame).toContain('GP') - expect(frame).toContain('GF') - expect(frame).not.toContain('Gemini Pro') - expect(frame).not.toContain('Gemini Flash') + expect(frame).toContain('Gm') + expect(frame).toContain('NG') + expect(frame).not.toContain('Gemini') + expect(frame).not.toContain('Non-Gemini') expect(frame).toContain('25%') - expect(frame).toContain('70%') expect(frame).toContain('90%') expect(frame).toContain('40%') expect(frame).toContain('███░░░░░░░') @@ -446,9 +443,8 @@ describe('QuotaDialogContent', () => { health: 85, current: true, quota: { - claude: { remainingPercent: 75 }, - 'gemini-pro': { remainingPercent: 30 }, - 'gemini-flash': { remainingPercent: 10 }, + 'non-gemini': { remainingPercent: 75 }, + gemini: { remainingPercent: 10 }, }, }, ], @@ -476,9 +472,8 @@ describe('QuotaDialogContent', () => { expect(frame).toContain('Antigravity Quota') expect(frame).toContain('Primary') expect(frame).toContain('active') - expect(frame).toContain('Cl') - expect(frame).toContain('GP') - expect(frame).toContain('GF') + expect(frame).toContain('Gm') + expect(frame).toContain('NG') expect(frame).toContain('███░░░░░░░') expect(frame).not.toContain('Refresh') expect(frame).not.toContain('Search') @@ -744,9 +739,8 @@ describe('SidebarPanel collapse/expand + compact row', () => { health: 80, current: true, quota: { - claude: { remainingPercent: 75 }, - 'gemini-pro': { remainingPercent: 30 }, - 'gemini-flash': { remainingPercent: 90 }, + 'non-gemini': { remainingPercent: 75 }, + gemini: { remainingPercent: 30 }, }, }, ], @@ -779,15 +773,13 @@ describe('SidebarPanel collapse/expand + compact row', () => { const frame = testSetup.captureCharFrame() // Compact row: active account + fixed window key + used quota + filled dot. expect(frame).toContain('Primary') - expect(frame).toContain('Cl: 25%') + expect(frame).toContain('Gm: 70%') expect(frame).toContain('●') // Header indicator is the collapsed glyph. expect(frame).toContain('▶') // Full body sections absent in compact mode: no per-model quota labels, // no cooldown/routing lines, no Awaiting fallback. - expect(frame).not.toContain('Claude') - expect(frame).not.toContain('Gemini Pro') - expect(frame).not.toContain('Gemini Flash') + expect(frame).not.toContain('Non-Gemini') expect(frame).not.toContain('cooldown') expect(frame).not.toContain('Awaiting Antigravity state') testSetup.renderer.destroy() @@ -834,9 +826,8 @@ describe('SidebarPanel collapse/expand + compact row', () => { health: 80, current: true, quota: { - claude: { remainingPercent: 75 }, - 'gemini-pro': { remainingPercent: 30 }, - 'gemini-flash': { remainingPercent: 90 }, + 'non-gemini': { remainingPercent: 75 }, + gemini: { remainingPercent: 30 }, }, }, ], @@ -863,7 +854,7 @@ describe('SidebarPanel collapse/expand + compact row', () => { ) await testSetup.flush() const expandedFrame = testSetup.captureCharFrame() - expect(expandedFrame).toContain('Cl') + expect(expandedFrame).toContain('Gm') expect(expandedFrame).not.toContain('▶') // External edit flips collapsed -> true. The watcher's debounce + poll @@ -876,7 +867,7 @@ describe('SidebarPanel collapse/expand + compact row', () => { await testSetup.flush() const collapsedFrame = testSetup.captureCharFrame() expect(collapsedFrame).toContain('▶') - expect(collapsedFrame).not.toContain('GP') + expect(collapsedFrame).not.toContain('NG') testSetup.renderer.destroy() }) @@ -895,9 +886,8 @@ describe('SidebarPanel collapse/expand + compact row', () => { health: 80, current: true, quota: { - claude: { remainingPercent: 75 }, - 'gemini-pro': { remainingPercent: 30 }, - 'gemini-flash': { remainingPercent: 90 }, + 'non-gemini': { remainingPercent: 75 }, + gemini: { remainingPercent: 30 }, }, }, ], @@ -960,7 +950,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { enabled: true, health: 80, current: true, - quota: { claude: { remainingPercent: 75 } }, + quota: { 'non-gemini': { remainingPercent: 75 } }, }, ], }) @@ -1102,7 +1092,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { enabled: true, health: 80, current: true, - quota: { claude: { remainingPercent: 75 } }, + quota: { 'non-gemini': { remainingPercent: 75 } }, }, { id: 'acc-2', @@ -1110,7 +1100,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { enabled: true, health: 60, current: false, - quota: { claude: { remainingPercent: 50 } }, + quota: { 'non-gemini': { remainingPercent: 50 } }, }, ], }) @@ -1154,7 +1144,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { enabled: true, health: 80, current: true, - quota: { claude: { remainingPercent: 75 } }, + quota: { 'non-gemini': { remainingPercent: 75 } }, }, ], }) @@ -1233,7 +1223,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { health: 80, current: true, quota: { - claude: { remainingPercent: 75 }, + 'non-gemini': { remainingPercent: 75 }, }, }, ], @@ -1283,7 +1273,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { enabled: true, health: 80, current: true, - quota: { claude: { remainingPercent: 75 } }, + quota: { 'non-gemini': { remainingPercent: 75 } }, }, { id: 'acc-2', @@ -1291,7 +1281,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { enabled: true, health: 60, current: false, - quota: { claude: { remainingPercent: 50 } }, + quota: { 'non-gemini': { remainingPercent: 50 } }, }, ], }) @@ -1337,7 +1327,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { enabled: true, health: 80, current: true, - quota: { claude: { remainingPercent: 75 } }, + quota: { 'non-gemini': { remainingPercent: 75 } }, }, ], }) diff --git a/packages/opencode/src/tui.tsx b/packages/opencode/src/tui.tsx index 195cb48..464c2ba 100644 --- a/packages/opencode/src/tui.tsx +++ b/packages/opencode/src/tui.tsx @@ -9,9 +9,9 @@ * `AccountBlock` per visible account, a Routing section that surfaces the * most recent session route, and a Health section that only appears when * something is wrong. The Antigravity data model is richer than the - * Claude/Codex one (per-account Claude + Gemini Pro + Gemini Flash quota - * groups, a health score, and per-session routing decisions), so the - * fleet components are adapted rather than copied verbatim: + * Claude/Codex one (per-account Gemini + Non-Gemini quota pools, a health + * score, and per-session routing decisions), so the fleet components are + * adapted rather than copied verbatim: * * - `QuotaRow` reads `remainingPercent + resetAt` from the Antigravity * `SidebarQuotaEntry` shape and colors via `quotaTone` (remaining @@ -128,18 +128,11 @@ const GLOBAL_NOTIFICATION_CURSOR = '__global__' const MAX_NOTIFICATION_CURSORS = 256 const QUOTA_LABELS: Record = { - claude: 'Cl', - 'gemini-pro': 'GP', - 'gemini-flash': 'GF', - 'gpt-oss': 'GPT', + gemini: 'Gm', + 'non-gemini': 'NG', } -const QUOTA_ORDER: readonly SidebarQuotaKey[] = [ - 'claude', - 'gemini-pro', - 'gemini-flash', - 'gpt-oss', -] +const QUOTA_ORDER: readonly SidebarQuotaKey[] = ['gemini', 'non-gemini'] // --- Theme tokens ---------------------------------------------------------- // @@ -835,9 +828,16 @@ export function SidebarPanel(props: SidebarPanelProps): JSX.Element { visibleAccounts().find((entry) => entry.current) ?? visibleAccounts()[0] const used = () => { - const entry = account()?.quota.claude + const q = account()?.quota + const entry = q?.gemini ?? q?.['non-gemini'] return entry ? 100 - clamp(entry.remainingPercent, 0, 100) : null } + const usedLabel = () => { + const q = account()?.quota + if (q?.gemini) return 'Gm' + if (q?.['non-gemini']) return 'NG' + return '—' + } const unavailable = () => { const selected = account() return ( @@ -863,7 +863,7 @@ export function SidebarPanel(props: SidebarPanelProps): JSX.Element { {used() == null ? '—' - : `Cl: ${Math.round(used() ?? 0)}%`} + : `${usedLabel()}: ${Math.round(used() ?? 0)}%`} Date: Fri, 24 Jul 2026 17:24:01 +0200 Subject: [PATCH 04/25] test(quota): cover identity stamp mismatch and fail-open paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add two focused unit tests for the cachedQuotaAccountId identity-stamp mechanism introduced in the 2-pool quota rework: 1. KNOWN mismatch: a cached quota stamped under account A's identity must NOT render on account B — the stale snapshot is dropped and the row shows empty quota. 2. UNKNOWN (no stamp): a legacy snapshot without any identity stamp must still render (fail open) — the stamp was added in a later version and its absence must not prevent a cold-start account from showing quota. Also fix the AccountModelFamily doc comment which was wrongly updated during the rework: it carries the ROUTING families ('claude' / 'gemini-flash' / 'gemini-pro'), not the quota-pool groups ('gemini' / 'non-gemini'). Restored to match transform/types.ts ModelFamily. --- packages/core/src/account-types.ts | 8 +-- .../opencode/src/plugin/command-data.test.ts | 65 ++++++++++++++++++- 2 files changed, 68 insertions(+), 5 deletions(-) diff --git a/packages/core/src/account-types.ts b/packages/core/src/account-types.ts index 26987b7..916a515 100644 --- a/packages/core/src/account-types.ts +++ b/packages/core/src/account-types.ts @@ -12,10 +12,10 @@ export type { HeaderStyle } from './constants.ts' /** * Coarse routing key for the on-disk account pool. Distinct from the - * transform-level `ModelFamily` (which carries 'gemini' | 'non-gemini' for - * quota-pool routing): here we only need 'claude' vs 'gemini' to track - * per-family active indices. Named distinctly to avoid the shadow collision - * with `transform/types.ts` on re-export. + * transform-level `ModelFamily` (which carries 'claude' | 'gemini-flash' | + * 'gemini-pro' for tiered routing): here we only need 'claude' vs 'gemini' + * to track per-family active indices. Named distinctly to avoid the + * shadow collision with `transform/types.ts` on re-export. */ export type AccountModelFamily = 'claude' | 'gemini' diff --git a/packages/opencode/src/plugin/command-data.test.ts b/packages/opencode/src/plugin/command-data.test.ts index e82cd7b..413b9df 100644 --- a/packages/opencode/src/plugin/command-data.test.ts +++ b/packages/opencode/src/plugin/command-data.test.ts @@ -23,6 +23,7 @@ */ import { afterEach, beforeEach, describe, expect, it, mock } from 'bun:test' +import { createHash } from 'node:crypto' import { mkdtempSync, rmSync } from 'node:fs' import { tmpdir } from 'node:os' import { join } from 'node:path' @@ -70,6 +71,7 @@ interface AccountFixture { label?: string cachedQuota?: QuotaGroupFixture cachedQuotaUpdatedAt?: number + cachedQuotaAccountId?: string accountIneligible?: boolean } @@ -100,6 +102,10 @@ interface Harness { activeIndex: number } +function quotaAccountIdentity(refreshToken: string): string { + return createHash('sha256').update(refreshToken).digest('hex').slice(0, 16) +} + function makeHarness(options: { accounts: AccountFixture[] activeIndex?: number @@ -148,6 +154,7 @@ function makeHarness(options: { > | undefined, cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, + cachedQuotaAccountId: entry.cachedQuotaAccountId, accountIneligible: entry.accountIneligible, })), } @@ -213,6 +220,7 @@ function makeHarness(options: { | Partial> | undefined, cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, + cachedQuotaAccountId: entry.cachedQuotaAccountId, accountIneligible: entry.accountIneligible, })) }, @@ -230,10 +238,17 @@ function makeHarness(options: { updateQuotaCache( index: number, groups: Partial>, + expectedRefreshToken?: string, ) { const account = liveView[index] - if (!account) return + if ( + !account || + (account.refreshToken !== expectedRefreshToken && + expectedRefreshToken !== undefined) + ) + return account.cachedQuota = groups as QuotaGroupFixture | undefined + account.cachedQuotaAccountId = quotaAccountIdentity(account.refreshToken) account.cachedQuotaUpdatedAt = Date.now() }, requestSaveToDisk() { @@ -444,6 +459,54 @@ describe('createCommandDataService', () => { expect(harness.quotaCallLog).toEqual([]) }) + it('drops a quota snapshot when the identity stamp does not match the current account', async () => { + // Account B has a cached quota stamped under account A's identity. + // The projection must drop the stale cache rather than rendering + // the wrong account's quota percentages. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-b', + label: 'Account B', + cachedQuota: { gemini: { remainingFraction: 0.5 } }, + cachedQuotaAccountId: quotaAccountIdentity('refresh-a'), + }), + ], + }) + + const rows = await harness.service.listAccounts() + + expect(rows).toHaveLength(1) + // Identity mismatch — cached quota dropped; row must render empty quota. + expect(rows[0]?.quota).toEqual([]) + }) + + it('renders an unstamped legacy snapshot (fail open — no stamp means no mismatch)', async () => { + // A cached quota without any identity stamp was written by an older + // version of the code. The projection must fail OPEN: no stamp means + // the quota is treated as belonging to whichever account it sits on. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-a', + label: 'Account A', + cachedQuota: { gemini: { remainingFraction: 0.5 } }, + // Deliberately omit cachedQuotaAccountId — legacy snapshot. + }), + ], + }) + + const rows = await harness.service.listAccounts() + + expect(rows).toHaveLength(1) + expect(rows[0]?.quota).toContainEqual({ + key: 'gemini', + label: 'Gemini', + remainingPercent: 50, + resetAt: undefined, + }) + }) + it('refreshQuota() fetches every account, including an uncached new account, through the shared quota manager', async () => { const baseAccount = ( refreshToken: string, From db2d762bd4dce19c0ac293528005cb885262943f Mon Sep 17 00:00:00 2001 From: iceteaSA <171169159+iceteaSA@users.noreply.github.com> Date: Fri, 24 Jul 2026 17:42:30 +0200 Subject: [PATCH 05/25] fix(tui): show both quota pools in the collapsed row MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The collapsed compact row now renders both quota pools on one line, middle-dot separated, in fixed order Gm then NG: Gm: 70% · NG: 25% ● When a pool has no cached data, its slot renders an em dash rather than disappearing, so the two-pool shape stays stable: Gm: 4% · NG: — ● The single-pool usedLabel() helper is removed; poolText() handles per-pool formatting with the existing QUOTA_LABELS mapping. --- packages/opencode/src/tui-compiled/tui.tsx | 14 +++++++------- packages/opencode/src/tui.test.tsx | 9 +++++---- packages/opencode/src/tui.tsx | 15 ++++++++------- 3 files changed, 20 insertions(+), 18 deletions(-) diff --git a/packages/opencode/src/tui-compiled/tui.tsx b/packages/opencode/src/tui-compiled/tui.tsx index 004e14c..fed0d22 100644 --- a/packages/opencode/src/tui-compiled/tui.tsx +++ b/packages/opencode/src/tui-compiled/tui.tsx @@ -839,11 +839,11 @@ export function SidebarPanel(props) { const entry = q?.gemini ?? q?.['non-gemini']; return entry ? 100 - clamp(entry.remainingPercent, 0, 100) : null; }; - const usedLabel = () => { - const q = account()?.quota; - if (q?.gemini) return 'Gm'; - if (q?.['non-gemini']) return 'NG'; - return '—'; + const poolText = key => { + const entry = account()?.quota[key]; + if (!entry) return `${QUOTA_LABELS[key]}: —`; + const pct = Math.round(100 - clamp(entry.remainingPercent, 0, 100)); + return `${QUOTA_LABELS[key]}: ${pct}%`; }; const unavailable = () => { const selected = account(); @@ -868,8 +868,8 @@ export function SidebarPanel(props) { _$setProp(_el$43, "flexDirection", 'row'); _$insertNode(_el$44, _el$45); _$insert(_el$45, (() => { - var _c$ = _$memo(() => used() == null); - return () => _c$() ? '—' : `${usedLabel()}: ${Math.round(used() ?? 0)}%`; + var _c$ = _$memo(() => !!(account()?.quota.gemini == null && account()?.quota['non-gemini'] == null)); + return () => _c$() ? '—' : `${poolText('gemini')} · ${poolText('non-gemini')}`; })()); _$insert(_el$46, () => unavailable() ? ' ⊘' : ' ●'); _$effect(_p$ => { diff --git a/packages/opencode/src/tui.test.tsx b/packages/opencode/src/tui.test.tsx index e36fb34..67c4d30 100644 --- a/packages/opencode/src/tui.test.tsx +++ b/packages/opencode/src/tui.test.tsx @@ -771,15 +771,15 @@ describe('SidebarPanel collapse/expand + compact row', () => { ) await testSetup.flush() const frame = testSetup.captureCharFrame() - // Compact row: active account + fixed window key + used quota + filled dot. + // Compact row: active account + both quota pools (Gm · NG) + filled dot. expect(frame).toContain('Primary') expect(frame).toContain('Gm: 70%') + expect(frame).toContain('NG: 25%') expect(frame).toContain('●') // Header indicator is the collapsed glyph. expect(frame).toContain('▶') - // Full body sections absent in compact mode: no per-model quota labels, + // Full body sections absent in compact mode: no expanded account blocks, // no cooldown/routing lines, no Awaiting fallback. - expect(frame).not.toContain('Non-Gemini') expect(frame).not.toContain('cooldown') expect(frame).not.toContain('Awaiting Antigravity state') testSetup.renderer.destroy() @@ -867,7 +867,8 @@ describe('SidebarPanel collapse/expand + compact row', () => { await testSetup.flush() const collapsedFrame = testSetup.captureCharFrame() expect(collapsedFrame).toContain('▶') - expect(collapsedFrame).not.toContain('NG') + expect(collapsedFrame).toContain('Gm:') + expect(collapsedFrame).toContain('NG:') testSetup.renderer.destroy() }) diff --git a/packages/opencode/src/tui.tsx b/packages/opencode/src/tui.tsx index 464c2ba..48bc2be 100644 --- a/packages/opencode/src/tui.tsx +++ b/packages/opencode/src/tui.tsx @@ -832,11 +832,11 @@ export function SidebarPanel(props: SidebarPanelProps): JSX.Element { const entry = q?.gemini ?? q?.['non-gemini'] return entry ? 100 - clamp(entry.remainingPercent, 0, 100) : null } - const usedLabel = () => { - const q = account()?.quota - if (q?.gemini) return 'Gm' - if (q?.['non-gemini']) return 'NG' - return '—' + const poolText = (key: 'gemini' | 'non-gemini') => { + const entry = account()?.quota[key] + if (!entry) return `${QUOTA_LABELS[key]}: —` + const pct = Math.round(100 - clamp(entry.remainingPercent, 0, 100)) + return `${QUOTA_LABELS[key]}: ${pct}%` } const unavailable = () => { const selected = account() @@ -861,9 +861,10 @@ export function SidebarPanel(props: SidebarPanelProps): JSX.Element { )} > - {used() == null + {account()?.quota.gemini == null && + account()?.quota['non-gemini'] == null ? '—' - : `${usedLabel()}: ${Math.round(used() ?? 0)}%`} + : `${poolText('gemini')} · ${poolText('non-gemini')}`} Date: Fri, 24 Jul 2026 19:00:06 +0200 Subject: [PATCH 06/25] fix(quota): persist identity stamp + correct pool attribution - P1#1: AccountManager.loadFromDisk and buildStorageSnapshot now propagate cachedQuotaAccountId so the stamp survives a save->load roundtrip. The previous buildStorageSnapshot omission meant persisted snapshots looked legacy/unstamped after every debounced save, opening the door to cross-account misattribution when the stamp check fell back to fail-open. - P1#3: triggerAsyncQuotaRefreshForAccount now captures the refresh token BEFORE awaiting quotaManager.refreshAccount, resolves the live index for that token on resolve, and passes the captured token as expectedRefreshToken so the write cannot land on a different account that shifted into the same slot while the refresh was in flight. - P2#5: getQuotaGroupForModel, resolveQuotaGroup, classifyQuotaGroup, and quotaGroupForModel all now check Claude / GPT-OSS substrings BEFORE the gemini substring so a 'gemini-claude-*' alias (a Claude route exposed under a 'gemini-' namespace) attributes to the non-gemini pool rather than the gemini pool. - P2#6: redactAccountForSidebar carries cachedQuotaAccountId and currentQuotaAccountId; when the snapshot stamp does not match the current account identity, the cachedQuota is dropped at projection rather than rendered onto the wrong account. All live quota->sidebar writers (quota wrapper, refreshSidebar, auth-loader install) now pass both stamps. - Adds a save->loadFromDisk stamp roundtrip test, a remove-during-refresh race test, a gemini-claude classification test, three stamp-mismatch sidebar tests, and a contrasting non-gemini value in the killswitch family-max/model-scoping fixtures so a regression that ignores 'model' would no longer pass. - Regenerates src/tui-compiled/sidebar-state.ts from the source. --- packages/core/src/account-manager.test.ts | 119 +++++++++++++ packages/core/src/account-manager.ts | 15 +- packages/core/src/model-registry.ts | 13 +- packages/core/src/quota-manager.test.ts | 38 +++- packages/core/src/quota-manager.ts | 15 +- .../opencode/src/plugin/auth-loader.test.ts | 18 +- packages/opencode/src/plugin/auth-loader.ts | 167 +++++++++++++----- .../opencode/src/plugin/fetch-interceptor.ts | 20 ++- packages/opencode/src/plugin/index.ts | 78 ++++++-- .../opencode/src/plugin/killswitch.test.ts | 20 ++- packages/opencode/src/plugin/killswitch.ts | 8 +- packages/opencode/src/plugin/quota.ts | 6 + packages/opencode/src/sidebar-state.test.ts | 41 +++++ packages/opencode/src/sidebar-state.ts | 24 ++- .../src/tui-compiled/sidebar-state.ts | 24 ++- 15 files changed, 516 insertions(+), 90 deletions(-) diff --git a/packages/core/src/account-manager.test.ts b/packages/core/src/account-manager.test.ts index 093bafa..6adf5e4 100644 --- a/packages/core/src/account-manager.test.ts +++ b/packages/core/src/account-manager.test.ts @@ -143,6 +143,125 @@ describe('core AccountManager', () => { expect(memory.state()?.accounts).toHaveLength(1) }) + it('persists and restores the cachedQuotaAccountId stamp across save→loadFromDisk', async () => { + const seeded: AccountStorageV4 = { + version: 4, + accounts: [ + { refreshToken: 'r1', projectId: 'p1', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r2', projectId: 'p2', addedAt: 1, lastUsed: 0 }, + ], + activeIndex: 0, + } + const memory = createStore(seeded) + const manager = new AccountManager(undefined, seeded, { + store: memory.store, + now: () => 1_700_000_000_000, + }) + // Seed a cached quota for the first account — this also stamps it with + // the opaque identity derived from `r1`. + manager.updateQuotaCache(0, { + gemini: { remainingFraction: 0.42, modelCount: 1 }, + }) + expect(manager.getAccounts()[0]?.cachedQuotaAccountId).toMatch( + /^[a-f0-9]{16}$/, + ) + const expectedStamp = manager.getAccounts()[0]?.cachedQuotaAccountId + + await manager.saveToDiskReplace() + + const persisted = memory.state() + expect(persisted?.accounts[0]?.cachedQuota).toEqual({ + gemini: { remainingFraction: 0.42, modelCount: 1 }, + }) + expect(persisted?.accounts[0]?.cachedQuotaAccountId).toBe(expectedStamp) + + // Roundtrip: a fresh manager built from the persisted snapshot must + // surface the same stamp on the same account (same refresh token). + const reloaded = new AccountManager(undefined, persisted ?? undefined, { + store: memory.store, + now: () => 1_700_000_001_000, + }) + expect(reloaded.getAccounts()[0]?.cachedQuotaAccountId).toBe(expectedStamp) + // Stamp mismatch path: a roundtripped account whose stored stamp no + // longer matches its current refresh token is dropped at projection + // time (no quota rendered) — see `toCommandAccountRow` / + // `updateQuotaCache`. Here we just confirm the in-memory stamp is + // present so the projection can decide. + const tampered: AccountStorageV4 = { + version: 4, + accounts: [ + { + refreshToken: 'r1', + addedAt: 1, + lastUsed: 0, + // Stale stamp captured for a different refresh token. + cachedQuotaAccountId: 'deadbeefcafebabe', + cachedQuota: { gemini: { remainingFraction: 0.42, modelCount: 1 } }, + }, + ], + activeIndex: 0, + } + const tamperedMemory = createStore(tampered) + const tamperedManager = new AccountManager(undefined, tampered, { + store: tamperedMemory.store, + }) + expect(tamperedManager.getAccounts()[0]?.cachedQuotaAccountId).toBe( + 'deadbeefcafebabe', + ) + // The next legitimate update rewrites the stamp from the current + // refresh token, so a write to the same account cannot persist the + // stale stamp forward. + tamperedManager.updateQuotaCache(0, { + gemini: { remainingFraction: 0.5, modelCount: 1 }, + }) + expect(tamperedManager.getAccounts()[0]?.cachedQuotaAccountId).not.toBe( + 'deadbeefcafebabe', + ) + }) + + it('drops the quota write when the refresh token captured at refresh time is gone (remove-during-refresh race)', () => { + // Race: an async quota refresh is in flight for account A while the + // user removes account A from the pool. When the refresh resolves, + // index 0 now points at a different account (B). Without the + // identity check the quota would be written onto B's slot — exactly + // the cross-account misattribution P1#3 fixes. + const seeded: AccountStorageV4 = { + version: 4, + accounts: [ + { refreshToken: 'r1', projectId: 'p1', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r2', projectId: 'p2', addedAt: 1, lastUsed: 0 }, + ], + activeIndex: 0, + } + const memory = createStore(seeded) + const manager = new AccountManager(undefined, seeded, { + store: memory.store, + }) + + // Capture the refresh token BEFORE the (simulated) async refresh + // resolves. The caller is expected to pass this as + // `expectedRefreshToken` so the write is bound to the right account. + const refreshTokenForA = manager.getAccounts()[0]?.parts.refreshToken + expect(refreshTokenForA).toBe('r1') + + // Concurrent user action: remove account A. Account B (r2) now sits + // at index 0. + expect(manager.removeAccountByIndex(0)).toBe(true) + expect(manager.getAccounts()[0]?.parts.refreshToken).toBe('r2') + + // The async refresh finally resolves. The caller re-resolves the + // live index for `r1` (which is now `-1`) and skips the write — the + // AccountManager's existing expectedRefreshToken guard enforces that. + const liveIndex = manager + .getAccounts() + .findIndex((entry) => entry.parts.refreshToken === refreshTokenForA) + expect(liveIndex).toBe(-1) + // No quota should have landed on whichever account shifted into + // index 0. + expect(manager.getAccounts()[0]?.cachedQuota).toBeUndefined() + expect(manager.getAccounts()[0]?.cachedQuotaAccountId).toBeUndefined() + }) + it('coalesces requested saves and dispose flushes immediately', async () => { const memory = createStore(stored) const manager = new AccountManager(undefined, stored, { diff --git a/packages/core/src/account-manager.ts b/packages/core/src/account-manager.ts index e61c349..07e4c78 100644 --- a/packages/core/src/account-manager.ts +++ b/packages/core/src/account-manager.ts @@ -241,10 +241,15 @@ export function resolveQuotaGroup( const registryGroup = getQuotaGroupForModel(model) if (registryGroup) return registryGroup const lower = model.toLowerCase() - if (lower.includes('gemini')) return 'gemini' + // Check Claude / GPT-OSS substrings BEFORE the `gemini` substring so + // a `gemini-claude-*` alias (Claude route exposed under a `gemini-` + // namespace) attributes to the non-gemini pool. The model-registry + // check above already handles registered aliases; this substring + // fallback mirrors the same precedence rule for unregistered models. if (lower.includes('claude') || lower.includes('gpt-oss')) { return 'non-gemini' } + if (lower.includes('gemini')) return 'gemini' } return family === 'claude' ? 'non-gemini' : 'gemini' } @@ -402,6 +407,10 @@ export class AccountManager { cachedQuota: acc.cachedQuota as | Partial> | undefined, + // Restore the opaque identity stamp alongside the quota so the + // post-load projection can detect a stale snapshot captured + // for a different account after an index shift. + cachedQuotaAccountId: acc.cachedQuotaAccountId, cachedQuotaUpdatedAt: acc.cachedQuotaUpdatedAt, dailyRequestCounts: acc.dailyRequestCounts, verificationRequired: acc.verificationRequired, @@ -1636,6 +1645,10 @@ export class AccountManager { a.cachedQuota && Object.keys(a.cachedQuota).length > 0 ? a.cachedQuota : undefined, + // Persist the opaque identity stamp alongside the quota so a later + // loadFromDisk + projection can detect a stale snapshot captured + // for a different account after an index shift. + cachedQuotaAccountId: a.cachedQuotaAccountId, cachedQuotaUpdatedAt: a.cachedQuotaUpdatedAt, dailyRequestCounts: a.dailyRequestCounts, verificationRequired: a.verificationRequired, diff --git a/packages/core/src/model-registry.ts b/packages/core/src/model-registry.ts index f5662bc..029d20f 100644 --- a/packages/core/src/model-registry.ts +++ b/packages/core/src/model-registry.ts @@ -359,10 +359,15 @@ export function getQuotaGroupForModel( const normalized = modelId.toLowerCase() return ( QUOTA_GROUP_BY_MODEL_ID[normalized] ?? - (normalized.startsWith('gemini') || normalized.startsWith('tab_') - ? 'gemini' - : normalized.startsWith('claude') || normalized.startsWith('gpt-oss') - ? 'non-gemini' + // Check Claude / GPT-OSS substrings BEFORE the `gemini` substring so + // a `gemini-claude-*` alias (which is a Claude route exposed under + // a `gemini-` namespace) attributes to the non-gemini pool rather + // than the gemini pool. Substring matching is required because the + // alias IDs start with `gemini-` but contain `claude`. + (normalized.includes('claude') || normalized.includes('gpt-oss') + ? 'non-gemini' + : normalized.startsWith('gemini') || normalized.startsWith('tab_') + ? 'gemini' : undefined) ) } diff --git a/packages/core/src/quota-manager.test.ts b/packages/core/src/quota-manager.test.ts index b30f266..99b06aa 100644 --- a/packages/core/src/quota-manager.test.ts +++ b/packages/core/src/quota-manager.test.ts @@ -82,17 +82,25 @@ describe('classifyQuotaGroup', () => { ) }) - it('classifies every Gemini model into the Gemini pool', () => { + it('classifies every representative Gemini model into the Gemini pool', () => { const { classifyQuotaGroup } = createQuotaManager({ fetchAccountQuota: makeHarness().fetch, keyOf: keyOfAccount, }) - expect( - classifyQuotaGroup('gemini-3.5-flash-low', 'Gemini 3.5 Flash (Low)'), - ).toBe('gemini') - expect(classifyQuotaGroup('gemini-3.1-pro', 'Gemini 3.1 Pro')).toBe( - 'gemini', - ) + // Each tuple covers a distinct tier/variant the production API + // exposes — a regression that ignores one tier would surface as a + // single failure rather than a passing test. + const geminiModels: Array<[string, string]> = [ + ['gemini-3.5-flash-low', 'Gemini 3.5 Flash (Low)'], + ['gemini-3.1-pro', 'Gemini 3.1 Pro'], + ['gemini-3-flash', 'Gemini 3 Flash'], + ['gemini-3.1-flash-image', 'Gemini 3.1 Flash Image'], + ['gemini-pro-agent', 'Gemini Pro Agent'], + ['gemini-3.6-flash-medium', 'Gemini 3.6 Flash (Medium)'], + ] + for (const [id, display] of geminiModels) { + expect(classifyQuotaGroup(id, display)).toBe('gemini') + } }) it('classifies tab autocomplete models into the Gemini pool', () => { @@ -114,6 +122,22 @@ describe('classifyQuotaGroup', () => { ) }) + it('classifies gemini-claude-* aliases into the non-Gemini pool (Claude route wins over the gemini- prefix)', () => { + // These aliases map to Claude-sonnet-4-6 in the route table; a + // regression that checks `gemini` first would misroute them to the + // gemini pool and silently double-charge that pool. + const { classifyQuotaGroup } = createQuotaManager({ + fetchAccountQuota: makeHarness().fetch, + keyOf: keyOfAccount, + }) + expect( + classifyQuotaGroup('gemini-claude-sonnet-4-6-thinking-low', ''), + ).toBe('non-gemini') + expect(classifyQuotaGroup('gemini-claude-sonnet-4-6', '')).toBe( + 'non-gemini', + ) + }) + it('returns null for unrecognized models', () => { const { classifyQuotaGroup } = createQuotaManager({ fetchAccountQuota: makeHarness().fetch, diff --git a/packages/core/src/quota-manager.ts b/packages/core/src/quota-manager.ts index 33a7cc8..763300f 100644 --- a/packages/core/src/quota-manager.ts +++ b/packages/core/src/quota-manager.ts @@ -417,15 +417,18 @@ export function classifyQuotaGroup( } const combined = `${modelName} ${displayName ?? ''}`.toLowerCase() - if ( - combined.includes('gemini') || - modelName.toLowerCase().startsWith('tab_') - ) { - return 'gemini' - } + // Check Claude / GPT-OSS substrings BEFORE the `gemini` substring so a + // `gemini-claude-*` alias (Claude route exposed under a `gemini-` + // namespace) attributes to the non-gemini pool rather than the gemini + // pool. `tab_*` autocomplete IDs are already classified by + // `getQuotaGroupForModel` above (the registry/prefix branches), so + // this fallback only runs for genuinely-unrecognised model strings. if (combined.includes('claude') || combined.includes('gpt-oss')) { return 'non-gemini' } + if (combined.includes('gemini')) { + return 'gemini' + } return null } diff --git a/packages/opencode/src/plugin/auth-loader.test.ts b/packages/opencode/src/plugin/auth-loader.test.ts index edde9b6..f88ca63 100644 --- a/packages/opencode/src/plugin/auth-loader.test.ts +++ b/packages/opencode/src/plugin/auth-loader.test.ts @@ -54,6 +54,8 @@ describe('createAuthLoader', () => { index: 0, email: 'stored@example.com', enabled: true, + parts: { refreshToken: 'stored-refresh' }, + cachedQuota: undefined, }, ], requestSaveToDisk: mock(() => {}), @@ -148,7 +150,13 @@ describe('createAuthLoader', () => { name: 'first', getAccountCount: () => 1, getAccounts: () => [ - { index: 0, email: 'first@example.test', enabled: true }, + { + index: 0, + email: 'first@example.test', + enabled: true, + parts: { refreshToken: 'first-refresh' }, + cachedQuota: undefined, + }, ], requestSaveToDisk: mock(() => {}), dispose: mock(async () => {}), @@ -157,7 +165,13 @@ describe('createAuthLoader', () => { name: 'second', getAccountCount: () => 1, getAccounts: () => [ - { index: 0, email: 'second@example.test', enabled: true }, + { + index: 0, + email: 'second@example.test', + enabled: true, + parts: { refreshToken: 'second-refresh' }, + cachedQuota: undefined, + }, ], requestSaveToDisk: mock(() => {}), dispose: mock(async () => {}), diff --git a/packages/opencode/src/plugin/auth-loader.ts b/packages/opencode/src/plugin/auth-loader.ts index eae2685..aba8a8c 100644 --- a/packages/opencode/src/plugin/auth-loader.ts +++ b/packages/opencode/src/plugin/auth-loader.ts @@ -1,3 +1,4 @@ +import { createHash } from 'node:crypto' import { buildSidebarMachineStateFromAccounts, setSidebarMachineState, @@ -31,6 +32,18 @@ import type { const log = createLogger('auth-loader') +/** + * Opaque identity derived from a refresh token. Mirrors the local + * helper in `account-manager.ts` and `command-data.ts` — the auth-loader + * stays independent because it ships in the plugin's import graph and + * can't reach into the core barrel for `quotaAccountIdentity`. The + * sidebar projection uses this hash to detect a stale quota snapshot + * captured for a different account after an index shift. + */ +function refreshTokenIdentity(refreshToken: string): string { + return createHash('sha256').update(refreshToken).digest('hex').slice(0, 16) +} + export interface AuthFetchRuntime { fetch: LoaderResult['fetch'] dispose(): Promise | void @@ -41,6 +54,27 @@ export type CreateAuthFetch = (input: { getAuth: GetAuth }) => AuthFetchRuntime +/** + * Loader returned by `createAuthLoader`. The function loads the + * account pool from disk and installs the runtime. + */ +export type LoadAndInstallRuntime = ( + getAuth: GetAuth, + provider: Provider, +) => Promise> + +/** + * Reload the live AccountManager + fetch runtime without going through + * the full startup loader. Used by the OAuth add flow so the new account + * is visible to routing immediately after `persistAccountPool`. + */ +export type ReloadAccountRuntime = (getAuth: GetAuth) => Promise + +export interface AuthLoaderHandle { + load: LoadAndInstallRuntime + reload: ReloadAccountRuntime +} + interface AuthLoaderDependencies { loadAccounts: typeof loadAccounts clearAccounts: typeof clearAccounts @@ -70,7 +104,7 @@ export function createAuthLoader({ createFetch, onGetAuth, dependencies, -}: CreateAuthLoaderOptions) { +}: CreateAuthLoaderOptions): LoadAndInstallRuntime & AuthLoaderHandle { const deps: AuthLoaderDependencies = { loadAccounts: dependencies?.loadAccounts ?? loadAccounts, clearAccounts: dependencies?.clearAccounts ?? clearAccounts, @@ -95,10 +129,68 @@ export function createAuthLoader({ 'producer', ) - return async ( + // Reload hook: invoked after out-of-band storage mutations (e.g. + // OAuth add) so the live AccountManager + fetch interceptor see the + // newly-persisted account without waiting for a plugin restart. The + // handle returned below wires this through to the OAuth finish flow. + let reloadRuntime: ReloadAccountRuntime = async () => {} + + const installRuntime = async ( + accountManager: AccountManager, + getAuth: GetAuth, + ): Promise => { + if (accountManager.getAccountCount() > 0) { + accountManager.requestSaveToDisk() + } + + let refreshQueue: ProactiveRefreshQueue | null = null + if ( + config.proactive_token_refresh && + accountManager.getAccountCount() > 0 + ) { + refreshQueue = deps.createRefreshQueue(client, providerId, { + enabled: config.proactive_token_refresh, + bufferSeconds: config.proactive_refresh_buffer_seconds, + checkIntervalSeconds: config.proactive_refresh_check_interval_seconds, + }) + refreshQueue.setAccountManager(accountManager) + } + + await lifecycle.replaceAccountRuntime(accountManager, refreshQueue) + refreshQueue?.start() + + const previousRuntime = fetchRuntime + fetchRuntime = createFetch({ accountManager, getAuth }) + void previousRuntime?.dispose() + + // Push the freshly materialized account pool into the sidebar so the + // TUI's next poll renders the labels / health / cooldown it needs + // without waiting for the first fetch to complete. + await setSidebarMachineState( + buildSidebarMachineStateFromAccounts( + accountManager.getAccounts().map((entry) => ({ + index: entry.index, + label: entry.label, + enabled: entry.enabled, + current: false, + coolingDownUntil: entry.coolingDownUntil, + cachedQuota: entry.cachedQuota, + // Stamp the sidebar snapshot so the projection can detect a + // stale cache that landed on the wrong account (the manager's + // `cachedQuotaAccountId` is keyed to whatever account actually + // produced the snapshot — the live refresh-token hash is the + // expected identity at this slot). + cachedQuotaAccountId: entry.cachedQuotaAccountId, + currentQuotaAccountId: refreshTokenIdentity(entry.parts.refreshToken), + })), + ), + ) + } + + async function runLoader( getAuth: GetAuth, provider: Provider, - ): Promise> => { + ): Promise> { onGetAuth?.(getAuth) let auth = await getAuth() @@ -164,45 +256,7 @@ export function createAuthLoader({ } const accountManager = await deps.loadAccountManager(auth) - if (accountManager.getAccountCount() > 0) { - accountManager.requestSaveToDisk() - } - - let refreshQueue: ProactiveRefreshQueue | null = null - if ( - config.proactive_token_refresh && - accountManager.getAccountCount() > 0 - ) { - refreshQueue = deps.createRefreshQueue(client, providerId, { - enabled: config.proactive_token_refresh, - bufferSeconds: config.proactive_refresh_buffer_seconds, - checkIntervalSeconds: config.proactive_refresh_check_interval_seconds, - }) - refreshQueue.setAccountManager(accountManager) - } - - await lifecycle.replaceAccountRuntime(accountManager, refreshQueue) - refreshQueue?.start() - - const previousRuntime = fetchRuntime - fetchRuntime = createFetch({ accountManager, getAuth }) - await previousRuntime?.dispose() - - // Push the freshly materialized account pool into the sidebar so the - // TUI's next poll renders the labels / health / cooldown it needs - // without waiting for the first fetch to complete. - await setSidebarMachineState( - buildSidebarMachineStateFromAccounts( - accountManager.getAccounts().map((entry) => ({ - index: entry.index, - label: entry.label, - enabled: entry.enabled, - current: false, - coolingDownUntil: entry.coolingDownUntil, - cachedQuota: entry.cachedQuota, - })), - ), - ) + await installRuntime(accountManager, getAuth) if (deps.isDebugEnabled()) { const logPath = deps.getLogFilePath() @@ -223,7 +277,34 @@ export function createAuthLoader({ return { apiKey: '', - fetch: fetchRuntime.fetch, + // `fetchRuntime` is guaranteed non-null by `installRuntime`'s + // `createFetch({ ... })` call above — the assignment replaced + // any prior value. + fetch: fetchRuntime!.fetch, } } + + reloadRuntime = async (getAuth: GetAuth): Promise => { + const auth = await getAuth() + if (!isOAuthAuth(auth)) return + const nextManager = await deps.loadAccountManager(auth) + await installRuntime(nextManager, getAuth) + } + + // Return a callable object: `plugin.auth.loader` is invoked with + // `(getAuth, provider)` (host contract), and `authLoader.reload(...)` + // rebuilds the runtime after out-of-band storage mutations (OAuth add). + // The callable closes over `runLoader` directly so the `this`-context + // binding stays unbound. + async function authLoaderCallable( + getAuth: GetAuth, + provider: Provider, + ): Promise> { + return runLoader(getAuth, provider) + } + async function reload(getAuth: GetAuth): Promise { + return reloadRuntime(getAuth) + } + const authLoader = Object.assign(authLoaderCallable, { reload }) + return authLoader as LoadAndInstallRuntime & AuthLoaderHandle } diff --git a/packages/opencode/src/plugin/fetch-interceptor.ts b/packages/opencode/src/plugin/fetch-interceptor.ts index e5f2dcb..fe42d1b 100644 --- a/packages/opencode/src/plugin/fetch-interceptor.ts +++ b/packages/opencode/src/plugin/fetch-interceptor.ts @@ -302,12 +302,30 @@ export function createFetchInterceptor( singleAccount = accountsForCheck[accountIndex] if (!singleAccount) return + // Capture the refresh token BEFORE the await so a concurrent + // add/remove that shifts the index can't attribute this result + // to the wrong account when the refresh resolves. + const expectedRefreshToken = singleAccount.refreshToken + const result = await quotaManager.refreshAccount(singleAccount, { index: accountIndex, }) if (result.status === 'ok' && result.quota?.groups) { - accountManager.updateQuotaCache(accountIndex, result.quota.groups) + // Re-resolve the live index for the captured refresh token. If + // the account was removed mid-refresh, drop the result rather + // than writing onto whichever account shifted into the slot. + const currentIndex = accountManager + .getAccounts() + .findIndex( + (entry) => entry.parts.refreshToken === expectedRefreshToken, + ) + if (currentIndex === -1) return + accountManager.updateQuotaCache( + currentIndex, + result.quota.groups, + expectedRefreshToken, + ) accountManager.requestSaveToDisk() } } catch (err) { diff --git a/packages/opencode/src/plugin/index.ts b/packages/opencode/src/plugin/index.ts index 0ef96b8..55bd2e5 100644 --- a/packages/opencode/src/plugin/index.ts +++ b/packages/opencode/src/plugin/index.ts @@ -1,5 +1,5 @@ +import { createHash } from 'node:crypto' import { join } from 'node:path' -import { authorizeAntigravity, exchangeAntigravity } from '../antigravity/oauth' import { ANTIGRAVITY_PROVIDER_ID } from '../constants' import { createAutoUpdateCheckerHook } from '../hooks/auto-update-checker' import { drainNotifications, pushNotification } from '../rpc/notifications' @@ -65,6 +65,19 @@ import { initAntigravityVersion } from './version' const logger = createLogger('plugin') +/** + * Opaque identity derived from a refresh token. Mirrors + * `command-data.ts`'s copy (the two callers stay independent because + * `command-data.ts` ships into the TUI's compiled tree and can't reach + * across the core barrel boundary). Used by the live quota→sidebar + * writers to stamp each quota snapshot with the account it belongs to + * so the sidebar projection can detect a stale cache after an index + * shift. + */ +function quotaAccountIdentity(refreshToken: string): string { + return createHash('sha256').update(refreshToken).digest('hex').slice(0, 16) +} + /** * High-level options for the plugin factory. Production callers omit it * entirely; the e2e workspace injects overrides so the same factory can @@ -151,6 +164,11 @@ export const createAntigravityPlugin = enabled: entry.enabled, coolingDownUntil: entry.coolingDownUntil, cachedQuota: entry.cachedQuota, + // Carry the identity stamp so the sidebar projection can + // detect a stale snapshot that landed on the wrong account + // after an index shift (see `redactAccountForSidebar`). + cachedQuotaAccountId: entry.cachedQuotaAccountId, + currentQuotaAccountId: quotaAccountIdentity(entry.parts.refreshToken), })) }, }) @@ -291,23 +309,6 @@ export const createAntigravityPlugin = confirmOpenVerificationUrl: promptOpenVerificationUrl, }, }) - const accountOAuth = createAccountCommandOAuthService({ - authorize: () => authorizeAntigravity(), - exchange: exchangeAntigravity, - persist: (result) => accountAccess.persistAccountPool([result], false), - listAccounts: async () => - projectCommandAccountRows(await accountAccess.loadAccounts()), - }) - lifecycle.register({ dispose: () => accountOAuth.dispose() }) - const oauthMethods = createOAuthMethods({ - client, - providerId, - config, - lifecycle, - accountAccess, - quotaManager, - getAuth: async () => (cachedGetAuth ? cachedGetAuth() : undefined), - }) const authLoader = createAuthLoader({ client, providerId, @@ -331,6 +332,40 @@ export const createAntigravityPlugin = fetchImpl: dependencies.fetchImpl, }), }) + const accountOAuth = createAccountCommandOAuthService({ + // Wire the OAuth primitives through the dependency seam (rather + // than the concrete imports above) so injected overrides from + // `dependencies.oauth.*` reach this path — the same composition + // pattern used by every other sub-factory. Without this seam the + // e2e harness / custom-host deployments would still hit the real + // Google OAuth endpoints when adding an account. + authorize: dependencies.oauth.authorize, + exchange: dependencies.oauth.exchange, + persist: (result) => accountAccess.persistAccountPool([result], false), + listAccounts: async () => + projectCommandAccountRows(await accountAccess.loadAccounts()), + // After the new account lands on disk, reload the live + // AccountManager + fetch interceptor so routing sees it + // immediately — without waiting for an auth reload. + onAfterPersist: () => + authLoader.reload(async () => { + const auth = cachedGetAuth ? await cachedGetAuth() : undefined + if (auth) return auth + throw new Error( + 'No live auth cached for OAuth finish reload — the host has not yet loaded any account.', + ) + }), + }) + lifecycle.register({ dispose: () => accountOAuth.dispose() }) + const oauthMethods = createOAuthMethods({ + client, + providerId, + config, + lifecycle, + accountAccess, + quotaManager, + getAuth: async () => (cachedGetAuth ? cachedGetAuth() : undefined), + }) const rpcServer = await startRpcServer({ dir: getRpcDir(directory), @@ -347,6 +382,13 @@ export const createAntigravityPlugin = enabled: entry.enabled, coolingDownUntil: entry.coolingDownUntil, cachedQuota: entry.cachedQuota, + // Carry the identity stamp so the sidebar projection can + // detect a stale snapshot that landed on the wrong account + // after an index shift (see `redactAccountForSidebar`). + cachedQuotaAccountId: entry.cachedQuotaAccountId, + currentQuotaAccountId: quotaAccountIdentity( + entry.parts.refreshToken, + ), })) }) const result = await applyCommand(request, { diff --git a/packages/opencode/src/plugin/killswitch.test.ts b/packages/opencode/src/plugin/killswitch.test.ts index 2475b9c..16b4bb4 100644 --- a/packages/opencode/src/plugin/killswitch.test.ts +++ b/packages/opencode/src/plugin/killswitch.test.ts @@ -126,14 +126,20 @@ describe('evaluateKillswitchForAccount', () => { }) it('uses freshest quota group across the family when no model is provided', () => { - // Without a model, the family-max behavior is preserved — - // gemini pool is at 2% so it's below the 5% threshold. + // Without a model, the family-max behavior is preserved — pick the + // highest remaining pool across gemini + non-gemini. The 90% + // non-gemini pool must outweigh the 2% gemini pool so the family + // is allowed. The model-scoped tests below pin the regression: + // if a future change accidentally collapses both branches to + // family-max, those assertions fail because they expect the gemini + // pool alone to drive the deny decision. const decision = evaluateKillswitchForAccount( { index: 0, refreshToken: 'rt', cachedQuota: { gemini: { remainingFraction: 0.02, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.9, modelCount: 1 }, }, cachedQuotaUpdatedAt: NOW, }, @@ -141,19 +147,22 @@ describe('evaluateKillswitchForAccount', () => { enabledWith(5), { now: NOW }, ) - expect(decision.allowed).toBe(false) - expect(decision.reason).toBe('below-threshold') - expect(decision.remainingPercent).toBe(2) + expect(decision.allowed).toBe(true) + expect(decision.reason).toBe('ok') + expect(decision.remainingPercent).toBe(90) }) it('scopes evaluation to gemini alone when the model is a gemini model', () => { // gemini pool is at 2% — must be DENIED because below 5% threshold. + // The non-gemini pool sits at 90% so a regression that ignored + // `model` (and fell back to family-max) would let this through. const decision = evaluateKillswitchForAccount( { index: 0, refreshToken: 'rt', cachedQuota: { gemini: { remainingFraction: 0.02, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.9, modelCount: 1 }, }, cachedQuotaUpdatedAt: NOW, }, @@ -174,6 +183,7 @@ describe('evaluateKillswitchForAccount', () => { refreshToken: 'rt', cachedQuota: { gemini: { remainingFraction: 0.02, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.9, modelCount: 1 }, }, cachedQuotaUpdatedAt: NOW, }, diff --git a/packages/opencode/src/plugin/killswitch.ts b/packages/opencode/src/plugin/killswitch.ts index 31fda95..0d805ee 100644 --- a/packages/opencode/src/plugin/killswitch.ts +++ b/packages/opencode/src/plugin/killswitch.ts @@ -78,8 +78,14 @@ function quotaGroupForModel( const lower = model.toLowerCase() if (family === 'claude') return 'non-gemini' if (family === 'gemini') { + // Check Claude / GPT-OSS substrings BEFORE the `gemini` substring so + // a `gemini-claude-*` alias (Claude route exposed under a `gemini-` + // namespace) attributes to the non-gemini pool rather than the + // gemini pool. + if (lower.includes('claude') || lower.includes('gpt-oss')) { + return 'non-gemini' + } if (lower.includes('gemini') || lower.startsWith('tab_')) return 'gemini' - if (lower.includes('gpt-oss') || lower.includes('oss')) return 'non-gemini' } return null } diff --git a/packages/opencode/src/plugin/quota.ts b/packages/opencode/src/plugin/quota.ts index 5fbe0c9..4217f18 100644 --- a/packages/opencode/src/plugin/quota.ts +++ b/packages/opencode/src/plugin/quota.ts @@ -114,6 +114,8 @@ export function createOpenCodeQuotaManager( enabled?: boolean coolingDownUntil?: number cachedQuota?: AccountMetadataV3['cachedQuota'] + cachedQuotaAccountId?: string + currentQuotaAccountId?: string }> | null /** * Optional transport adapter used for both `fetchAvailableModels` @@ -285,6 +287,8 @@ export async function pushSidebarQuotaSnapshot( enabled?: boolean coolingDownUntil?: number cachedQuota?: AccountMetadataV3['cachedQuota'] + cachedQuotaAccountId?: string + currentQuotaAccountId?: string }> | null, backoffUntil: number = 0, ): Promise { @@ -300,6 +304,8 @@ export async function pushSidebarQuotaSnapshot( current: false, coolingDownUntil: entry.coolingDownUntil, cachedQuota: entry.cachedQuota, + cachedQuotaAccountId: entry.cachedQuotaAccountId, + currentQuotaAccountId: entry.currentQuotaAccountId, })), { checkedAt: Date.now(), diff --git a/packages/opencode/src/sidebar-state.test.ts b/packages/opencode/src/sidebar-state.test.ts index 03eb562..a9caa5e 100644 --- a/packages/opencode/src/sidebar-state.test.ts +++ b/packages/opencode/src/sidebar-state.test.ts @@ -511,4 +511,45 @@ describe('redaction', () => { }) expect(redacted.quota['non-gemini']?.remainingPercent).toBe(42) }) + + it('drops cachedQuota when the persisted stamp does not match the current account identity', () => { + // Persisted stamp from a different refresh token than the live one. + // The projection must drop the quota rather than render the wrong + // account's percentages after an index shift or token replacement. + const redacted = redactAccountForSidebar({ + index: 0, + cachedQuota: { + gemini: { remainingFraction: 0.42 }, + 'non-gemini': { remainingFraction: 0.8 }, + }, + cachedQuotaAccountId: 'deadbeefcafebabe', + currentQuotaAccountId: '0123456789abcdef', + }) + expect(redacted.quota).toEqual({}) + }) + + it('preserves cachedQuota when the stamp matches the current account identity', () => { + const stamp = 'a'.repeat(16) + const redacted = redactAccountForSidebar({ + index: 0, + cachedQuota: { + 'non-gemini': { remainingFraction: 0.5 }, + }, + cachedQuotaAccountId: stamp, + currentQuotaAccountId: stamp, + }) + expect(redacted.quota['non-gemini']?.remainingPercent).toBe(50) + }) + + it('preserves cachedQuota when only one of the stamps is provided (fail open for legacy)', () => { + // Legacy snapshots omit the stamp; pre-stamp live views also omit + // the current identity. The projection must not silently drop the + // quota in either half-missing case. + const redacted = redactAccountForSidebar({ + index: 0, + cachedQuota: { gemini: { remainingFraction: 0.3 } }, + cachedQuotaAccountId: 'a'.repeat(16), + }) + expect(redacted.quota.gemini?.remainingPercent).toBe(30) + }) }) diff --git a/packages/opencode/src/sidebar-state.ts b/packages/opencode/src/sidebar-state.ts index 83f5a29..108aeb5 100644 --- a/packages/opencode/src/sidebar-state.ts +++ b/packages/opencode/src/sidebar-state.ts @@ -435,6 +435,20 @@ export interface SidebarAccountRedactionInput { gemini?: { remainingFraction?: number; resetTime?: string } 'non-gemini'?: { remainingFraction?: number; resetTime?: string } } + /** + * Opaque identity stamp that was attached to the persisted quota snapshot. + * Used together with `currentQuotaAccountId` to detect a stale cache that + * landed on the wrong account after an index shift or token replacement. + * PII-safe — it is a 16-char hash, not the refresh token itself. + */ + cachedQuotaAccountId?: string + /** + * Opaque identity stamp for the account that is currently at this index. + * The redactor drops `cachedQuota` when `cachedQuotaAccountId` is set + * AND does not match this value, mirroring `toCommandAccountRow` in the + * command-data service. Omitted in the persisted sidebar state. + */ + currentQuotaAccountId?: string } /** @@ -462,7 +476,15 @@ export function redactAccountForSidebar( ) const quota: SidebarAccountState['quota'] = {} - const cached = source.cachedQuota + // Stamp mismatch: the persisted quota snapshot was captured for a + // different account than the one currently at this index. Drop the + // stale cache rather than rendering the wrong account's quota + // percentages. Mirrors `toCommandAccountRow` in command-data. + const staleCachedQuota = + typeof source.cachedQuotaAccountId === 'string' && + typeof source.currentQuotaAccountId === 'string' && + source.cachedQuotaAccountId !== source.currentQuotaAccountId + const cached = staleCachedQuota ? undefined : source.cachedQuota if (cached) { for (const key of ['gemini', 'non-gemini'] as const) { const entry = cached[key] diff --git a/packages/opencode/src/tui-compiled/sidebar-state.ts b/packages/opencode/src/tui-compiled/sidebar-state.ts index 83f5a29..108aeb5 100644 --- a/packages/opencode/src/tui-compiled/sidebar-state.ts +++ b/packages/opencode/src/tui-compiled/sidebar-state.ts @@ -435,6 +435,20 @@ export interface SidebarAccountRedactionInput { gemini?: { remainingFraction?: number; resetTime?: string } 'non-gemini'?: { remainingFraction?: number; resetTime?: string } } + /** + * Opaque identity stamp that was attached to the persisted quota snapshot. + * Used together with `currentQuotaAccountId` to detect a stale cache that + * landed on the wrong account after an index shift or token replacement. + * PII-safe — it is a 16-char hash, not the refresh token itself. + */ + cachedQuotaAccountId?: string + /** + * Opaque identity stamp for the account that is currently at this index. + * The redactor drops `cachedQuota` when `cachedQuotaAccountId` is set + * AND does not match this value, mirroring `toCommandAccountRow` in the + * command-data service. Omitted in the persisted sidebar state. + */ + currentQuotaAccountId?: string } /** @@ -462,7 +476,15 @@ export function redactAccountForSidebar( ) const quota: SidebarAccountState['quota'] = {} - const cached = source.cachedQuota + // Stamp mismatch: the persisted quota snapshot was captured for a + // different account than the one currently at this index. Drop the + // stale cache rather than rendering the wrong account's quota + // percentages. Mirrors `toCommandAccountRow` in command-data. + const staleCachedQuota = + typeof source.cachedQuotaAccountId === 'string' && + typeof source.currentQuotaAccountId === 'string' && + source.cachedQuotaAccountId !== source.currentQuotaAccountId + const cached = staleCachedQuota ? undefined : source.cachedQuota if (cached) { for (const key of ['gemini', 'non-gemini'] as const) { const entry = cached[key] From 52b78021ec2a4cdfc3b3a5c67180c86b9c390ae6 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 19:00:47 +0200 Subject: [PATCH 07/25] fix(commands): add account to live runtime + oauth concurrency/seam - P1#2: accountOAuth now accepts an onAfterPersist hook the plugin entry wires to authLoader.reload, so a successful OAuth add refreshes the live AccountManager + fetch interceptor immediately instead of waiting for the next auth reload. Reuses the auth-loader install path (load -> replaceAccountRuntime -> fetch rebuild -> sidebar publish) rather than inventing a second reload path. authLoader exposes .reload alongside the callable loader shape the host contract requires (auth.auth.loader). - P2#4: takePending now consumes the entry atomically with the take (deleted inside takePending BEFORE the async exchange starts) so two concurrent add-oauth-finish calls for the same session cannot both exchange the same auth code. The previous peek-then- finally-del pattern was the long-deferred 'Should' from a11. - P2#7: accountOAuth.authorize / .exchange now go through dependencies.oauth.* (the same seam every other sub-factory uses) rather than the concrete authorizeAntigravity / exchangeAntigravity imports, so injected overrides (e2e harness, custom hosts) reach the OAuth add path. - P2#10: account-command-oauth.finish separates the four error stages - callback parse, exchange, persist, listAccounts - so a persistence failure no longer reports as 'OAuth exchange failed due to a network error'. The post-persist listAccounts failure still reports success with an empty accounts payload (the on-disk write already landed). - P3: commands.ts now imports the shared AccountCommandOAuthService type from account-command-oauth instead of declaring its own. Adds onAfterPersist invocation tests, the concurrent finish() race test, the persistence-vs-exchange error stage tests, and asserts exchange is NOT called when no pending session exists. --- .../src/plugin/account-command-oauth.test.ts | 194 +++++++++++++++++- .../src/plugin/account-command-oauth.ts | 98 +++++++-- packages/opencode/src/plugin/commands.test.ts | 12 +- packages/opencode/src/plugin/commands.ts | 16 +- 4 files changed, 284 insertions(+), 36 deletions(-) diff --git a/packages/opencode/src/plugin/account-command-oauth.test.ts b/packages/opencode/src/plugin/account-command-oauth.test.ts index 7df18ce..95b852d 100644 --- a/packages/opencode/src/plugin/account-command-oauth.test.ts +++ b/packages/opencode/src/plugin/account-command-oauth.test.ts @@ -87,7 +87,136 @@ describe('createAccountCommandOAuthService', () => { }) }) - it('returns a friendly error when the session has no pending OAuth flow', async () => { + it('invokes onAfterPersist with the exchanged success result so the live runtime can reload', async () => { + const persisted = mock(async () => {}) + const onAfterPersist = mock(async () => {}) + const service = createAccountCommandOAuthService({ + authorize: async () => ({ + url: 'https://accounts.google.test/authorize?state=signed-state&redirect_uri=http%3A%2F%2Flocalhost%3A51121%2Foauth-callback', + verifier: 'pkce-verifier', + projectId: '', + }), + exchange: mock(async () => ({ + type: 'success' as const, + refresh: 'refresh-token|project', + access: 'access-token', + expires: 123, + email: 'private@example.test', + label: 'OAuth display name', + projectId: 'project', + })), + persist: persisted, + listAccounts: mock(async () => initialRows), + onAfterPersist, + }) + + await service.start('session-1') + await service.finish('session-1', 'callback-code', 'Work account') + + expect(persisted).toHaveBeenCalledTimes(1) + expect(onAfterPersist).toHaveBeenCalledTimes(1) + // The reload hook gets the same success result persist saw, so the + // runtime reload (auth-loader path) can re-read the freshly-persisted + // account pool without re-running the OAuth exchange. + expect(onAfterPersist).toHaveBeenCalledWith( + expect.objectContaining({ refresh: 'refresh-token|project' }), + ) + }) + + it('still reports success when onAfterPersist throws — the on-disk write already landed', async () => { + const persisted = mock(async () => {}) + const service = createAccountCommandOAuthService({ + authorize: async () => ({ + url: 'https://accounts.google.test/authorize?state=signed-state&redirect_uri=http%3A%2F%2Flocalhost%3A51121%2Foauth-callback', + verifier: 'pkce-verifier', + projectId: '', + }), + exchange: mock(async () => ({ + type: 'success' as const, + refresh: 'refresh-token|project', + access: 'access-token', + expires: 123, + email: 'private@example.test', + label: 'OAuth display name', + projectId: 'project', + })), + persist: persisted, + listAccounts: mock(async () => initialRows), + onAfterPersist: mock(async () => { + throw new Error('runtime reload blew up') + }), + }) + + await service.start('session-1') + const finished = await service.finish('session-1', 'callback-code') + + expect(persisted).toHaveBeenCalledTimes(1) + expect(finished.text).toBe('OAuth account added.') + }) + + it('consumes the pending entry inside takePending so two concurrent finish() calls do not both exchange', async () => { + // Build a service where the exchange hangs until we resolve it + // manually. This guarantees the first finish() is still mid-flight + // when the second one runs — the only way a regression that + // peeks-then-finally-dels would surface. + let releaseExchange: (() => void) | undefined + const exchangeGate = new Promise((resolve) => { + releaseExchange = resolve + }) + const exchange = mock(async () => { + await exchangeGate + return { + type: 'success' as const, + refresh: 'refresh-token|project', + access: 'access-token', + expires: 123, + email: 'private@example.test', + label: 'OAuth display name', + projectId: 'project', + } + }) + const persist = mock(async () => {}) + const service = createAccountCommandOAuthService({ + authorize: async () => ({ + url: 'https://accounts.google.test/authorize?state=signed-state&redirect_uri=http%3A%2F%2Flocalhost%3A51121%2Foauth-callback', + verifier: 'pkce-verifier', + projectId: '', + }), + exchange, + persist, + listAccounts: mock(async () => initialRows), + }) + + await service.start('session-1') + + // First finish() hangs inside exchange. + const firstFinish = service.finish('session-1', 'callback-code-1') + + // Give the first finish() a microtask tick to enter exchange. + await Promise.resolve() + + // Second finish() before the first resolves — the pending entry + // should already be consumed, so the second call must observe + // "no pending entry" and short-circuit. + const secondFinish = await service.finish('session-1', 'callback-code-2') + expect(secondFinish.text).toBe('OAuth session expired. Please start again.') + + // Release the first call. + releaseExchange?.() + const first = await firstFinish + expect(first.text).toBe('OAuth account added.') + expect(exchange).toHaveBeenCalledTimes(1) + expect(persist).toHaveBeenCalledTimes(1) + }) + + it('reports a persistence-stage error distinctly from the exchange-stage error', async () => { + // The previous implementation collapsed every caught error into + // "OAuth exchange failed due to a network error". When persist + // throws AFTER a successful exchange, the operator needs to know + // the account was NOT stored — not that the OAuth itself failed. + const persist = mock(async () => { + throw new Error('disk full') + }) const service = createAccountCommandOAuthService({ authorize: async () => ({ url: 'https://accounts.google.test/authorize?state=signed-state&redirect_uri=http%3A%2F%2Flocalhost%3A51121%2Foauth-callback', @@ -95,9 +224,65 @@ describe('createAccountCommandOAuthService', () => { projectId: '', }), exchange: mock(async () => ({ - type: 'failed' as const, - error: 'unused', + type: 'success' as const, + refresh: 'refresh-token|project', + access: 'access-token', + expires: 123, + email: 'private@example.test', + label: 'OAuth display name', + projectId: 'project', })), + persist, + listAccounts: mock(async () => initialRows), + }) + + await service.start('session-1') + const finished = await service.finish('session-1', 'callback-code') + + expect(persist).toHaveBeenCalledTimes(1) + expect(finished.text).toBe( + 'OAuth account could not be saved to disk. Please try again.', + ) + }) + + it('reports an exchange-stage error (network) without attempting persist', async () => { + // The exchange throws before persist runs. The dialog must surface + // the network-stage message and skip the persistence call entirely. + const persist = mock(async () => {}) + const service = createAccountCommandOAuthService({ + authorize: async () => ({ + url: 'https://accounts.google.test/authorize?state=signed-state&redirect_uri=http%3A%2F%2Flocalhost%3A51121%2Foauth-callback', + verifier: 'pkce-verifier', + projectId: '', + }), + exchange: mock(async () => { + throw new Error('network down') + }), + persist, + listAccounts: mock(async () => initialRows), + }) + + await service.start('session-1') + const finished = await service.finish('session-1', 'callback-code') + + expect(finished.text).toBe( + 'OAuth exchange failed due to a network error. Please try again.', + ) + expect(persist).not.toHaveBeenCalled() + }) + + it('returns a friendly error when the session has no pending OAuth flow', async () => { + const exchange = mock(async () => ({ + type: 'failed' as const, + error: 'unused', + })) + const service = createAccountCommandOAuthService({ + authorize: async () => ({ + url: 'https://accounts.google.test/authorize?state=signed-state&redirect_uri=http%3A%2F%2Flocalhost%3A51121%2Foauth-callback', + verifier: 'pkce-verifier', + projectId: '', + }), + exchange, persist: mock(async () => {}), listAccounts: mock(async () => initialRows), }) @@ -108,5 +293,8 @@ describe('createAccountCommandOAuthService', () => { text: 'OAuth session expired. Please start again.', accounts: initialRows, }) + // A finish() for a session that never started must not even touch + // the exchange network path — there is no auth code to redeem. + expect(exchange).not.toHaveBeenCalled() }) }) diff --git a/packages/opencode/src/plugin/account-command-oauth.ts b/packages/opencode/src/plugin/account-command-oauth.ts index 5a9feef..095e7c4 100644 --- a/packages/opencode/src/plugin/account-command-oauth.ts +++ b/packages/opencode/src/plugin/account-command-oauth.ts @@ -22,6 +22,14 @@ export interface AccountCommandOAuthServiceOptions { ) => Promise persist: (result: OAuthSuccess) => Promise listAccounts: () => Promise + /** + * Optional hook invoked AFTER persist completes successfully. The plugin + * uses it to refresh the live AccountManager so the new account is + * visible to routing immediately, without waiting for an auth reload. + * Failures are swallowed — the OAuth flow still reports success and the + * next periodic AccountManager reload picks up the new account. + */ + onAfterPersist?: (result: OAuthSuccess) => Promise | void now?: () => number } @@ -66,6 +74,12 @@ export function createAccountCommandOAuthService( pendingBySession.delete(sessionId) return undefined } + // Consume the entry atomically with the take. A second concurrent + // `add-oauth-finish` call must observe "no pending entry" even + // before the first call's exchange completes — the previous + // peek-then-finally pattern allowed two finish() calls to both + // exchange and persist the same auth code. + pendingBySession.delete(sessionId) return entry } @@ -105,33 +119,83 @@ export function createAccountCommandOAuthService( } } + // Stage 1: parse the callback URL/code. Failures here mean the + // user pasted a malformed callback — exchange has not started. + let callback: ReturnType try { - const callback = parseOAuthCallbackInput(callbackInput, pending.state) - if ('error' in callback) { - return { - text: `OAuth authentication failed: ${callback.error}`, - accounts: await options.listAccounts(), - } - } - const result = await options.exchange(callback.code, callback.state) - if (result.type === 'failed') { - return { - text: 'OAuth authentication failed. Please check the code and try again.', - accounts: await options.listAccounts(), - } + callback = parseOAuthCallbackInput(callbackInput, pending.state) + } catch { + return { + text: 'OAuth authentication failed: could not parse the callback. Please try again.', + accounts: await options.listAccounts(), } - await options.persist({ ...result, label: label || result.label }) + } + if ('error' in callback) { return { - text: 'OAuth account added.', + text: `OAuth authentication failed: ${callback.error}`, accounts: await options.listAccounts(), } + } + + // Stage 2: exchange the auth code with Google. Failures here mean + // the network or token endpoint rejected the code — nothing has + // been persisted yet. + let result: Awaited> + try { + result = await options.exchange(callback.code, callback.state) } catch { return { text: 'OAuth exchange failed due to a network error. Please try again.', accounts: await options.listAccounts(), } - } finally { - pendingBySession.delete(sessionId) + } + if (result.type === 'failed') { + return { + text: 'OAuth authentication failed. Please check the code and try again.', + accounts: await options.listAccounts(), + } + } + + // Stage 3: persist the new account to disk. A failure here means + // the account is NOT stored — surface the stage so the operator + // can retry without thinking the previous error already landed it. + const persisted: OAuthSuccess = { + ...result, + label: label || result.label, + } + try { + await options.persist(persisted) + } catch { + return { + text: 'OAuth account could not be saved to disk. Please try again.', + accounts: await options.listAccounts(), + } + } + + // Refresh the live AccountManager so routing sees the new account + // immediately instead of waiting for the next auth reload. Errors + // here are non-fatal — the on-disk write already landed. + try { + await options.onAfterPersist?.(persisted) + } catch { + // Best-effort refresh; swallow failures. + } + + // Stage 4: re-list the account pool. If the post-persist read + // fails for any reason (lock contention, I/O) we still report the + // successful add — the dialog will refresh on its next open. + let accounts: CommandAccountRow[] + try { + accounts = await options.listAccounts() + } catch { + return { + text: 'OAuth account added.', + accounts: [], + } + } + return { + text: 'OAuth account added.', + accounts, } }, diff --git a/packages/opencode/src/plugin/commands.test.ts b/packages/opencode/src/plugin/commands.test.ts index 95bb104..92c39cb 100644 --- a/packages/opencode/src/plugin/commands.test.ts +++ b/packages/opencode/src/plugin/commands.test.ts @@ -769,7 +769,16 @@ describe('applyCommand', () => { quota: [], }, ]) - const refreshQuota = mock(async () => []) + // refreshQuota returns a never-resolving promise. If buildDialogPayload + // ever accidentally `await`s the refresh path (instead of the fire- + // and-forget `void` it does today) this test will hang and the + // runner will report it as a timeout — exactly the regression we + // want to surface. + let refreshStarted = false + const refreshQuota = mock(() => { + refreshStarted = true + return new Promise(() => {}) + }) const payload = await buildDialogPayload('antigravity-quota', '', { client: {} as never, @@ -779,6 +788,7 @@ describe('applyCommand', () => { }) expect(payload.knobs.accounts).toHaveLength(1) + expect(refreshStarted).toBe(true) expect(refreshQuota).toHaveBeenCalledTimes(1) }) diff --git a/packages/opencode/src/plugin/commands.ts b/packages/opencode/src/plugin/commands.ts index 1498506..4e84844 100644 --- a/packages/opencode/src/plugin/commands.ts +++ b/packages/opencode/src/plugin/commands.ts @@ -27,6 +27,7 @@ import { buildSidebarMachineStateFromAccounts, setSidebarMachineState, } from '../sidebar-state' +import type { AccountCommandOAuthService } from './account-command-oauth' import type { CommandAccountRow, CommandDataService } from './command-data' import { executeGeminiDumpCommand, @@ -129,21 +130,6 @@ interface CommandContext { accountOAuth?: AccountCommandOAuthService } -export interface AccountCommandOAuthService { - start(sessionId: string): Promise<{ - url: string - accounts: CommandAccountRow[] - }> - finish( - sessionId: string, - code: string, - label?: string, - ): Promise<{ - text: string - accounts: CommandAccountRow[] - }> -} - /** * Build the dialog payload the TUI renders for `command`. * From 90337d91ac83df590cca3dc2eb06c3ce5851297b Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 19:01:23 +0200 Subject: [PATCH 08/25] fix(quota): preserve reset times + current-index on sidebar/remove - P2#8: writeSidebar now preserves each pool's q.resetAt as an ISO resetTime string so the sidebar can render a reset countdown for the gemini + non-gemini pools. The previous toFraction only carried remainingPercent, so every dialog re-render dropped the freshest reset deadline. - P2#9: mutateLiveAndStorage.remove now captures the live current account's refresh token per family BEFORE the storage mutation and resolves the post-removal storage index by token, so a non-current account removal persists the same current the live AccountManager.removeAccountByIndex preserves. The previous hardcoded activeIndex: 0 re-elected whichever account shifted into slot 0 on every restart. - Pins both fixes with new tests: the refreshed-account stamp assertion (command-data.test.ts) and the non-current-removal remap test that asserts the persisted activeIndex follows the shift instead of resetting to 0. - Regenerates src/tui-compiled/plugin/command-data.ts from source. --- .../opencode/src/plugin/command-data.test.ts | 68 ++++++++++++++++- packages/opencode/src/plugin/command-data.ts | 75 +++++++++++++++++-- .../src/tui-compiled/plugin/command-data.ts | 75 +++++++++++++++++-- 3 files changed, 203 insertions(+), 15 deletions(-) diff --git a/packages/opencode/src/plugin/command-data.test.ts b/packages/opencode/src/plugin/command-data.test.ts index 413b9df..74f43fb 100644 --- a/packages/opencode/src/plugin/command-data.test.ts +++ b/packages/opencode/src/plugin/command-data.test.ts @@ -666,6 +666,12 @@ describe('createCommandDataService', () => { ?.remainingFraction, ).toBe(0.7) expect(harness.storage.accounts[0]?.cachedQuotaUpdatedAt).toBeGreaterThan(1) + // The refreshed account's persisted snapshot must carry the identity + // stamp derived from its refresh token — pins the P1#1 fix that + // propagates `cachedQuotaAccountId` through `buildStorageSnapshot`. + expect(harness.storage.accounts[0]?.cachedQuotaAccountId).toBe( + quotaAccountIdentity('refresh-a'), + ) // The non-refreshed account's cached state must remain untouched. expect( harness.storage.accounts[1]?.cachedQuota?.gemini?.remainingFraction, @@ -883,7 +889,67 @@ describe('createCommandDataService', () => { 'refresh-a', 'refresh-c', ]) - // CLI menu semantics: active index resets to 0 after removal. + // Active index tracks the same refresh token in the post-remove + // array. Removing a non-current account keeps the current's slot + // unchanged; the previous implementation reset to 0, which would + // re-elect the (now first) account on every restart even though + // the user explicitly removed a different one. + expect(harness.storage.activeIndex).toBe(0) + expect(harness.storage.activeIndexByFamily).toEqual({ + claude: 0, + gemini: 0, + }) + }) + + it("removeAccount() persists the current account's remapped index, not a hardcoded 0", async () => { + // Current is refresh-c (index 2). Removing a NON-current account + // (refresh-b at index 1) must persist the same current index as + // the live manager — refresh-c shifted from index 2 to index 1, so + // a restart must re-elect refresh-c, not the hardcoded 0. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha' }), + makeAccountFixture({ refreshToken: 'refresh-b', label: 'Beta' }), + makeAccountFixture({ refreshToken: 'refresh-c', label: 'Gamma' }), + ], + activeIndex: 2, + }) + + await harness.service.removeAccount(1) + + expect(harness.storage.accounts.map((a) => a.refreshToken)).toEqual([ + 'refresh-a', + 'refresh-c', + ]) + // Live manager shifts refresh-c from index 2 to index 1 because + // the removed refresh-b sat at index 1. The persisted activeIndex + // must follow that shift. + expect(harness.storage.activeIndex).toBe(1) + expect(harness.storage.activeIndexByFamily).toEqual({ + claude: 1, + gemini: 1, + }) + }) + + it('removeAccount() keeps the current account at its unchanged index when removing an account after it', async () => { + // Current is refresh-a (index 0). Removing refresh-b at index 1 + // doesn't touch the current's slot; the persisted index must stay + // at 0 (the current account is unchanged). + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha' }), + makeAccountFixture({ refreshToken: 'refresh-b', label: 'Beta' }), + makeAccountFixture({ refreshToken: 'refresh-c', label: 'Gamma' }), + ], + activeIndex: 0, + }) + + await harness.service.removeAccount(1) + + expect(harness.storage.accounts.map((a) => a.refreshToken)).toEqual([ + 'refresh-a', + 'refresh-c', + ]) expect(harness.storage.activeIndex).toBe(0) expect(harness.storage.activeIndexByFamily).toEqual({ claude: 0, diff --git a/packages/opencode/src/plugin/command-data.ts b/packages/opencode/src/plugin/command-data.ts index 6d3a84f..210310e 100644 --- a/packages/opencode/src/plugin/command-data.ts +++ b/packages/opencode/src/plugin/command-data.ts @@ -419,10 +419,19 @@ export function createCommandDataService( const gemini = row.quota.find((q) => q.key === 'gemini') const nonGemini = row.quota.find((q) => q.key === 'non-gemini') const toFraction = ( - q: { remainingPercent: number | null } | undefined, + q: { remainingPercent: number | null; resetAt?: number } | undefined, ): { remainingFraction?: number; resetTime?: string } | undefined => { if (!q || q.remainingPercent == null) return undefined - return { remainingFraction: q.remainingPercent / 100 } + return { + remainingFraction: q.remainingPercent / 100, + // Preserve resetTime so the sidebar can render reset countdowns + // for each pool. Without this the TUI loses the freshest reset + // deadline every time the dialog re-renders. + resetTime: + typeof q.resetAt === 'number' && Number.isFinite(q.resetAt) + ? new Date(q.resetAt).toISOString() + : undefined, + } } return { index: row.index, @@ -433,6 +442,9 @@ export function createCommandDataService( gemini: toFraction(gemini), 'non-gemini': toFraction(nonGemini), }, + // The stamp check has already been done by `toCommandAccountRow` + // before the rows reach this writer; nothing further for the + // sidebar projection to validate here. } }) // Fire-and-forget — the sidebar writer is fenced by its own queue, @@ -586,9 +598,29 @@ export function createCommandDataService( ) } + // Capture the live view's current-account identity per family BEFORE + // the storage mutation. The remove action must persist the index that + // the same account will occupy AFTER the removal — unconditionally + // resetting to 0 made a non-current removal promote whichever account + // shifted into slot 0 to "active" on the next restart. + const liveCurrentTokens: { claude?: string; gemini?: string } = {} + if (action === 'remove') { + const liveAccounts = accountManagerView.getAccounts() + const liveCurrentIndex = accountManagerView.activeIndex() + const liveCurrentToken = + liveAccounts[liveCurrentIndex]?.refreshToken ?? undefined + liveCurrentTokens.claude = liveCurrentToken + liveCurrentTokens.gemini = liveCurrentToken + } + let foundInStorage = false let desiredEnabled: boolean | undefined let previousEnabled: boolean | undefined + let nextActiveIndex = 0 + let nextActiveIndexByFamily: { claude?: number; gemini?: number } = { + claude: 0, + gemini: 0, + } await storage.mutate((current) => { const tokenIdx = current.accounts.findIndex( (entry) => entry.refreshToken === refreshToken, @@ -624,13 +656,42 @@ export function createCommandDataService( } } + // remove: build the post-removal account list, then resolve the + // current-account's NEW index in that list. If the removed + // account was the live current, the current token falls out of + // the list entirely and we fall back to index 0 — matching the + // live AccountManager.removeAccount() behavior (which leaves the + // current index pointing at the same numeric slot, now occupied + // by whichever account shifted in). + const nextAccounts = current.accounts.filter( + (account) => account.refreshToken !== refreshToken, + ) + const resolveNextIndex = ( + liveToken: string | undefined, + legacyIndex: number, + ): number => { + if (!liveToken) + return Math.max(0, Math.min(legacyIndex, nextAccounts.length - 1)) + const found = nextAccounts.findIndex( + (account) => account.refreshToken === liveToken, + ) + if (found === -1) return 0 + return found + } + const legacyClaude = + current.activeIndexByFamily?.claude ?? current.activeIndex + const legacyGemini = + current.activeIndexByFamily?.gemini ?? current.activeIndex + nextActiveIndex = resolveNextIndex(liveCurrentTokens.claude, legacyClaude) + nextActiveIndexByFamily = { + claude: resolveNextIndex(liveCurrentTokens.claude, legacyClaude), + gemini: resolveNextIndex(liveCurrentTokens.gemini, legacyGemini), + } return { ...current, - accounts: current.accounts.filter( - (account) => account.refreshToken !== refreshToken, - ), - activeIndex: 0, - activeIndexByFamily: { claude: 0, gemini: 0 }, + accounts: nextAccounts, + activeIndex: nextActiveIndex, + activeIndexByFamily: nextActiveIndexByFamily, } }) diff --git a/packages/opencode/src/tui-compiled/plugin/command-data.ts b/packages/opencode/src/tui-compiled/plugin/command-data.ts index 6d3a84f..210310e 100644 --- a/packages/opencode/src/tui-compiled/plugin/command-data.ts +++ b/packages/opencode/src/tui-compiled/plugin/command-data.ts @@ -419,10 +419,19 @@ export function createCommandDataService( const gemini = row.quota.find((q) => q.key === 'gemini') const nonGemini = row.quota.find((q) => q.key === 'non-gemini') const toFraction = ( - q: { remainingPercent: number | null } | undefined, + q: { remainingPercent: number | null; resetAt?: number } | undefined, ): { remainingFraction?: number; resetTime?: string } | undefined => { if (!q || q.remainingPercent == null) return undefined - return { remainingFraction: q.remainingPercent / 100 } + return { + remainingFraction: q.remainingPercent / 100, + // Preserve resetTime so the sidebar can render reset countdowns + // for each pool. Without this the TUI loses the freshest reset + // deadline every time the dialog re-renders. + resetTime: + typeof q.resetAt === 'number' && Number.isFinite(q.resetAt) + ? new Date(q.resetAt).toISOString() + : undefined, + } } return { index: row.index, @@ -433,6 +442,9 @@ export function createCommandDataService( gemini: toFraction(gemini), 'non-gemini': toFraction(nonGemini), }, + // The stamp check has already been done by `toCommandAccountRow` + // before the rows reach this writer; nothing further for the + // sidebar projection to validate here. } }) // Fire-and-forget — the sidebar writer is fenced by its own queue, @@ -586,9 +598,29 @@ export function createCommandDataService( ) } + // Capture the live view's current-account identity per family BEFORE + // the storage mutation. The remove action must persist the index that + // the same account will occupy AFTER the removal — unconditionally + // resetting to 0 made a non-current removal promote whichever account + // shifted into slot 0 to "active" on the next restart. + const liveCurrentTokens: { claude?: string; gemini?: string } = {} + if (action === 'remove') { + const liveAccounts = accountManagerView.getAccounts() + const liveCurrentIndex = accountManagerView.activeIndex() + const liveCurrentToken = + liveAccounts[liveCurrentIndex]?.refreshToken ?? undefined + liveCurrentTokens.claude = liveCurrentToken + liveCurrentTokens.gemini = liveCurrentToken + } + let foundInStorage = false let desiredEnabled: boolean | undefined let previousEnabled: boolean | undefined + let nextActiveIndex = 0 + let nextActiveIndexByFamily: { claude?: number; gemini?: number } = { + claude: 0, + gemini: 0, + } await storage.mutate((current) => { const tokenIdx = current.accounts.findIndex( (entry) => entry.refreshToken === refreshToken, @@ -624,13 +656,42 @@ export function createCommandDataService( } } + // remove: build the post-removal account list, then resolve the + // current-account's NEW index in that list. If the removed + // account was the live current, the current token falls out of + // the list entirely and we fall back to index 0 — matching the + // live AccountManager.removeAccount() behavior (which leaves the + // current index pointing at the same numeric slot, now occupied + // by whichever account shifted in). + const nextAccounts = current.accounts.filter( + (account) => account.refreshToken !== refreshToken, + ) + const resolveNextIndex = ( + liveToken: string | undefined, + legacyIndex: number, + ): number => { + if (!liveToken) + return Math.max(0, Math.min(legacyIndex, nextAccounts.length - 1)) + const found = nextAccounts.findIndex( + (account) => account.refreshToken === liveToken, + ) + if (found === -1) return 0 + return found + } + const legacyClaude = + current.activeIndexByFamily?.claude ?? current.activeIndex + const legacyGemini = + current.activeIndexByFamily?.gemini ?? current.activeIndex + nextActiveIndex = resolveNextIndex(liveCurrentTokens.claude, legacyClaude) + nextActiveIndexByFamily = { + claude: resolveNextIndex(liveCurrentTokens.claude, legacyClaude), + gemini: resolveNextIndex(liveCurrentTokens.gemini, legacyGemini), + } return { ...current, - accounts: current.accounts.filter( - (account) => account.refreshToken !== refreshToken, - ), - activeIndex: 0, - activeIndexByFamily: { claude: 0, gemini: 0 }, + accounts: nextAccounts, + activeIndex: nextActiveIndex, + activeIndexByFamily: nextActiveIndexByFamily, } }) From 50340ded9ab2b861adc8cddc8eb6f24f9e200c2f Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 19:02:02 +0200 Subject: [PATCH 09/25] test: strengthen quota/killswitch/oauth coverage + small ui fixes - accounts.test.ts 'model takes precedence over family' now uses a cross-family pair (claude family + gemini model, gemini family + claude/gpt-oss model) so a regression that drops the model-vs- family check would no longer pass. - quota-status.test.ts gains a legacy/unknown-key fixture asserting that the 'claude', 'gemini-antigravity', and 'gemini-cli' keys left over from the 3/4-key model are dropped by the collapsed two-pool projection while the supported keys still render. - command-dialogs.tsx empty OAuth code-prompt submit now toasts a validation message and reopens the prompt instead of silently dropping the user back to the account list. - command-dialogs.test.tsx asserts the OAuth URL surfaces in the Copy-URL option's description so the user can paste the mocked authorize URL into a browser. - Regenerates src/tui-compiled/tui/command-dialogs.tsx from source. --- packages/opencode/src/plugin/accounts.test.ts | 14 +++++++++++--- .../opencode/src/plugin/ui/quota-status.test.ts | 14 ++++++++++++++ .../src/tui-compiled/tui/command-dialogs.tsx | 9 ++++++++- packages/opencode/src/tui/command-dialogs.test.tsx | 8 ++++++++ packages/opencode/src/tui/command-dialogs.tsx | 9 ++++++++- 5 files changed, 49 insertions(+), 5 deletions(-) diff --git a/packages/opencode/src/plugin/accounts.test.ts b/packages/opencode/src/plugin/accounts.test.ts index 775434d..8423130 100644 --- a/packages/opencode/src/plugin/accounts.test.ts +++ b/packages/opencode/src/plugin/accounts.test.ts @@ -2680,9 +2680,17 @@ describe('resolveQuotaGroup', () => { }) it('model takes precedence over family', () => { - // Even if family says claude, model determines the quota group - expect(resolveQuotaGroup('gemini', 'gemini-2.5-flash')).toBe('gemini') - expect(resolveQuotaGroup('gemini', 'gemini-3-pro')).toBe('gemini') + // Cross-family pair: a `claude` family carrying a `gemini-` model + // string must resolve to the gemini pool, and vice versa. The + // previous version only exercised gemini-family + gemini-model, + // which would pass even if a regression collapsed the model-vs- + // family check entirely. + expect(resolveQuotaGroup('claude', 'gemini-2.5-flash')).toBe('gemini') + expect(resolveQuotaGroup('claude', 'gemini-3-pro')).toBe('gemini') + expect(resolveQuotaGroup('gemini', 'claude-sonnet-4-6')).toBe('non-gemini') + expect(resolveQuotaGroup('gemini', 'gpt-oss-120b-medium')).toBe( + 'non-gemini', + ) }) }) diff --git a/packages/opencode/src/plugin/ui/quota-status.test.ts b/packages/opencode/src/plugin/ui/quota-status.test.ts index 69a8103..0df7c1f 100644 --- a/packages/opencode/src/plugin/ui/quota-status.test.ts +++ b/packages/opencode/src/plugin/ui/quota-status.test.ts @@ -263,6 +263,20 @@ describe('formatCachedQuotaWithStatus', () => { expect(formatCachedQuotaWithStatus({})).toBeUndefined() }) + it('ignores unknown / legacy quota keys while keeping the supported ones', () => { + // Older snapshots may carry `claude`, `gemini-antigravity`, or + // `gemini-cli` keys left over from the 3/4-key model. The collapsed + // two-pool projection must drop those (not crash, not surface them + // as extra groups) and render only the supported pools. + const result = formatCachedQuotaWithStatus({ + claude: { remainingFraction: 0.5 }, + 'gemini-antigravity': { remainingFraction: 0.5 }, + 'gemini-cli': { remainingFraction: 0.5 }, + 'non-gemini': { remainingFraction: 0.8 }, + } as never) + expect(result).toBe('Non-Gemini 80%') + }) + it('formats READY groups without status label', () => { const result = formatCachedQuotaWithStatus({ 'non-gemini': { remainingFraction: 0.8 }, diff --git a/packages/opencode/src/tui-compiled/tui/command-dialogs.tsx b/packages/opencode/src/tui-compiled/tui/command-dialogs.tsx index 1ca7418..0031ca3 100644 --- a/packages/opencode/src/tui-compiled/tui/command-dialogs.tsx +++ b/packages/opencode/src/tui-compiled/tui/command-dialogs.tsx @@ -284,7 +284,14 @@ function renderAccountDialog(api, initialPayload, apply) { onConfirm: value => { const code = value.trim(); if (!code) { - renderMain(); + // Empty submit must keep the user inside the flow, not + // silently drop them back to the account list. Surface the + // validation message and reopen the same prompt so the + // user can paste again. + api.ui.toast({ + message: 'Please paste the callback URL or authorization code.' + }); + openOAuthCodePrompt(oauthUrl); return; } openOAuthLabelPrompt(code, oauthUrl); diff --git a/packages/opencode/src/tui/command-dialogs.test.tsx b/packages/opencode/src/tui/command-dialogs.test.tsx index 6daa9ea..3615254 100644 --- a/packages/opencode/src/tui/command-dialogs.test.tsx +++ b/packages/opencode/src/tui/command-dialogs.test.tsx @@ -978,6 +978,14 @@ describe('openCommandDialog (imperative dispatcher)', () => { expect(applyMock.mock.calls.at(0)?.[1]).toBe('add-oauth-start') expect(applyMock.mock.calls.at(0)?.[2]?.timeoutMs).toBe(120_000) expect(localFake.capturedSelectProps?.title).toBe('OAuth sign-in') + // The OAuth URL from the mocked start result must surface in the + // dialog's Copy-URL option so the user can paste it into a browser. + const copyUrlOption = localFake.capturedSelectProps?.options?.find( + (option) => option.title === 'Copy URL to clipboard', + ) + expect(copyUrlOption?.description).toContain( + 'https://accounts.google.test/authorize', + ) const enterCode = localFake.capturedSelectProps?.options?.find( (option) => option.title === 'Enter sign-in code', diff --git a/packages/opencode/src/tui/command-dialogs.tsx b/packages/opencode/src/tui/command-dialogs.tsx index 9a2b1a8..b8bac5c 100644 --- a/packages/opencode/src/tui/command-dialogs.tsx +++ b/packages/opencode/src/tui/command-dialogs.tsx @@ -309,7 +309,14 @@ function renderAccountDialog( onConfirm={(value: string) => { const code = value.trim() if (!code) { - renderMain() + // Empty submit must keep the user inside the flow, not + // silently drop them back to the account list. Surface the + // validation message and reopen the same prompt so the + // user can paste again. + api.ui.toast({ + message: 'Please paste the callback URL or authorization code.', + }) + openOAuthCodePrompt(oauthUrl) return } openOAuthLabelPrompt(code, oauthUrl) From c0ae5f2c5f2387fe751a42714d003d81727770fd Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 20:47:37 +0200 Subject: [PATCH 10/25] feat(quota): adopt retrieveUserQuotaSummary as the windowed quota source MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add QuotaWindowEntry to QuotaGroupSummary + AccountMetadataV3 cachedQuota type - Add fetchQuotaSummary() calling v1internal:retrieveUserQuotaSummary - Map RUQS groups to pools by bucketId prefix (gemini-* → gemini, 3p-* → non-gemini) - Derive pool remainingFraction/resetTime from most-constrained window - Swap primary quota fetch to RUQS; fall back to fetchAvailableModels on 403/failure - Use managedProjectId from refresh token parts as primary; fall back to projectId - Fix quota.test.ts event count for the added fallback fetch call --- packages/core/src/account-types.ts | 12 +- packages/core/src/quota-manager.ts | 236 +++++++++++++++++++++ packages/core/src/quota-types.ts | 16 ++ packages/opencode/src/plugin/quota.test.ts | 1 + packages/opencode/src/plugin/quota.ts | 80 +++++-- 5 files changed, 321 insertions(+), 24 deletions(-) diff --git a/packages/core/src/account-types.ts b/packages/core/src/account-types.ts index 916a515..b85c510 100644 --- a/packages/core/src/account-types.ts +++ b/packages/core/src/account-types.ts @@ -69,7 +69,17 @@ export interface AccountMetadataV3 { /** Cached soft quota data (group-level aggregation) */ cachedQuota?: Record< string, - { remainingFraction?: number; resetTime?: string; modelCount: number } + { + remainingFraction?: number + resetTime?: string + modelCount: number + /** Per-window breakdown from the RUQS response. Omitted in legacy shapes. */ + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } > /** Cached per-model quota data (individual model granularity) */ cachedPerModelQuota?: { diff --git a/packages/core/src/quota-manager.ts b/packages/core/src/quota-manager.ts index 763300f..821a854 100644 --- a/packages/core/src/quota-manager.ts +++ b/packages/core/src/quota-manager.ts @@ -26,6 +26,7 @@ import type { QuotaGroup, QuotaGroupSummary, QuotaSummary, + QuotaWindowEntry, } from './quota-types.ts' const log = createLogger('quota-manager') @@ -550,6 +551,241 @@ export interface FetchAvailableModelsResponse { models?: Record } +// ============================================================================ +// retrieveUserQuotaSummary — windowed quota source (2 pools × variable windows) +// ============================================================================ + +export interface RetrieveUserQuotaSummaryBucket { + bucketId: string + displayName: string + window: 'weekly' | '5h' + resetTime: string + remainingFraction: number + description?: string +} + +export interface RetrieveUserQuotaSummaryGroup { + displayName: string + description?: string + buckets: RetrieveUserQuotaSummaryBucket[] +} + +export interface RetrieveUserQuotaSummaryResponse { + groups: RetrieveUserQuotaSummaryGroup[] + description?: string +} + +/** + * Derive the most-constrained window from a set of window entries. + * Returns the entry with the smallest `remainingFraction` — this is the + * binding constraint for the pool. `resetTime` comes from the same window. + * Returns `undefined` when there are no windows. + */ +function mostConstrainedWindow( + windows: QuotaWindowEntry[], +): { remainingFraction: number; resetTime: string } | undefined { + if (windows.length === 0) return undefined + let best = windows[0]! + for (let i = 1; i < windows.length; i++) { + if (windows[i]!.remainingFraction < best.remainingFraction) { + best = windows[i]! + } + } + return { + remainingFraction: best.remainingFraction, + resetTime: best.resetTime, + } +} + +/** + * Map a retrieveUserQuotaSummary bucketId prefix to our internal pool. + * + * bucketId prefixes: + * `gemini-*` → gemini + * `3p-*` → non-gemini + */ +function poolForBucketId(bucketId: string): QuotaGroup | null { + if (bucketId.startsWith('gemini-')) return 'gemini' + if (bucketId.startsWith('3p-')) return 'non-gemini' + return null +} + +/** + * Aggregate a retrieveUserQuotaSummary response into a QuotaSummary. + * + * Each RUQS group maps to a pool via bucketId prefix. Within a pool, + * windows are stored ordered (weekly first, then 5h). The pool's + * `remainingFraction`/`resetTime` derive from the most-constrained window. + */ +export function aggregateQuotaSummary( + response: RetrieveUserQuotaSummaryResponse, +): QuotaSummary { + const groups: Partial> = {} + let totalCount = 0 + + for (const group of response.groups) { + const windows: QuotaWindowEntry[] = [] + for (const bucket of group.buckets) { + const pool = poolForBucketId(bucket.bucketId) + if (!pool) continue + windows.push({ + window: bucket.window, + remainingFraction: normalizeRemainingFraction(bucket.remainingFraction), + resetTime: bucket.resetTime, + }) + } + if (windows.length === 0) continue + + // Order: weekly first, then 5h. + windows.sort((a, b) => { + const order: Record = { weekly: 0, '5h': 1 } + return (order[a.window] ?? 2) - (order[b.window] ?? 2) + }) + + const constrained = mostConstrainedWindow(windows) + const firstBucket = group.buckets[0] + if (!firstBucket) continue + const pool = poolForBucketId(firstBucket.bucketId) + if (!pool || !constrained) continue + + const modelCount = group.description + ? (group.description.match(/[^,:]+/g)?.length ?? 0) + : 0 + + groups[pool] = { + remainingFraction: constrained.remainingFraction, + resetTime: constrained.resetTime, + modelCount, + windows, + } + totalCount += modelCount + } + + return { groups, modelCount: totalCount } +} + +export interface FetchQuotaSummaryOptions { + accessToken: string + /** Managed project ID. Falls back to regular projectId on 403. */ + managedProjectId?: string + /** Regular project ID — used as fallback when managedProjectId is missing or returns 403. */ + projectId?: string + endpoints: readonly string[] + timeoutMs?: number + userAgent?: string + fetchVia?: ( + url: string, + options: RequestInit, + extra: { timeoutMs: number; signal?: AbortSignal | null }, + ) => Promise +} + +export interface FetchQuotaSummaryResult { + summary: RetrieveUserQuotaSummaryResponse + /** True when the result came from the legacy fallback path. */ + fellBackToLegacy?: boolean +} + +/** + * Fetch the windowed quota summary via `retrieveUserQuotaSummary`. + * + * Uses the same transport/UA/timeout conventions as `fetchAvailableModels`. + * On 403 with the managedProjectId, retries with the regular projectId. + * If that also 403s, falls back to `fetchAvailableModels` so quota never + * goes dark. On missing managedProjectId, tries projectId first. + */ +export async function fetchQuotaSummary( + options: FetchQuotaSummaryOptions, +): Promise { + const timeoutMs = options.timeoutMs ?? QUOTA_MANAGER_DEFAULT_TIMEOUT_MS + const userAgent = options.userAgent ?? buildAntigravityHarnessUserAgent() + const transport = options.fetchVia ?? defaultTransport + const errors: string[] = [] + + const endpoint = options.endpoints[0] + if (!endpoint) { + throw new Error('No endpoints configured for fetchQuotaSummary') + } + + const tryBody = async ( + projectId: string, + ): Promise => { + const body = { project: projectId } + try { + const response = await transport( + `${endpoint}/v1internal:retrieveUserQuotaSummary`, + { + method: 'POST', + headers: { + 'User-Agent': userAgent, + Authorization: `Bearer ${options.accessToken}`, + 'Content-Type': 'application/json', + 'Accept-Encoding': 'gzip', + }, + body: JSON.stringify(body), + }, + { timeoutMs }, + ) + + if (response.ok) { + return (await response.json()) as RetrieveUserQuotaSummaryResponse + } + + const status = response.status + if (status === 403) { + errors.push( + `retrieveUserQuotaSummary 403 at ${endpoint} (project=${projectId.slice(0, 12)}…)`, + ) + return null + } + + if (status === 429 || status >= 500) { + const message = await response.text().catch(() => '') + errors.push( + `retrieveUserQuotaSummary ${status} at ${endpoint}${message ? `: ${message.trim().slice(0, 200)}` : ''}`, + ) + return null + } + + const message = await response.text().catch(() => '') + errors.push( + `retrieveUserQuotaSummary ${status} at ${endpoint}${message ? `: ${message.trim().slice(0, 200)}` : ''}`, + ) + return null + } catch (error) { + errors.push( + `retrieveUserQuotaSummary network error at ${endpoint}: ${error instanceof Error ? error.message : String(error)}`, + ) + return null + } + } + + // Try managedProjectId first, then projectId, then legacy fallback. + const primary = options.managedProjectId ?? options.projectId + if (primary) { + const result = await tryBody(primary) + if (result) return { summary: result } + } + + // If primary failed with 403 and we used managedProjectId, + // retry with regular projectId as fallback (only when distinct). + const fallbackId = + options.managedProjectId && + options.projectId && + options.managedProjectId !== options.projectId + ? options.projectId + : undefined + if (fallbackId) { + const result = await tryBody(fallbackId) + if (result) return { summary: result } + } + + // Give up — the caller should fall back to fetchAvailableModels. + throw new Error( + errors.join('; ') || 'fetchQuotaSummary failed: no project ID available', + ) +} + /** * Aggregate Gemini CLI quota buckets into a summary. */ diff --git a/packages/core/src/quota-types.ts b/packages/core/src/quota-types.ts index 53456ed..c1b3adc 100644 --- a/packages/core/src/quota-types.ts +++ b/packages/core/src/quota-types.ts @@ -11,10 +11,26 @@ import type { AccountMetadataV3 } from './account-types.ts' export type QuotaGroup = 'gemini' | 'non-gemini' +export type QuotaWindow = 'weekly' | '5h' + +export interface QuotaWindowEntry { + window: QuotaWindow + remainingFraction: number + resetTime: string +} + export interface QuotaGroupSummary { + /** Most-constrained window's remainingFraction (derived from `windows`). */ remainingFraction?: number + /** Most-constrained window's resetTime (derived from `windows`). */ resetTime?: string modelCount: number + /** + * Per-window breakdown from the retrieveUserQuotaSummary response. + * `weekly` first, then `5h`. Omitted in legacy cached shapes — + * consumers treat a single remainingFraction as one unlabeled window. + */ + windows?: QuotaWindowEntry[] } export interface PerModelQuotaEntry { diff --git a/packages/opencode/src/plugin/quota.test.ts b/packages/opencode/src/plugin/quota.test.ts index 4e7e3cd..e142ecb 100644 --- a/packages/opencode/src/plugin/quota.test.ts +++ b/packages/opencode/src/plugin/quota.test.ts @@ -256,6 +256,7 @@ describe('pushSidebarQuotaSnapshot', () => { }) await drainSidebarWrites() expect(events).toEqual([ + 'fetch:start', 'fetch:start', 'fetch:start', 'sidebar:write-start', diff --git a/packages/opencode/src/plugin/quota.ts b/packages/opencode/src/plugin/quota.ts index 4217f18..355cabd 100644 --- a/packages/opencode/src/plugin/quota.ts +++ b/packages/opencode/src/plugin/quota.ts @@ -20,13 +20,14 @@ import { type AccountQuotaResult, aggregateGeminiCliQuota, aggregateQuota, + aggregateQuotaSummary, createQuotaManager, defaultKeyOf, type FetchAccountQuota, type FetchAvailableModelsOptions, - type FetchAvailableModelsResponse, fetchAvailableModels, fetchGeminiCliQuota, + fetchQuotaSummary, type GeminiCliQuotaSummary, type QuotaManager, type QuotaSummary, @@ -372,36 +373,68 @@ function makeFetchAccountQuota( await persistRotatedRefresh(client, providerId, auth).catch(() => {}) } - const [antigravityResponse, geminiCliResponse] = await Promise.all([ - fetchAvailableModels({ + let quotaResult: QuotaSummary + let fellBackToLegacy = false + + const authParts = parseRefreshParts(auth.refresh) + const managedProjectId = authParts.managedProjectId + + // Primary: retrieveUserQuotaSummary (windowed quota). + // Fall back to fetchAvailableModels on any failure (403, network, etc.) + // so quota never goes dark. + try { + const summaryResult = await fetchQuotaSummary({ accessToken: auth.access ?? '', + managedProjectId, projectId: projectContext.effectiveProjectId, endpoints: ANTIGRAVITY_ENDPOINT_FALLBACKS, userAgent: buildAntigravityHarnessUserAgent(), timeoutMs: 10_000, ...(fetchVia ? { fetchVia } : {}), - }).catch((): FetchAvailableModelsResponse => ({ models: undefined })), - fetchGeminiCliQuota({ - accessToken: auth.access ?? '', - projectId: projectContext.effectiveProjectId, - endpoints: ANTIGRAVITY_ENDPOINT_FALLBACKS, - userAgent: buildGeminiCliUserAgent(), - timeoutMs: 10_000, - ...(fetchVia ? { fetchVia } : {}), - }), - ]) - - let quotaResult: QuotaSummary - if (antigravityResponse.models === undefined) { - quotaResult = { - groups: {}, - modelCount: 0, - error: 'Failed to fetch Antigravity quota', + }) + quotaResult = aggregateQuotaSummary(summaryResult.summary) + if (summaryResult.fellBackToLegacy) { + fellBackToLegacy = true + } + } catch { + // Fall back to fetchAvailableModels-based path. + fellBackToLegacy = true + try { + const modelsResponse = await fetchAvailableModels({ + accessToken: auth.access ?? '', + projectId: projectContext.effectiveProjectId, + endpoints: ANTIGRAVITY_ENDPOINT_FALLBACKS, + userAgent: buildAntigravityHarnessUserAgent(), + timeoutMs: 10_000, + ...(fetchVia ? { fetchVia } : {}), + }) + if (modelsResponse.models) { + quotaResult = aggregateQuota(modelsResponse.models) + } else { + quotaResult = { + groups: {}, + modelCount: 0, + error: 'Failed to fetch Antigravity quota (legacy fallback)', + } + } + } catch { + quotaResult = { + groups: {}, + modelCount: 0, + error: 'Failed to fetch Antigravity quota', + } } - } else { - quotaResult = aggregateQuota(antigravityResponse.models) } + const geminiCliResponse = await fetchGeminiCliQuota({ + accessToken: auth.access ?? '', + projectId: projectContext.effectiveProjectId, + endpoints: ANTIGRAVITY_ENDPOINT_FALLBACKS, + userAgent: buildGeminiCliUserAgent(), + timeoutMs: 10_000, + ...(fetchVia ? { fetchVia } : {}), + }) + const geminiCliQuotaResult = aggregateGeminiCliQuota(geminiCliResponse) const annotated: GeminiCliQuotaSummary = geminiCliResponse.buckets === undefined || @@ -420,7 +453,8 @@ function makeFetchAccountQuota( logQuotaStatus(account.email, index, remainingPercent, family) } - logQuotaFetch('complete', 1, 'ok=1 errors=0') + const legacyTag = fellBackToLegacy ? ' legacy=1' : '' + logQuotaFetch('complete', 1, `ok=1 errors=0${legacyTag}`) return { index, From 2c0d8be9ec2ea62fa6bef8ac85bd258ce8036fd6 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 20:47:48 +0200 Subject: [PATCH 11/25] feat(tui): render per-window quota bars (weekly/5h) per pool - AccountBlock iterates per-window entries, one QuotaRow per window - Pool prefix (Gm/NG) on first window row; indented gutter labels for rest - Gutter labels: 7d (weekly) / 5h (5h) matching the fleet's convention - Collapsed row shows pool + binding-window label + used% - SidebarQuotaEntry gains windows array (redacted) - SidebarAccountRedactionInput.cachedQuota supports windows shape - Tolerant read: pre-windows snapshots render a single QuotaRow as before --- packages/opencode/src/sidebar-state.ts | 46 ++++++++++++- .../src/tui-compiled/sidebar-state.ts | 46 ++++++++++++- packages/opencode/src/tui-compiled/tui.tsx | 67 ++++++++++++++----- packages/opencode/src/tui.tsx | 66 +++++++++++++++--- 4 files changed, 192 insertions(+), 33 deletions(-) diff --git a/packages/opencode/src/sidebar-state.ts b/packages/opencode/src/sidebar-state.ts index 108aeb5..68524db 100644 --- a/packages/opencode/src/sidebar-state.ts +++ b/packages/opencode/src/sidebar-state.ts @@ -54,9 +54,18 @@ export const SIDEBAR_STATE_VERSION = 1 as const export type SidebarQuotaKey = 'gemini' | 'non-gemini' +export interface SidebarQuotaWindowEntry { + window: 'weekly' | '5h' + remainingPercent: number + resetAt?: number +} + export interface SidebarQuotaEntry { + /** Most-constrained window's remaining % (derived, for back-compat). */ remainingPercent: number resetAt?: number + /** Per-window breakdown. Omitted in pre-windows snapshots. */ + windows?: SidebarQuotaWindowEntry[] } export interface SidebarAccountState { @@ -432,8 +441,24 @@ export interface SidebarAccountRedactionInput { /** Health score in `[0, 100]`. Defaults to 100 when missing. */ healthScore?: number cachedQuota?: { - gemini?: { remainingFraction?: number; resetTime?: string } - 'non-gemini'?: { remainingFraction?: number; resetTime?: string } + gemini?: { + remainingFraction?: number + resetTime?: string + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } + 'non-gemini'?: { + remainingFraction?: number + resetTime?: string + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } } /** * Opaque identity stamp that was attached to the persisted quota snapshot. @@ -497,7 +522,22 @@ export function redactAccountForSidebar( const parsed = Date.parse(entry.resetTime) if (Number.isFinite(parsed)) resetAt = parsed } - quota[key] = { remainingPercent, resetAt } + const windows = entry.windows?.map((w) => ({ + window: w.window, + remainingPercent: clampNumber( + Math.round(w.remainingFraction * 100), + 0, + 100, + ), + resetAt: + typeof w.resetTime === 'string' && w.resetTime.length > 0 + ? (() => { + const parsed = Date.parse(w.resetTime) + return Number.isFinite(parsed) ? parsed : undefined + })() + : undefined, + })) + quota[key] = { remainingPercent, resetAt, windows } } } diff --git a/packages/opencode/src/tui-compiled/sidebar-state.ts b/packages/opencode/src/tui-compiled/sidebar-state.ts index 108aeb5..68524db 100644 --- a/packages/opencode/src/tui-compiled/sidebar-state.ts +++ b/packages/opencode/src/tui-compiled/sidebar-state.ts @@ -54,9 +54,18 @@ export const SIDEBAR_STATE_VERSION = 1 as const export type SidebarQuotaKey = 'gemini' | 'non-gemini' +export interface SidebarQuotaWindowEntry { + window: 'weekly' | '5h' + remainingPercent: number + resetAt?: number +} + export interface SidebarQuotaEntry { + /** Most-constrained window's remaining % (derived, for back-compat). */ remainingPercent: number resetAt?: number + /** Per-window breakdown. Omitted in pre-windows snapshots. */ + windows?: SidebarQuotaWindowEntry[] } export interface SidebarAccountState { @@ -432,8 +441,24 @@ export interface SidebarAccountRedactionInput { /** Health score in `[0, 100]`. Defaults to 100 when missing. */ healthScore?: number cachedQuota?: { - gemini?: { remainingFraction?: number; resetTime?: string } - 'non-gemini'?: { remainingFraction?: number; resetTime?: string } + gemini?: { + remainingFraction?: number + resetTime?: string + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } + 'non-gemini'?: { + remainingFraction?: number + resetTime?: string + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } } /** * Opaque identity stamp that was attached to the persisted quota snapshot. @@ -497,7 +522,22 @@ export function redactAccountForSidebar( const parsed = Date.parse(entry.resetTime) if (Number.isFinite(parsed)) resetAt = parsed } - quota[key] = { remainingPercent, resetAt } + const windows = entry.windows?.map((w) => ({ + window: w.window, + remainingPercent: clampNumber( + Math.round(w.remainingFraction * 100), + 0, + 100, + ), + resetAt: + typeof w.resetTime === 'string' && w.resetTime.length > 0 + ? (() => { + const parsed = Date.parse(w.resetTime) + return Number.isFinite(parsed) ? parsed : undefined + })() + : undefined, + })) + quota[key] = { remainingPercent, resetAt, windows } } } diff --git a/packages/opencode/src/tui-compiled/tui.tsx b/packages/opencode/src/tui-compiled/tui.tsx index fed0d22..d5c6428 100644 --- a/packages/opencode/src/tui-compiled/tui.tsx +++ b/packages/opencode/src/tui-compiled/tui.tsx @@ -103,6 +103,10 @@ const QUOTA_LABELS = { 'non-gemini': 'NG' }; const QUOTA_ORDER = ['gemini', 'non-gemini']; +const WINDOW_GUTTER = { + weekly: '7d', + '5h': '5h' +}; // --- Theme tokens ---------------------------------------------------------- // @@ -570,23 +574,49 @@ function AccountBlock(props) { _$insert(_el$24, statusWord); _$insert(_el$19, _$createComponent(For, { each: QUOTA_ORDER, - children: key => _$createComponent(QuotaRow, { - get theme() { - return props.theme; - }, - get appearance() { - return props.appearance; - }, - get label() { - return QUOTA_LABELS[key]; - }, - get entry() { - return props.account.quota[key]; - }, - get now() { - return props.now; + children: key => { + const poolEntry = props.account.quota[key]; + const windows = poolEntry?.windows; + // Legacy: no windows array — render a single row with the pool label. + const hasWindows = windows && windows.length > 0; + if (!hasWindows) { + return _$createComponent(QuotaRow, { + get theme() { + return props.theme; + }, + get appearance() { + return props.appearance; + }, + get label() { + return QUOTA_LABELS[key]; + }, + entry: poolEntry, + get now() { + return props.now; + } + }); } - }) + return _$memo(() => windows.map((w, wi) => _$createComponent(QuotaRow, { + get theme() { + return props.theme; + }, + get appearance() { + return props.appearance; + }, + get label() { + return wi === 0 ? `${QUOTA_LABELS[key]} ${WINDOW_GUTTER[w.window] ?? w.window}` : ` ${WINDOW_GUTTER[w.window] ?? w.window}`; + }, + get entry() { + return { + remainingPercent: w.remainingPercent, + resetAt: w.resetAt + }; + }, + get now() { + return props.now; + } + }))); + } }), _el$25); _$insertNode(_el$25, _el$26); _$setProp(_el$25, "width", '100%'); @@ -843,7 +873,10 @@ export function SidebarPanel(props) { const entry = account()?.quota[key]; if (!entry) return `${QUOTA_LABELS[key]}: —`; const pct = Math.round(100 - clamp(entry.remainingPercent, 0, 100)); - return `${QUOTA_LABELS[key]}: ${pct}%`; + // Show binding window when available, else just pool label. + const bindingWindow = entry.windows?.reduce((best, w) => best === null || w.remainingPercent < best.remainingPercent ? w : best, null); + const gutter = bindingWindow ? ` ${WINDOW_GUTTER[bindingWindow.window] ?? bindingWindow.window}` : ''; + return `${QUOTA_LABELS[key]}${gutter}: ${pct}%`; }; const unavailable = () => { const selected = account(); diff --git a/packages/opencode/src/tui.tsx b/packages/opencode/src/tui.tsx index 48bc2be..e172fee 100644 --- a/packages/opencode/src/tui.tsx +++ b/packages/opencode/src/tui.tsx @@ -134,6 +134,11 @@ const QUOTA_LABELS: Record = { const QUOTA_ORDER: readonly SidebarQuotaKey[] = ['gemini', 'non-gemini'] +const WINDOW_GUTTER: Record = { + weekly: '7d', + '5h': '5h', +} + // --- Theme tokens ---------------------------------------------------------- // // The sidebar pulls its colors from the live host theme (`api.theme.current` @@ -603,15 +608,43 @@ function AccountBlock(props: { - {(key) => ( - - )} + {(key) => { + const poolEntry = props.account.quota[key] + const windows = poolEntry?.windows + // Legacy: no windows array — render a single row with the pool label. + const hasWindows = windows && windows.length > 0 + if (!hasWindows) { + return ( + + ) + } + return ( + <> + {windows.map((w, wi) => ( + + ))} + + ) + }} {` ${healthText()}`} @@ -836,7 +869,20 @@ export function SidebarPanel(props: SidebarPanelProps): JSX.Element { const entry = account()?.quota[key] if (!entry) return `${QUOTA_LABELS[key]}: —` const pct = Math.round(100 - clamp(entry.remainingPercent, 0, 100)) - return `${QUOTA_LABELS[key]}: ${pct}%` + // Show binding window when available, else just pool label. + const bindingWindow = entry.windows?.reduce< + (typeof entry.windows)[0] | null + >( + (best, w) => + best === null || w.remainingPercent < best.remainingPercent + ? w + : best, + null, + ) + const gutter = bindingWindow + ? ` ${WINDOW_GUTTER[bindingWindow.window] ?? bindingWindow.window}` + : '' + return `${QUOTA_LABELS[key]}${gutter}: ${pct}%` } const unavailable = () => { const selected = account() From d645a7c8b2bfb5d6a4ac0eb6beaa1fdd54daa580 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 20:47:53 +0200 Subject: [PATCH 12/25] test: window-aware quota fixtures + e2e summary route MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Copy real Pro (weekly+5h) and Free (weekly-only) RUQS response dumps - aggregateQuotaSummary tests: pool mapping, window order, most-constrained derivation - fetchQuotaSummary tests: managedProjectId primary, 403→projectId fallback, missing ID - Add quotaSummaryWindow fixture to mock-antigravity-server for e2e coverage --- .../src/__fixtures__/quota/free-ruqs.json | 32 +++ .../core/src/__fixtures__/quota/pro-ruqs.json | 49 ++++ packages/core/src/quota-manager.test.ts | 247 +++++++++++++++++- .../e2e-tests/src/mock-antigravity-server.ts | 26 ++ 4 files changed, 353 insertions(+), 1 deletion(-) create mode 100644 packages/core/src/__fixtures__/quota/free-ruqs.json create mode 100644 packages/core/src/__fixtures__/quota/pro-ruqs.json diff --git a/packages/core/src/__fixtures__/quota/free-ruqs.json b/packages/core/src/__fixtures__/quota/free-ruqs.json new file mode 100644 index 0000000..fae59ea --- /dev/null +++ b/packages/core/src/__fixtures__/quota/free-ruqs.json @@ -0,0 +1,32 @@ +{ + "groups": [ + { + "buckets": [ + { + "bucketId": "gemini-weekly", + "displayName": "Weekly Limit", + "window": "weekly", + "resetTime": "2026-07-31T15:54:18Z", + "description": "You have used some of your weekly limit, it will fully refresh in 6 days, 22 hours.", + "remainingFraction": 0.8875432 + } + ], + "displayName": "Gemini Models", + "description": "Models within this group: Gemini Flash, Gemini Pro" + }, + { + "buckets": [ + { + "bucketId": "3p-weekly", + "displayName": "Weekly Limit", + "window": "weekly", + "resetTime": "2026-07-31T17:19:52Z", + "remainingFraction": 1 + } + ], + "displayName": "Claude and GPT models", + "description": "Models within this group: Claude Opus, Claude Sonnet, GPT-OSS" + } + ], + "description": "Within each group, models share a weekly limit. Quota is consumed proportionally to the cost of the tokens. Thus, limits will last longer with shorter tasks or using more cost-effective models. Your weekly limit is tied directly to your individual tier." +} diff --git a/packages/core/src/__fixtures__/quota/pro-ruqs.json b/packages/core/src/__fixtures__/quota/pro-ruqs.json new file mode 100644 index 0000000..26f83d8 --- /dev/null +++ b/packages/core/src/__fixtures__/quota/pro-ruqs.json @@ -0,0 +1,49 @@ +{ + "groups": [ + { + "buckets": [ + { + "bucketId": "gemini-weekly", + "displayName": "Weekly Limit", + "window": "weekly", + "resetTime": "2026-07-28T18:24:21Z", + "description": "You have used some of your weekly limit, it will fully refresh in 4 days, 1 hour.", + "remainingFraction": 0.9214085 + }, + { + "bucketId": "gemini-5h", + "displayName": "Five Hour Limit", + "window": "5h", + "resetTime": "2026-07-24T20:43:21Z", + "description": "You have used some of your 5-hour limit, it will fully refresh in 3 hours, 25 minutes.", + "remainingFraction": 0.9885728 + } + ], + "displayName": "Gemini Models", + "description": "Models within this group: Gemini Flash, Gemini Pro" + }, + { + "buckets": [ + { + "bucketId": "3p-weekly", + "displayName": "Weekly Limit", + "window": "weekly", + "resetTime": "2026-07-31T13:41:52Z", + "description": "You have used some of your weekly limit, it will fully refresh in 6 days, 20 hours.", + "remainingFraction": 0.9852099 + }, + { + "bucketId": "3p-5h", + "displayName": "Five Hour Limit", + "window": "5h", + "resetTime": "2026-07-24T18:41:52Z", + "description": "You have used some of your 5-hour limit, it will fully refresh in 1 hour, 23 minutes.", + "remainingFraction": 0.9556296 + } + ], + "displayName": "Claude and GPT models", + "description": "Models within this group: Claude Opus, Claude Sonnet, GPT-OSS" + } + ], + "description": "Within each group, models share a weekly limit and a 5-hour limit. Quota is consumed proportionally to the cost of the tokens. Thus, limits will last longer with shorter tasks or using more cost-effective models. The 5-hour limit smooths out aggregate demand to fairly distribute global capacity across all users, while your weekly limit is tied directly to your individual tier." +} diff --git a/packages/core/src/quota-manager.test.ts b/packages/core/src/quota-manager.test.ts index 99b06aa..77dcb33 100644 --- a/packages/core/src/quota-manager.test.ts +++ b/packages/core/src/quota-manager.test.ts @@ -1,6 +1,14 @@ import { afterEach, describe, expect, it } from 'bun:test' +import { readFileSync } from 'node:fs' +import { join } from 'node:path' import type { AccountMetadataV3 } from './account-types.ts' -import { createQuotaManager, type FetchAccountQuota } from './quota-manager.ts' +import { + aggregateQuotaSummary, + createQuotaManager, + type FetchAccountQuota, + fetchQuotaSummary, + type RetrieveUserQuotaSummaryResponse, +} from './quota-manager.ts' import type { AccountQuotaResult } from './quota-types.ts' function makeAccount( @@ -776,3 +784,240 @@ describe('hashed log labels', () => { expect(hashedLogLabel('x', '')).toMatch(/^x [a-f0-9]{8}$/) }) }) + +// ============================================================================ +// aggregateQuotaSummary — windowed quota from retrieveUserQuotaSummary +// ============================================================================ + +function loadFixture(name: string): RetrieveUserQuotaSummaryResponse { + const path = join(import.meta.dir, '__fixtures__', 'quota', name) + const raw = readFileSync(path, 'utf8') + return JSON.parse(raw) as RetrieveUserQuotaSummaryResponse +} + +describe('aggregateQuotaSummary', () => { + it('maps Pro (weekly+5h) groups to pools by bucketId prefix', () => { + const response = loadFixture('pro-ruqs.json') + const summary = aggregateQuotaSummary(response) + + expect(summary.groups.gemini).toBeDefined() + expect(summary.groups['non-gemini']).toBeDefined() + + const gemini = summary.groups.gemini! + expect(gemini.windows).toHaveLength(2) + // weekly first + expect(gemini.windows![0]!.window).toBe('weekly') + expect(gemini.windows![0]!.remainingFraction).toBeCloseTo(0.9214, 3) + expect(gemini.windows![1]!.window).toBe('5h') + expect(gemini.windows![1]!.remainingFraction).toBeCloseTo(0.9886, 3) + + const nonGemini = summary.groups['non-gemini']! + expect(nonGemini.windows).toHaveLength(2) + expect(nonGemini.windows![0]!.window).toBe('weekly') + expect(nonGemini.windows![0]!.remainingFraction).toBeCloseTo(0.9852, 3) + expect(nonGemini.windows![1]!.window).toBe('5h') + expect(nonGemini.windows![1]!.remainingFraction).toBeCloseTo(0.9556, 3) + }) + + it('maps Free (weekly-only) groups to pools', () => { + const response = loadFixture('free-ruqs.json') + const summary = aggregateQuotaSummary(response) + + const gemini = summary.groups.gemini! + // Weekly-only: one window + expect(gemini.windows).toHaveLength(1) + expect(gemini.windows![0]!.window).toBe('weekly') + expect(gemini.windows![0]!.remainingFraction).toBeCloseTo(0.8875, 3) + + const nonGemini = summary.groups['non-gemini']! + expect(nonGemini.windows).toHaveLength(1) + expect(nonGemini.windows![0]!.window).toBe('weekly') + expect(nonGemini.windows![0]!.remainingFraction).toBe(1) + }) + + it('derives most-constrained remainingFraction and resetTime per pool', () => { + const response = loadFixture('pro-ruqs.json') + const summary = aggregateQuotaSummary(response) + + // Gemini: weekly=0.9214, 5h=0.9886 → weekly is more constrained + expect(summary.groups.gemini!.remainingFraction).toBeCloseTo(0.9214, 3) + expect(summary.groups.gemini!.resetTime).toBe('2026-07-28T18:24:21Z') + + // Non-gemini: weekly=0.9852, 5h=0.9556 → 5h is more constrained + expect(summary.groups['non-gemini']!.remainingFraction).toBeCloseTo( + 0.9556, + 3, + ) + expect(summary.groups['non-gemini']!.resetTime).toBe('2026-07-24T18:41:52Z') + }) + + it('preserves window order: weekly first, then 5h', () => { + const response = loadFixture('pro-ruqs.json') + const summary = aggregateQuotaSummary(response) + + for (const group of Object.values(summary.groups)) { + if (!group?.windows) continue + for (let i = 1; i < group.windows.length; i++) { + const prev = group.windows[i - 1]! + const curr = group.windows[i]! + // All weeklies come before any 5h + if (prev.window === 'weekly' && curr.window === '5h') continue + if (prev.window === 'weekly' && curr.window === 'weekly') continue + if (prev.window === '5h' && curr.window === '5h') continue + // A 5h before a weekly = wrong order + expect(prev.window).not.toBe('5h') + } + } + }) + + it('back-compat: treats a windows-less QuotaGroupSummary gracefully', () => { + // Legacy shapes may have no `windows` array — consumers read + // `remainingFraction`/`resetTime` directly which are the derived values. + const response = loadFixture('free-ruqs.json') + const summary = aggregateQuotaSummary(response) + + // `remainingFraction` is the most-constrained (and only) window + expect(summary.groups.gemini!.remainingFraction).toBe( + summary.groups.gemini!.windows![0]!.remainingFraction, + ) + // Consumers that check `if (entry.windows)` skip per-window rendering + // and fall back to a single QuotaRow — this is the tolerant-read path. + }) +}) + +// ============================================================================ +// fetchQuotaSummary — network layer with managedProjectId fallback +// ============================================================================ + +describe('fetchQuotaSummary', () => { + const ENDPOINTS = ['http://127.0.0.1:1'] as const + + it('returns the summary when the server responds 200', async () => { + const summary: RetrieveUserQuotaSummaryResponse = { + groups: [ + { + displayName: 'Gemini Models', + buckets: [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly Limit', + window: 'weekly', + resetTime: '2026-01-01T00:00:00Z', + remainingFraction: 0.5, + }, + ], + }, + ], + } + const fetchVia = async () => + new Response(JSON.stringify(summary), { status: 200 }) + const result = await fetchQuotaSummary({ + accessToken: 'tok', + managedProjectId: 'mp', + endpoints: ENDPOINTS, + fetchVia: fetchVia as any, + }) + expect(result.summary.groups).toHaveLength(1) + expect(result.summary.groups[0]!.buckets[0]!.bucketId).toBe('gemini-weekly') + }) + + it('falls back to projectId when managedProjectId returns 403', async () => { + const summary: RetrieveUserQuotaSummaryResponse = { + groups: [ + { + displayName: 'Claude and GPT models', + buckets: [ + { + bucketId: '3p-weekly', + displayName: 'Weekly Limit', + window: 'weekly', + resetTime: '2026-01-01T00:00:00Z', + remainingFraction: 0.3, + }, + ], + }, + ], + } + let triedManaged = false + const fetchVia = async ( + _url: string, + init: RequestInit, + ): Promise => { + const body = JSON.parse((init as any).body ?? '{}') + if (body.project === 'mp') { + triedManaged = true + return new Response(JSON.stringify({ error: 'PERMISSION_DENIED' }), { + status: 403, + }) + } + return new Response(JSON.stringify(summary), { status: 200 }) + } + const result = await fetchQuotaSummary({ + accessToken: 'tok', + managedProjectId: 'mp', + projectId: 'regular', + endpoints: ENDPOINTS, + fetchVia: fetchVia as any, + }) + expect(triedManaged).toBe(true) + expect(result.summary.groups[0]!.buckets[0]!.bucketId).toBe('3p-weekly') + }) + + it('throws when all project IDs fail (caller handles fallback)', async () => { + const fetchVia = async () => + new Response('{"error":"PERMISSION_DENIED"}', { status: 403 }) + await expect( + fetchQuotaSummary({ + accessToken: 'tok', + projectId: 'regular', + endpoints: ENDPOINTS, + fetchVia: fetchVia as any, + }), + ).rejects.toThrow() + }) + + it('uses managedProjectId as primary when both IDs are present', async () => { + let receivedProjectId = '' + const summary: RetrieveUserQuotaSummaryResponse = { + groups: [], + } + const fetchVia = async ( + _url: string, + init: RequestInit, + ): Promise => { + const body = JSON.parse((init as any).body ?? '{}') + receivedProjectId = body.project as string + return new Response(JSON.stringify(summary), { status: 200 }) + } + await fetchQuotaSummary({ + accessToken: 'tok', + managedProjectId: 'managed-proj', + projectId: 'regular-proj', + endpoints: ENDPOINTS, + fetchVia: fetchVia as any, + }) + expect(receivedProjectId).toBe('managed-proj') + }) + + it('falls back to projectId when managedProjectId is missing', async () => { + let receivedProjectId = '' + const summary: RetrieveUserQuotaSummaryResponse = { + groups: [], + } + const fetchVia = async ( + _url: string, + init: RequestInit, + ): Promise => { + const body = JSON.parse((init as any).body ?? '{}') + receivedProjectId = body.project as string + return new Response(JSON.stringify(summary), { status: 200 }) + } + await fetchQuotaSummary({ + accessToken: 'tok', + projectId: 'regular-proj', + endpoints: ENDPOINTS, + fetchVia: fetchVia as any, + }) + expect(receivedProjectId).toBe('regular-proj') + }) +}) diff --git a/packages/e2e-tests/src/mock-antigravity-server.ts b/packages/e2e-tests/src/mock-antigravity-server.ts index 15f475e..ec3d0b7 100644 --- a/packages/e2e-tests/src/mock-antigravity-server.ts +++ b/packages/e2e-tests/src/mock-antigravity-server.ts @@ -87,6 +87,23 @@ export interface GeminiCliQuotaFixture extends BaseFixture { buckets: Array<{ model: string; remainingFraction: number }> } +/** Windowed quota summary envelope — fed to `retrieveUserQuotaSummary`. */ +export interface QuotaSummaryWindowFixture extends BaseFixture { + kind: 'quotaSummaryWindow' + /** Groups with their per-window buckets. */ + groups: Array<{ + displayName: string + description?: string + buckets: Array<{ + bucketId: string + displayName: string + window: 'weekly' | '5h' + resetTime: string + remainingFraction: number + }> + }> +} + /** Chunked SSE stream — one chunk per `chunks` element. */ export interface StreamChunkedFixture extends BaseFixture { kind: 'streamChunked' @@ -130,6 +147,7 @@ export type Fixture = | GenerateContentFixture | ProjectDiscoveryFixture | QuotaSummaryFixture + | QuotaSummaryWindowFixture | GeminiCliQuotaFixture | StreamChunkedFixture | TokenExpiryFixture @@ -312,6 +330,14 @@ export async function startMockAntigravityServer(): Promise { }) return } + case 'quotaSummaryWindow': { + sendJson(response, fixture.status ?? 200, { + groups: fixture.groups, + description: + 'Within each group, models share a weekly limit and a 5-hour limit.', + }) + return + } case 'generateContent': { const candidates = [ { From d69078b24aa486ff3845a616cdaf5b68b729522c Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 21:11:32 +0200 Subject: [PATCH 13/25] fix(tui): carry quota windows through the sidebar read path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - normalizeQuota: deserialize windows array from on-disk record (was stripping it) - QuotaRow label width: 3→5 chars to fit pool+gutter labels (Gm 7d, Gm 5h, etc.) - AccountBlock: pool prefix on every window row (not just first) - Adopt tui-windows-frames.test.tsx (reviewer's render harness, 4 frame tests) - e2e rpc-tui-flow: enqueue quotaSummaryWindow fixture instead of legacy quotaSummary --- .../e2e-tests/src/rpc-tui-flow.e2e.test.ts | 30 +- packages/opencode/src/sidebar-state.ts | 25 ++ .../src/tui-compiled/sidebar-state.ts | 25 ++ packages/opencode/src/tui-compiled/tui.tsx | 8 +- .../opencode/src/tui-windows-frames.test.tsx | 352 ++++++++++++++++++ packages/opencode/src/tui.tsx | 12 +- 6 files changed, 438 insertions(+), 14 deletions(-) create mode 100644 packages/opencode/src/tui-windows-frames.test.tsx diff --git a/packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts b/packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts index 5682455..a364e07 100644 --- a/packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts +++ b/packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts @@ -121,9 +121,35 @@ describe('rpc / tui flow (e2e)', () => { kind: 'projectDiscovery', projectId: 'project-rpc', }) + // Primary: retrieveUserQuotaSummary (windowed). h.server.enqueue({ - kind: 'quotaSummary', - models: [{ id: 'gemini-3-flash', displayName: 'Gemini 3 Flash' }], + kind: 'quotaSummaryWindow', + groups: [ + { + displayName: 'Gemini Models', + buckets: [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly Limit', + window: 'weekly', + resetTime: '2026-07-31T00:00:00Z', + remainingFraction: 0.7, + }, + ], + }, + { + displayName: 'Claude and GPT models', + buckets: [ + { + bucketId: '3p-weekly', + displayName: 'Weekly Limit', + window: 'weekly', + resetTime: '2026-07-31T00:00:00Z', + remainingFraction: 0.8, + }, + ], + }, + ], }) h.server.enqueue({ kind: 'geminiCliQuota', diff --git a/packages/opencode/src/sidebar-state.ts b/packages/opencode/src/sidebar-state.ts index 68524db..edd7417 100644 --- a/packages/opencode/src/sidebar-state.ts +++ b/packages/opencode/src/sidebar-state.ts @@ -322,9 +322,34 @@ function normalizeQuota(input: unknown): SidebarQuotaEntry | null { const remaining = toFiniteNumber(record.remainingPercent) if (remaining === null) return null const resetAt = toFiniteNumber(record.resetAt) ?? undefined + + // Tolerant: deserialize windows array if present; drop malformed entries. + const windowsRaw = record.windows + let windows: SidebarQuotaEntry['windows'] + if (Array.isArray(windowsRaw)) { + const parsed = windowsRaw + .filter(isObject) + .map((w) => { + const win = (w as Record).window + const rp = toFiniteNumber( + (w as Record).remainingPercent, + ) + const ra = toFiniteNumber((w as Record).resetAt) + if ((win !== 'weekly' && win !== '5h') || rp === null) return null + return { + window: win as 'weekly' | '5h', + remainingPercent: clampNumber(rp, 0, 100), + resetAt: ra ?? undefined, + } + }) + .filter((v): v is NonNullable => v !== null) + if (parsed.length > 0) windows = parsed + } + return { remainingPercent: clampNumber(remaining, 0, 100), resetAt, + windows, } } diff --git a/packages/opencode/src/tui-compiled/sidebar-state.ts b/packages/opencode/src/tui-compiled/sidebar-state.ts index 68524db..edd7417 100644 --- a/packages/opencode/src/tui-compiled/sidebar-state.ts +++ b/packages/opencode/src/tui-compiled/sidebar-state.ts @@ -322,9 +322,34 @@ function normalizeQuota(input: unknown): SidebarQuotaEntry | null { const remaining = toFiniteNumber(record.remainingPercent) if (remaining === null) return null const resetAt = toFiniteNumber(record.resetAt) ?? undefined + + // Tolerant: deserialize windows array if present; drop malformed entries. + const windowsRaw = record.windows + let windows: SidebarQuotaEntry['windows'] + if (Array.isArray(windowsRaw)) { + const parsed = windowsRaw + .filter(isObject) + .map((w) => { + const win = (w as Record).window + const rp = toFiniteNumber( + (w as Record).remainingPercent, + ) + const ra = toFiniteNumber((w as Record).resetAt) + if ((win !== 'weekly' && win !== '5h') || rp === null) return null + return { + window: win as 'weekly' | '5h', + remainingPercent: clampNumber(rp, 0, 100), + resetAt: ra ?? undefined, + } + }) + .filter((v): v is NonNullable => v !== null) + if (parsed.length > 0) windows = parsed + } + return { remainingPercent: clampNumber(remaining, 0, 100), resetAt, + windows, } } diff --git a/packages/opencode/src/tui-compiled/tui.tsx b/packages/opencode/src/tui-compiled/tui.tsx index d5c6428..710f5dd 100644 --- a/packages/opencode/src/tui-compiled/tui.tsx +++ b/packages/opencode/src/tui-compiled/tui.tsx @@ -442,7 +442,7 @@ function QuotaRow(props) { _$setProp(_el$14, "width", '100%'); _$setProp(_el$14, "flexDirection", 'row'); _$setProp(_el$14, "justifyContent", 'space-between'); - _$insert(_el$15, () => props.label.padEnd(3)); + _$insert(_el$15, () => props.label.padEnd(5)); _$insertNode(_el$16, _$createTextNode(`—`)); _$effect(_p$ => { var _v$6 = props.theme().textMuted, @@ -471,7 +471,7 @@ function QuotaRow(props) { _$insertNode(_el$1, _el$11); _$insertNode(_el$1, _el$12); _$setProp(_el$1, "flexDirection", 'row'); - _$setProp(_el$10, "width", 3); + _$setProp(_el$10, "width", 5); _$setProp(_el$10, "flexShrink", 0); _$insert(_el$10, () => props.label); _$setProp(_el$11, "flexShrink", 0); @@ -596,7 +596,7 @@ function AccountBlock(props) { } }); } - return _$memo(() => windows.map((w, wi) => _$createComponent(QuotaRow, { + return _$memo(() => windows.map(w => _$createComponent(QuotaRow, { get theme() { return props.theme; }, @@ -604,7 +604,7 @@ function AccountBlock(props) { return props.appearance; }, get label() { - return wi === 0 ? `${QUOTA_LABELS[key]} ${WINDOW_GUTTER[w.window] ?? w.window}` : ` ${WINDOW_GUTTER[w.window] ?? w.window}`; + return `${QUOTA_LABELS[key]} ${WINDOW_GUTTER[w.window] ?? w.window}`; }, get entry() { return { diff --git a/packages/opencode/src/tui-windows-frames.test.tsx b/packages/opencode/src/tui-windows-frames.test.tsx new file mode 100644 index 0000000..780907b --- /dev/null +++ b/packages/opencode/src/tui-windows-frames.test.tsx @@ -0,0 +1,352 @@ +/** @jsxImportSource @opentui/solid */ + +/** + * Reviewer-only render harness for the windows rework. Mirrors the + * pattern used by `tui.test.tsx` (same `testRender` import, same + * `SidebarStateV1` shape, same `createSidebarController` factory) so the + * JSX pragma + Solid transform apply through the project's + * `@opentui/solid/preload` bunfig configuration. + */ + +import { afterEach, beforeEach, describe, expect, it } from 'bun:test' +import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { testRender } from '@opentui/solid' +import { + DEFAULT_SIDEBAR_STATE, + SIDEBAR_STATE_ENV, + SIDEBAR_STATE_VERSION, + type SidebarStateV1, +} from './sidebar-state' +import { + createSidebarController, + QuotaDialogContent, + SidebarPanel, +} from './tui' +import type { TuiLogger } from './tui/file-logger' +import { + type AntigravityAuthTuiPrefs, + DEFAULT_PREFS, + PLUGIN_KEY, + TUI_PREFS_FILE_ENV, +} from './tui-preferences' + +interface Fixture { + statePath: string + prefsPath: string + cleanup: () => void +} + +function makeFixture(): Fixture { + const root = mkdtempSync(join(tmpdir(), 'agy-tui-windows-')) + const statePath = join(root, 'sidebar-state.json') + const prefsPath = join(root, 'tui-preferences.jsonc') + process.env[SIDEBAR_STATE_ENV] = statePath + process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE = join(root, 'tui.log') + process.env[TUI_PREFS_FILE_ENV] = prefsPath + return { + statePath, + prefsPath, + cleanup: () => { + delete process.env[SIDEBAR_STATE_ENV] + delete process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE + delete process.env[TUI_PREFS_FILE_ENV] + rmSync(root, { recursive: true, force: true }) + }, + } +} + +const NOOP_LOGGER: TuiLogger = { + debug: () => undefined, + info: () => undefined, + warn: () => undefined, + error: () => undefined, + getLogPath: () => undefined, +} + +function writeState(state: Partial): SidebarStateV1 { + return { + ...DEFAULT_SIDEBAR_STATE, + ...state, + version: SIDEBAR_STATE_VERSION, + } +} + +function writePrefs( + prefsPath: string, + overrides: Partial, +): AntigravityAuthTuiPrefs { + const merged: AntigravityAuthTuiPrefs = { + ...DEFAULT_PREFS, + ...overrides, + header: { ...DEFAULT_PREFS.header, ...(overrides.header ?? {}) }, + sections: { ...DEFAULT_PREFS.sections, ...(overrides.sections ?? {}) }, + appearance: { + ...DEFAULT_PREFS.appearance, + ...(overrides.appearance ?? {}), + }, + } + const root = { [PLUGIN_KEY]: merged } + mkdirSync(join(prefsPath, '..'), { recursive: true }) + writeFileSync(prefsPath, JSON.stringify(root), 'utf-8') + return merged +} + +async function settle(): Promise { + await new Promise((resolve) => setTimeout(resolve, 20)) +} + +function dumpFrame(label: string, frame: string): string { + const banner = `=== ${label} ===` + const bar = '='.repeat(banner.length) + return [`${bar}`, banner, `${bar}`, frame, `${bar}`].join('\n') +} + +describe('windows rework — reviewer frames', () => { + let fixture: Fixture + + beforeEach(() => { + fixture = makeFixture() + }) + + afterEach(() => { + fixture.cleanup() + }) + + it('(a) Pro account — 2 pools × 2 windows (7d + 5h) with aligned gutters', async () => { + const future = Date.now() + 5 * 60 * 1000 + const future2 = Date.now() + 7 * 24 * 60 * 60 * 1000 + const payload = writeState({ + checkedAt: Date.now(), + routingAuthoritative: true, + accounts: [ + { + id: 'acc-1', + label: 'Pro Account', + enabled: true, + health: 80, + current: true, + quota: { + gemini: { + remainingPercent: 92, + resetAt: future2, + windows: [ + { window: 'weekly', remainingPercent: 92, resetAt: future2 }, + { window: '5h', remainingPercent: 99, resetAt: future }, + ], + }, + 'non-gemini': { + remainingPercent: 96, + resetAt: future2, + windows: [ + { window: 'weekly', remainingPercent: 99, resetAt: future2 }, + { window: '5h', remainingPercent: 96, resetAt: future }, + ], + }, + }, + }, + ], + }) + mkdirSync(join(fixture.statePath, '..'), { recursive: true }) + writeFileSync(fixture.statePath, JSON.stringify(payload), 'utf-8') + + const testSetup = await testRender( + () => , + { width: 60, height: 24 }, + ) + await testSetup.flush() + const frame = testSetup.captureCharFrame() + console.log(`\n${dumpFrame('(a) Pro Account frame', frame)}`) + expect(frame).toContain('Pro Account') + expect(frame).toContain('Gm') + expect(frame).toContain('NG') + expect(frame).toContain('7d') + expect(frame).toContain('5h') + expect(frame).toContain('Gm 7d') + expect(frame).toContain('Gm 5h') + expect(frame).toContain('NG 7d') + expect(frame).toContain('NG 5h') + testSetup.renderer.destroy() + }) + + it('(b) Free account — 2 pools × 1 window (weekly only, no fake 5h)', async () => { + const future = Date.now() + 7 * 24 * 60 * 60 * 1000 + const payload = writeState({ + checkedAt: Date.now(), + routingAuthoritative: true, + accounts: [ + { + id: 'acc-1', + label: 'Free Account', + enabled: true, + health: 80, + current: true, + quota: { + gemini: { + remainingPercent: 89, + resetAt: future, + windows: [ + { window: 'weekly', remainingPercent: 89, resetAt: future }, + ], + }, + 'non-gemini': { + remainingPercent: 100, + windows: [ + { window: 'weekly', remainingPercent: 100, resetAt: future }, + ], + }, + }, + }, + ], + }) + mkdirSync(join(fixture.statePath, '..'), { recursive: true }) + writeFileSync(fixture.statePath, JSON.stringify(payload), 'utf-8') + + const testSetup = await testRender( + () => , + { width: 60, height: 24 }, + ) + await testSetup.flush() + const frame = testSetup.captureCharFrame() + console.log(`\n${dumpFrame('(b) Free Account frame', frame)}`) + expect(frame).toContain('Free Account') + expect(frame).toContain('Gm') + expect(frame).toContain('NG') + expect(frame).toContain('Gm 7d') + expect(frame).toContain('NG 7d') + expect(frame).not.toContain('5h') + expect(frame).not.toContain('Gm 5h') + expect(frame).not.toContain('NG 5h') + testSetup.renderer.destroy() + }) + + it('(c) Collapsed row — pool + binding-window label + used%', async () => { + const future = Date.now() + 5 * 60 * 1000 + const future2 = Date.now() + 7 * 24 * 60 * 60 * 1000 + const payload = writeState({ + checkedAt: Date.now(), + routingAuthoritative: true, + accounts: [ + { + id: 'acc-1', + label: 'Pro Account', + enabled: true, + health: 80, + current: true, + quota: { + gemini: { + remainingPercent: 92, + resetAt: future2, + windows: [ + { window: 'weekly', remainingPercent: 92, resetAt: future2 }, + { window: '5h', remainingPercent: 99, resetAt: future }, + ], + }, + 'non-gemini': { + remainingPercent: 96, + resetAt: future2, + windows: [ + { window: 'weekly', remainingPercent: 99, resetAt: future2 }, + { window: '5h', remainingPercent: 96, resetAt: future }, + ], + }, + }, + }, + ], + }) + mkdirSync(join(fixture.statePath, '..'), { recursive: true }) + writeFileSync(fixture.statePath, JSON.stringify(payload), 'utf-8') + + const prefs = writePrefs(fixture.prefsPath, { + collapsed: true, + rememberCollapsed: true, + }) + const controller = createSidebarController(prefs) + + const testSetup = await testRender( + () => ( + + ), + { width: 80, height: 12 }, + ) + await testSetup.flush() + const frame = testSetup.captureCharFrame() + console.log(`\n${dumpFrame('(c) Collapsed row frame', frame)}`) + // Collapsed row should show binding-window label + used% per pool. + expect(frame).toContain('Gm 7d:') + expect(frame).toContain('NG') + expect(frame).toContain('●') + expect(frame).toContain('8%') + expect(frame).toContain('4%') + testSetup.renderer.destroy() + }) + + it('(d) QuotaDialogContent — modal scope, same per-window rows', async () => { + const future = Date.now() + 5 * 60 * 1000 + const future2 = Date.now() + 7 * 24 * 60 * 60 * 1000 + const payload = writeState({ + checkedAt: Date.now(), + routingAuthoritative: true, + accounts: [ + { + id: 'acc-1', + label: 'Pro Account', + enabled: true, + health: 80, + current: true, + quota: { + gemini: { + remainingPercent: 92, + resetAt: future2, + windows: [ + { window: 'weekly', remainingPercent: 92, resetAt: future2 }, + { window: '5h', remainingPercent: 99, resetAt: future }, + ], + }, + 'non-gemini': { + remainingPercent: 96, + resetAt: future2, + windows: [ + { window: 'weekly', remainingPercent: 99, resetAt: future2 }, + { window: '5h', remainingPercent: 96, resetAt: future }, + ], + }, + }, + }, + ], + }) + mkdirSync(join(fixture.statePath, '..'), { recursive: true }) + writeFileSync(fixture.statePath, JSON.stringify(payload), 'utf-8') + + const controller = createSidebarController(DEFAULT_PREFS) + const testSetup = await testRender( + () => ( + + ), + { width: 80, height: 20 }, + ) + await settle() + const frame = testSetup.captureCharFrame() + console.log(`\n${dumpFrame('(d) QuotaDialogContent frame', frame)}`) + expect(frame).toContain('Antigravity Quota') + expect(frame).toContain('Pro Account') + expect(frame).toContain('Gm 7d') + expect(frame).toContain('Gm 5h') + expect(frame).toContain('NG 7d') + expect(frame).toContain('NG 5h') + testSetup.renderer.destroy() + }) +}) diff --git a/packages/opencode/src/tui.tsx b/packages/opencode/src/tui.tsx index e172fee..b8e0deb 100644 --- a/packages/opencode/src/tui.tsx +++ b/packages/opencode/src/tui.tsx @@ -515,14 +515,14 @@ function QuotaRow(props: { when={used() != null} fallback={ - {props.label.padEnd(3)} + {props.label.padEnd(5)} {'\u2014'} } > - + {props.label} - {windows.map((w, wi) => ( + {windows.map((w) => ( Date: Fri, 24 Jul 2026 21:36:19 +0200 Subject: [PATCH 14/25] fix(quota): project cached windows through every sidebar producer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Centralize pool→sidebar projection in projectQuotaPoolForSidebar (sidebar-state.ts). redactAccountForSidebar uses it for the canonical {remainingPercent, resetAt, windows} conversion. toCommandAccountRow now carries windows from cachedQuota entries. writeSidebar projects them through into the sidebar input shape. Add 5 producer-seam tests: windowed→redacted, legacy tolerance, Free weekly-only, full write→read round-trip, identity-mismatch drops entire cachedQuota. --- packages/opencode/src/plugin/command-data.ts | 56 +++++- packages/opencode/src/sidebar-state.test.ts | 165 ++++++++++++++++++ packages/opencode/src/sidebar-state.ts | 79 ++++++--- .../src/tui-compiled/plugin/command-data.ts | 56 +++++- .../src/tui-compiled/sidebar-state.ts | 79 ++++++--- 5 files changed, 371 insertions(+), 64 deletions(-) diff --git a/packages/opencode/src/plugin/command-data.ts b/packages/opencode/src/plugin/command-data.ts index 210310e..bf80ef0 100644 --- a/packages/opencode/src/plugin/command-data.ts +++ b/packages/opencode/src/plugin/command-data.ts @@ -82,6 +82,11 @@ type CommandDataQuotaGroupSummary = { remainingFraction?: number resetTime?: string modelCount: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> } type CommandDataAccountQuotaResult = { @@ -137,6 +142,11 @@ export interface CommandAccountRow { label: string remainingPercent: number | null resetAt?: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingPercent: number + resetAt?: number + }> }> } @@ -198,11 +208,27 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { const parsed = Date.parse(cachedEntry.resetTime) if (Number.isFinite(parsed)) resetAt = parsed } + // Carry per-window breakdown when the cache was populated from the + // windowed RUQS path — the sidebar writer projects through here. + const windows = cachedEntry.windows?.length + ? cachedEntry.windows.map((w) => ({ + window: w.window, + remainingPercent: Math.round(w.remainingFraction * 100), + resetAt: + typeof w.resetTime === 'string' && w.resetTime.length > 0 + ? (() => { + const parsed = Date.parse(w.resetTime) + return Number.isFinite(parsed) ? parsed : undefined + })() + : undefined, + })) + : undefined quota.push({ key, label: QUOTA_GROUP_LABELS[key], remainingPercent, resetAt, + windows, }) } const label = `Account ${entry.index + 1}` @@ -418,19 +444,33 @@ export function createCommandDataService( const accounts: SidebarAccountRedactionInput[] = rows.map((row) => { const gemini = row.quota.find((q) => q.key === 'gemini') const nonGemini = row.quota.find((q) => q.key === 'non-gemini') - const toFraction = ( - q: { remainingPercent: number | null; resetAt?: number } | undefined, - ): { remainingFraction?: number; resetTime?: string } | undefined => { + // Map a CommandAccountRow quota entry → cachedQuota pool shape + // so projectQuotaPoolForSidebar (the canonical projection seam) + // carries the per-window breakdown into the sidebar state. + const toPool = ( + q: + | { + remainingPercent: number | null + resetAt?: number + windows?: CommandAccountRow['quota'][number]['windows'] + } + | undefined, + ) => { if (!q || q.remainingPercent == null) return undefined return { remainingFraction: q.remainingPercent / 100, - // Preserve resetTime so the sidebar can render reset countdowns - // for each pool. Without this the TUI loses the freshest reset - // deadline every time the dialog re-renders. resetTime: typeof q.resetAt === 'number' && Number.isFinite(q.resetAt) ? new Date(q.resetAt).toISOString() : undefined, + windows: q.windows?.map((w) => ({ + window: w.window, + remainingFraction: w.remainingPercent / 100, + resetTime: + typeof w.resetAt === 'number' && Number.isFinite(w.resetAt) + ? new Date(w.resetAt).toISOString() + : '', + })), } } return { @@ -439,8 +479,8 @@ export function createCommandDataService( enabled: row.enabled, current: row.current, cachedQuota: { - gemini: toFraction(gemini), - 'non-gemini': toFraction(nonGemini), + gemini: toPool(gemini), + 'non-gemini': toPool(nonGemini), }, // The stamp check has already been done by `toCommandAccountRow` // before the rows reach this writer; nothing further for the diff --git a/packages/opencode/src/sidebar-state.test.ts b/packages/opencode/src/sidebar-state.test.ts index a9caa5e..a664ea2 100644 --- a/packages/opencode/src/sidebar-state.test.ts +++ b/packages/opencode/src/sidebar-state.test.ts @@ -24,6 +24,7 @@ import { join } from 'node:path' import { acquireFencedFileLock } from '@cortexkit/antigravity-auth-core' import { + buildSidebarMachineStateFromAccounts, DEFAULT_SIDEBAR_STATE, pruneActiveRouting, readSidebarState, @@ -553,3 +554,167 @@ describe('redaction', () => { expect(redacted.quota.gemini?.remainingPercent).toBe(30) }) }) + +describe('windows rework — producer seam tests', () => { + it('redactAccountForSidebar carries windows when cachedQuota has them', () => { + const redacted = redactAccountForSidebar({ + index: 0, + cachedQuota: { + gemini: { + remainingFraction: 0.92, + resetTime: '2026-07-28T18:24:21Z', + windows: [ + { + window: 'weekly', + remainingFraction: 0.92, + resetTime: '2026-07-28T18:24:21Z', + }, + { + window: '5h', + remainingFraction: 0.99, + resetTime: '2026-07-24T20:43:21Z', + }, + ], + }, + 'non-gemini': { + remainingFraction: 0.96, + resetTime: '2026-07-24T18:41:52Z', + windows: [ + { + window: 'weekly', + remainingFraction: 0.99, + resetTime: '2026-07-31T13:41:52Z', + }, + { + window: '5h', + remainingFraction: 0.96, + resetTime: '2026-07-24T18:41:52Z', + }, + ], + }, + }, + }) + + const gemini = redacted.quota.gemini + expect(gemini).toBeDefined() + expect(gemini!.remainingPercent).toBe(92) + expect(gemini!.windows).toHaveLength(2) + expect(gemini!.windows![0]!.window).toBe('weekly') + expect(gemini!.windows![0]!.remainingPercent).toBe(92) + expect(gemini!.windows![1]!.window).toBe('5h') + expect(gemini!.windows![1]!.remainingPercent).toBe(99) + + const nonGemini = redacted.quota['non-gemini'] + expect(nonGemini).toBeDefined() + expect(nonGemini!.windows).toHaveLength(2) + expect(nonGemini!.windows![0]!.remainingPercent).toBe(99) + expect(nonGemini!.windows![1]!.remainingPercent).toBe(96) + }) + + it('redactAccountForSidebar handles legacy cachedQuota without windows', () => { + const redacted = redactAccountForSidebar({ + index: 0, + cachedQuota: { + gemini: { remainingFraction: 0.5 }, + }, + }) + expect(redacted.quota.gemini?.remainingPercent).toBe(50) + expect(redacted.quota.gemini?.windows).toBeUndefined() + }) + + it('redactAccountForSidebar handles Free (weekly-only) windows', () => { + const redacted = redactAccountForSidebar({ + index: 0, + cachedQuota: { + gemini: { + remainingFraction: 0.89, + windows: [ + { + window: 'weekly', + remainingFraction: 0.89, + resetTime: '2026-07-31T15:54:18Z', + }, + ], + }, + }, + }) + + const gemini = redacted.quota.gemini + expect(gemini).toBeDefined() + expect(gemini!.windows).toHaveLength(1) + expect(gemini!.windows![0]!.window).toBe('weekly') + expect(gemini!.windows![0]!.remainingPercent).toBe(89) + }) + + it('full round-trip: write windows through setSidebarMachineState → readSidebarState', async () => { + const root = mkdtempSync(join(tmpdir(), 'agy-seam-windows-')) + const statePath = join(root, 'sidebar-state.json') + try { + // Write state with windowed data through the real machine-state path. + await setSidebarMachineState( + buildSidebarMachineStateFromAccounts( + [ + { + index: 0, + cachedQuota: { + gemini: { + remainingFraction: 0.7, + resetTime: '2026-08-01T00:00:00Z', + windows: [ + { + window: 'weekly', + remainingFraction: 0.7, + resetTime: '2026-08-01T00:00:00Z', + }, + { + window: '5h', + remainingFraction: 0.85, + resetTime: '2026-07-25T00:00:00Z', + }, + ], + }, + }, + }, + ], + { checkedAt: Date.now() }, + ), + { stateFile: statePath }, + ) + + // Read back and assert windows survived the full cycle. + const state = readSidebarState(statePath) + expect(state.accounts).toHaveLength(1) + const gemini = state.accounts[0]!.quota.gemini + expect(gemini).toBeDefined() + expect(gemini!.windows).toHaveLength(2) + expect(gemini!.windows![0]!.window).toBe('weekly') + expect(gemini!.windows![0]!.remainingPercent).toBe(70) + expect(gemini!.windows![1]!.window).toBe('5h') + expect(gemini!.windows![1]!.remainingPercent).toBe(85) + } finally { + rmSync(root, { recursive: true, force: true }) + } + }) + + it('identity mismatch drops the entire cachedQuota including windows', () => { + const redacted = redactAccountForSidebar({ + index: 0, + cachedQuota: { + gemini: { + remainingFraction: 0.9, + windows: [ + { + window: 'weekly', + remainingFraction: 0.9, + resetTime: '2026-08-01T00:00:00Z', + }, + ], + }, + }, + cachedQuotaAccountId: 'stamp-a', + currentQuotaAccountId: 'stamp-b', // mismatch + }) + // Stamp mismatch → cachedQuota dropped → no quota in the redacted output. + expect(Object.keys(redacted.quota)).toHaveLength(0) + }) +}) diff --git a/packages/opencode/src/sidebar-state.ts b/packages/opencode/src/sidebar-state.ts index edd7417..8eedca4 100644 --- a/packages/opencode/src/sidebar-state.ts +++ b/packages/opencode/src/sidebar-state.ts @@ -501,6 +501,59 @@ export interface SidebarAccountRedactionInput { currentQuotaAccountId?: string } +/** + * Project a raw cachedQuota pool entry into the sidebar-safe shape. + * + * Centralized seam — every producer (quota.ts pushSidebarQuotaSnapshot, + * auth-loader.ts materializer, command-data.ts writeSidebar) routes + * through here so the `windows` array is never dropped by an inline + * `{ remainingFraction, resetTime }` literal. + * + * Tolerant: returns `undefined` when the source is missing or the + * `remainingFraction` is not a finite number. Legacy entries without + * `windows` produce `{ remainingPercent, resetAt }` only — the TUI's + * legacy path then renders a single bar. + */ +export function projectQuotaPoolForSidebar(source: { + remainingFraction?: number + resetTime?: string + windows?: ReadonlyArray<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> +}): SidebarQuotaEntry | undefined { + const fraction = source.remainingFraction + if (typeof fraction !== 'number' || !Number.isFinite(fraction)) + return undefined + const remainingPercent = clampNumber(Math.round(fraction * 100), 0, 100) + let resetAt: number | undefined + if (typeof source.resetTime === 'string' && source.resetTime.length > 0) { + const parsed = Date.parse(source.resetTime) + if (Number.isFinite(parsed)) resetAt = parsed + } + + const windows: SidebarQuotaEntry['windows'] = source.windows?.length + ? source.windows.map((w) => ({ + window: w.window, + remainingPercent: clampNumber( + Math.round(w.remainingFraction * 100), + 0, + 100, + ), + resetAt: + typeof w.resetTime === 'string' && w.resetTime.length > 0 + ? (() => { + const parsed = Date.parse(w.resetTime) + return Number.isFinite(parsed) ? parsed : undefined + })() + : undefined, + })) + : undefined + + return { remainingPercent, resetAt, windows } +} + /** * Convert a live account snapshot into the redacted shape the TUI renders. * The redacted `SidebarAccountState` carries no email, refresh token, access @@ -539,30 +592,8 @@ export function redactAccountForSidebar( for (const key of ['gemini', 'non-gemini'] as const) { const entry = cached[key] if (!entry) continue - const fraction = entry.remainingFraction - if (typeof fraction !== 'number' || !Number.isFinite(fraction)) continue - const remainingPercent = clampNumber(Math.round(fraction * 100), 0, 100) - let resetAt: number | undefined - if (typeof entry.resetTime === 'string' && entry.resetTime.length > 0) { - const parsed = Date.parse(entry.resetTime) - if (Number.isFinite(parsed)) resetAt = parsed - } - const windows = entry.windows?.map((w) => ({ - window: w.window, - remainingPercent: clampNumber( - Math.round(w.remainingFraction * 100), - 0, - 100, - ), - resetAt: - typeof w.resetTime === 'string' && w.resetTime.length > 0 - ? (() => { - const parsed = Date.parse(w.resetTime) - return Number.isFinite(parsed) ? parsed : undefined - })() - : undefined, - })) - quota[key] = { remainingPercent, resetAt, windows } + const projected = projectQuotaPoolForSidebar(entry) + if (projected) quota[key] = projected } } diff --git a/packages/opencode/src/tui-compiled/plugin/command-data.ts b/packages/opencode/src/tui-compiled/plugin/command-data.ts index 210310e..bf80ef0 100644 --- a/packages/opencode/src/tui-compiled/plugin/command-data.ts +++ b/packages/opencode/src/tui-compiled/plugin/command-data.ts @@ -82,6 +82,11 @@ type CommandDataQuotaGroupSummary = { remainingFraction?: number resetTime?: string modelCount: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> } type CommandDataAccountQuotaResult = { @@ -137,6 +142,11 @@ export interface CommandAccountRow { label: string remainingPercent: number | null resetAt?: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingPercent: number + resetAt?: number + }> }> } @@ -198,11 +208,27 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { const parsed = Date.parse(cachedEntry.resetTime) if (Number.isFinite(parsed)) resetAt = parsed } + // Carry per-window breakdown when the cache was populated from the + // windowed RUQS path — the sidebar writer projects through here. + const windows = cachedEntry.windows?.length + ? cachedEntry.windows.map((w) => ({ + window: w.window, + remainingPercent: Math.round(w.remainingFraction * 100), + resetAt: + typeof w.resetTime === 'string' && w.resetTime.length > 0 + ? (() => { + const parsed = Date.parse(w.resetTime) + return Number.isFinite(parsed) ? parsed : undefined + })() + : undefined, + })) + : undefined quota.push({ key, label: QUOTA_GROUP_LABELS[key], remainingPercent, resetAt, + windows, }) } const label = `Account ${entry.index + 1}` @@ -418,19 +444,33 @@ export function createCommandDataService( const accounts: SidebarAccountRedactionInput[] = rows.map((row) => { const gemini = row.quota.find((q) => q.key === 'gemini') const nonGemini = row.quota.find((q) => q.key === 'non-gemini') - const toFraction = ( - q: { remainingPercent: number | null; resetAt?: number } | undefined, - ): { remainingFraction?: number; resetTime?: string } | undefined => { + // Map a CommandAccountRow quota entry → cachedQuota pool shape + // so projectQuotaPoolForSidebar (the canonical projection seam) + // carries the per-window breakdown into the sidebar state. + const toPool = ( + q: + | { + remainingPercent: number | null + resetAt?: number + windows?: CommandAccountRow['quota'][number]['windows'] + } + | undefined, + ) => { if (!q || q.remainingPercent == null) return undefined return { remainingFraction: q.remainingPercent / 100, - // Preserve resetTime so the sidebar can render reset countdowns - // for each pool. Without this the TUI loses the freshest reset - // deadline every time the dialog re-renders. resetTime: typeof q.resetAt === 'number' && Number.isFinite(q.resetAt) ? new Date(q.resetAt).toISOString() : undefined, + windows: q.windows?.map((w) => ({ + window: w.window, + remainingFraction: w.remainingPercent / 100, + resetTime: + typeof w.resetAt === 'number' && Number.isFinite(w.resetAt) + ? new Date(w.resetAt).toISOString() + : '', + })), } } return { @@ -439,8 +479,8 @@ export function createCommandDataService( enabled: row.enabled, current: row.current, cachedQuota: { - gemini: toFraction(gemini), - 'non-gemini': toFraction(nonGemini), + gemini: toPool(gemini), + 'non-gemini': toPool(nonGemini), }, // The stamp check has already been done by `toCommandAccountRow` // before the rows reach this writer; nothing further for the diff --git a/packages/opencode/src/tui-compiled/sidebar-state.ts b/packages/opencode/src/tui-compiled/sidebar-state.ts index edd7417..8eedca4 100644 --- a/packages/opencode/src/tui-compiled/sidebar-state.ts +++ b/packages/opencode/src/tui-compiled/sidebar-state.ts @@ -501,6 +501,59 @@ export interface SidebarAccountRedactionInput { currentQuotaAccountId?: string } +/** + * Project a raw cachedQuota pool entry into the sidebar-safe shape. + * + * Centralized seam — every producer (quota.ts pushSidebarQuotaSnapshot, + * auth-loader.ts materializer, command-data.ts writeSidebar) routes + * through here so the `windows` array is never dropped by an inline + * `{ remainingFraction, resetTime }` literal. + * + * Tolerant: returns `undefined` when the source is missing or the + * `remainingFraction` is not a finite number. Legacy entries without + * `windows` produce `{ remainingPercent, resetAt }` only — the TUI's + * legacy path then renders a single bar. + */ +export function projectQuotaPoolForSidebar(source: { + remainingFraction?: number + resetTime?: string + windows?: ReadonlyArray<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> +}): SidebarQuotaEntry | undefined { + const fraction = source.remainingFraction + if (typeof fraction !== 'number' || !Number.isFinite(fraction)) + return undefined + const remainingPercent = clampNumber(Math.round(fraction * 100), 0, 100) + let resetAt: number | undefined + if (typeof source.resetTime === 'string' && source.resetTime.length > 0) { + const parsed = Date.parse(source.resetTime) + if (Number.isFinite(parsed)) resetAt = parsed + } + + const windows: SidebarQuotaEntry['windows'] = source.windows?.length + ? source.windows.map((w) => ({ + window: w.window, + remainingPercent: clampNumber( + Math.round(w.remainingFraction * 100), + 0, + 100, + ), + resetAt: + typeof w.resetTime === 'string' && w.resetTime.length > 0 + ? (() => { + const parsed = Date.parse(w.resetTime) + return Number.isFinite(parsed) ? parsed : undefined + })() + : undefined, + })) + : undefined + + return { remainingPercent, resetAt, windows } +} + /** * Convert a live account snapshot into the redacted shape the TUI renders. * The redacted `SidebarAccountState` carries no email, refresh token, access @@ -539,30 +592,8 @@ export function redactAccountForSidebar( for (const key of ['gemini', 'non-gemini'] as const) { const entry = cached[key] if (!entry) continue - const fraction = entry.remainingFraction - if (typeof fraction !== 'number' || !Number.isFinite(fraction)) continue - const remainingPercent = clampNumber(Math.round(fraction * 100), 0, 100) - let resetAt: number | undefined - if (typeof entry.resetTime === 'string' && entry.resetTime.length > 0) { - const parsed = Date.parse(entry.resetTime) - if (Number.isFinite(parsed)) resetAt = parsed - } - const windows = entry.windows?.map((w) => ({ - window: w.window, - remainingPercent: clampNumber( - Math.round(w.remainingFraction * 100), - 0, - 100, - ), - resetAt: - typeof w.resetTime === 'string' && w.resetTime.length > 0 - ? (() => { - const parsed = Date.parse(w.resetTime) - return Number.isFinite(parsed) ? parsed : undefined - })() - : undefined, - })) - quota[key] = { remainingPercent, resetAt, windows } + const projected = projectQuotaPoolForSidebar(entry) + if (projected) quota[key] = projected } } From fe5d93645e25bb839d46c648380723272a6ba0e4 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 21:44:38 +0200 Subject: [PATCH 15/25] fix(quota): source managedProjectId from the account record MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bare refresh tokens have no packed project IDs — the real managedProjectId lives on the persisted account record, not in the refresh string. Fall back to account.managedProjectId (mirroring the existing pattern at quota.ts:511). Add regression test: fetchQuotaSummary posts managedProjectId from options. e2e mock: quotaSummaryWindow fixture enforces 403 when the posted project does not match the managedProjectId — catches exactly this class of bug. rpc-tui-flow account now has distinct projectId vs managedProjectId. --- packages/core/src/quota-manager.test.ts | 25 +++++++++++++++ .../e2e-tests/src/mock-antigravity-server.ts | 32 ++++++++++++++++++- .../e2e-tests/src/rpc-tui-flow.e2e.test.ts | 5 +++ packages/opencode/src/plugin/quota.ts | 6 +++- 4 files changed, 66 insertions(+), 2 deletions(-) diff --git a/packages/core/src/quota-manager.test.ts b/packages/core/src/quota-manager.test.ts index 77dcb33..34090f4 100644 --- a/packages/core/src/quota-manager.test.ts +++ b/packages/core/src/quota-manager.test.ts @@ -1020,4 +1020,29 @@ describe('fetchQuotaSummary', () => { }) expect(receivedProjectId).toBe('regular-proj') }) + + it('uses managedProjectId from options when present (the record fallback path)', async () => { + // Regression: bare refresh tokens have no packed managedProjectId. + // The caller must supply it from the account record. This test + // asserts that fetchQuotaSummary posts the options.managedProjectId + // when it is provided. + let receivedProjectId = '' + const summary: RetrieveUserQuotaSummaryResponse = { groups: [] } + const fetchVia = async ( + _url: string, + init: RequestInit, + ): Promise => { + const body = JSON.parse((init as any).body ?? '{}') + receivedProjectId = body.project as string + return new Response(JSON.stringify(summary), { status: 200 }) + } + await fetchQuotaSummary({ + accessToken: 'tok', + managedProjectId: 'managed-from-record', + projectId: 'regular-proj', + endpoints: ENDPOINTS, + fetchVia: fetchVia as any, + }) + expect(receivedProjectId).toBe('managed-from-record') + }) }) diff --git a/packages/e2e-tests/src/mock-antigravity-server.ts b/packages/e2e-tests/src/mock-antigravity-server.ts index ec3d0b7..d987448 100644 --- a/packages/e2e-tests/src/mock-antigravity-server.ts +++ b/packages/e2e-tests/src/mock-antigravity-server.ts @@ -90,6 +90,13 @@ export interface GeminiCliQuotaFixture extends BaseFixture { /** Windowed quota summary envelope — fed to `retrieveUserQuotaSummary`. */ export interface QuotaSummaryWindowFixture extends BaseFixture { kind: 'quotaSummaryWindow' + /** + * When set, the server returns 403 if the posted `project` does not + * equal this value — mirroring the real API's PERMISSION_DENIED for + * non-managed project IDs. This makes the e2e fail on exactly the + * class of bug where the caller sends the wrong project ID. + */ + managedProjectId?: string /** Groups with their per-window buckets. */ groups: Array<{ displayName: string @@ -259,7 +266,7 @@ export async function startMockAntigravityServer(): Promise { } openSockets.add(response) response.once('close', () => openSockets.delete(response)) - await dispatch(fixture, request, response) + await dispatch(fixture, request, response, body) } catch (error) { try { sendJson(response, 500, { @@ -299,6 +306,7 @@ export async function startMockAntigravityServer(): Promise { fixture: Fixture, _request: IncomingMessage, response: ServerResponse, + reqBody: string, ): Promise { switch (fixture.kind) { case 'json': { @@ -331,6 +339,28 @@ export async function startMockAntigravityServer(): Promise { return } case 'quotaSummaryWindow': { + // Honor the real API's 403 behavior: if managedProjectId is set + // on the fixture and the posted project does not match, return + // PERMISSION_DENIED so the test exercises the fallback path. + if (fixture.managedProjectId) { + let postedProject = '' + try { + postedProject = + (JSON.parse(reqBody) as { project?: string }).project ?? '' + } catch { + /* malformed — let the 200 path handle it */ + } + if (postedProject && postedProject !== fixture.managedProjectId) { + sendJson(response, 403, { + error: { + code: 7, + message: 'PERMISSION_DENIED', + status: 'PERMISSION_DENIED', + }, + }) + return + } + } sendJson(response, fixture.status ?? 200, { groups: fixture.groups, description: diff --git a/packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts b/packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts index a364e07..a00d0a3 100644 --- a/packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts +++ b/packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts @@ -122,8 +122,13 @@ describe('rpc / tui flow (e2e)', () => { projectId: 'project-rpc', }) // Primary: retrieveUserQuotaSummary (windowed). + // managedProjectId enforces the real API's 403: the caller must + // post the managed project id, not the regular project id. + // If it posts the wrong id, the mock returns 403 and the test + // exercises the fallback path instead — failing the window assertion. h.server.enqueue({ kind: 'quotaSummaryWindow', + managedProjectId: 'managed-rpc', groups: [ { displayName: 'Gemini Models', diff --git a/packages/opencode/src/plugin/quota.ts b/packages/opencode/src/plugin/quota.ts index 355cabd..f98489e 100644 --- a/packages/opencode/src/plugin/quota.ts +++ b/packages/opencode/src/plugin/quota.ts @@ -377,7 +377,11 @@ function makeFetchAccountQuota( let fellBackToLegacy = false const authParts = parseRefreshParts(auth.refresh) - const managedProjectId = authParts.managedProjectId + // Bare refresh tokens have no packed project IDs — fall back to the + // account record. The real managedProjectId lives on the persisted + // account, not in the packed refresh string. + const managedProjectId = + authParts.managedProjectId ?? account.managedProjectId // Primary: retrieveUserQuotaSummary (windowed quota). // Fall back to fetchAvailableModels on any failure (403, network, etc.) From 1cbcfd66acde4ace05514cfeeaa5bc05e003e4e4 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 21:56:20 +0200 Subject: [PATCH 16/25] fix(core): project record-level managedProjectId through account metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add standalone projectId/managedProjectId fields to ManagedAccount, populated from the authoritative stored account record during construction. These survive bare-refresh-token rotations where parts.* may be lost. getAccountsForQuotaCheck + buildStorageSnapshot now fall back to the record-level fields via a.parts.* ?? a.*. updateFromAuth keeps both parts.* and record-level fields in sync. Regression tests: getAccountsForQuotaCheck returns managedProjectId from bare-token accounts; save→reload round-trip preserves it. --- packages/core/src/account-manager.test.ts | 62 +++++++++++++++++++++++ packages/core/src/account-manager.ts | 23 +++++++-- 2 files changed, 81 insertions(+), 4 deletions(-) diff --git a/packages/core/src/account-manager.test.ts b/packages/core/src/account-manager.test.ts index 6adf5e4..55b060d 100644 --- a/packages/core/src/account-manager.test.ts +++ b/packages/core/src/account-manager.test.ts @@ -294,3 +294,65 @@ describe('AccountManager instance dependencies', () => { expect(second.getAccounts()[0]?.coolingDownUntil).toBe(9_500) }) }) + +describe('managedProjectId projection', () => { + it('getAccountsForQuotaCheck returns managedProjectId from the record when refresh token is bare', () => { + const stored: AccountStorageV4 = { + version: 4, + accounts: [ + { + email: 'test@example.com', + refreshToken: 'bare-refresh-token', + projectId: 'my-project', + managedProjectId: 'my-managed-project', + addedAt: 1_000, + lastUsed: 2_000, + }, + ], + activeIndex: 0, + } + const manager = new AccountManager(undefined, stored, { + store: createStore(stored).store, + now: () => 1_000, + }) + const accounts = manager.getAccountsForQuotaCheck() + expect(accounts).toHaveLength(1) + expect(accounts[0]!.projectId).toBe('my-project') + expect(accounts[0]!.managedProjectId).toBe('my-managed-project') + }) + + it('save→reload round-trip preserves managedProjectId from the record', async () => { + const { store, state } = createStore(null) + const stored: AccountStorageV4 = { + version: 4, + accounts: [ + { + email: 'test@example.com', + refreshToken: 'bare-refresh-token', + projectId: 'my-project', + managedProjectId: 'my-managed-project', + addedAt: 1_000, + lastUsed: 2_000, + }, + ], + activeIndex: 0, + } + const manager = new AccountManager(undefined, stored, { + store, + now: () => 1_000, + }) + // Trigger a save — buildStorageSnapshot must preserve managedProjectId. + manager.requestSaveToDisk() + await manager.flushSaveToDisk() + const saved = state() + expect(saved?.accounts[0]?.managedProjectId).toBe('my-managed-project') + + // Reload and verify getAccountsForQuotaCheck still returns it. + const manager2 = new AccountManager(undefined, saved, { + store, + now: () => 1_000, + }) + const accounts = manager2.getAccountsForQuotaCheck() + expect(accounts[0]!.managedProjectId).toBe('my-managed-project') + }) +}) diff --git a/packages/core/src/account-manager.ts b/packages/core/src/account-manager.ts index 07e4c78..04c5ba1 100644 --- a/packages/core/src/account-manager.ts +++ b/packages/core/src/account-manager.ts @@ -73,6 +73,13 @@ export interface ManagedAccount { addedAt: number lastUsed: number parts: RefreshParts + /** Authoritative project ID from the persisted account record. Survives + * bare-refresh-token rotations where `parts.projectId` may be lost. */ + projectId?: string + /** Authoritative managed project ID from the persisted account record. + * Survives bare-refresh-token rotations where `parts.managedProjectId` + * may be lost. */ + managedProjectId?: string access?: string expires?: number enabled: boolean @@ -394,6 +401,10 @@ export class AccountManager { projectId: acc.projectId, managedProjectId: acc.managedProjectId, }, + // Authoritative record-level fields that survive bare-refresh-token + // rotations where `parts.*` may be overwritten with undefined. + projectId: acc.projectId, + managedProjectId: acc.managedProjectId, access: matchesFallback ? authFallback?.access : undefined, expires: matchesFallback ? authFallback?.expires : undefined, enabled: acc.enabled !== false, @@ -1548,6 +1559,10 @@ export class AccountManager { managedProjectId: parts.managedProjectId ?? account.parts.managedProjectId, } + // Keep the record-level fields in sync with the authoritative source. + account.projectId = parts.projectId ?? account.projectId + account.managedProjectId = + parts.managedProjectId ?? account.managedProjectId account.access = auth.access account.expires = auth.expires } @@ -1628,8 +1643,8 @@ export class AccountManager { email: a.email, label: a.label, refreshToken: a.parts.refreshToken, - projectId: a.parts.projectId, - managedProjectId: a.parts.managedProjectId, + projectId: a.parts.projectId ?? a.projectId, + managedProjectId: a.parts.managedProjectId ?? a.managedProjectId, addedAt: a.addedAt, lastUsed: a.lastUsed, enabled: a.enabled, @@ -2045,8 +2060,8 @@ export class AccountManager { return this.accounts.map((a) => ({ email: a.email, refreshToken: a.parts.refreshToken, - projectId: a.parts.projectId, - managedProjectId: a.parts.managedProjectId, + projectId: a.parts.projectId ?? a.projectId, + managedProjectId: a.parts.managedProjectId ?? a.managedProjectId, addedAt: a.addedAt, lastUsed: a.lastUsed, enabled: a.enabled, From 8bff3ec826a67a95ce9c273fc80e93e6aaaff7ca Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 22:55:57 +0200 Subject: [PATCH 17/25] fix(auth): delegate loader fetch through the live runtime P1 #1: returned fetch captured the current runtime's interceptor by reference, so the host kept routing requests through the OLD AccountManager after an OAuth-driven reload. Replace with a call-time delegating wrapper that reads the live runtime on every call. P1 #4: AuthLoaderHandle.load was advertised on the type but never attached to the callable. Add load so the contract matches the documented public surface. P1 #5: installRuntime's prior-runtime dispose was fire-and-forget (void), so two overlapping reloads could interleave teardown. Await the dispose and serialize reloads through a shared promise chain so each install waits for the previous one to settle. --- .../opencode/src/plugin/auth-loader.test.ts | 212 ++++++++++++++++++ packages/opencode/src/plugin/auth-loader.ts | 39 +++- 2 files changed, 244 insertions(+), 7 deletions(-) diff --git a/packages/opencode/src/plugin/auth-loader.test.ts b/packages/opencode/src/plugin/auth-loader.test.ts index f88ca63..fdd3bc1 100644 --- a/packages/opencode/src/plugin/auth-loader.test.ts +++ b/packages/opencode/src/plugin/auth-loader.test.ts @@ -243,4 +243,216 @@ describe('createAuthLoader', () => { expect(secondManager.dispose).toHaveBeenCalledTimes(1) expect(secondDispose).toHaveBeenCalledTimes(1) }) + + it('returned fetch delegates to the live runtime after reload()', async () => { + const firstFetch = mock(async () => new Response('first')) + const secondFetch = mock(async () => new Response('second')) + const firstDispose = mock(async () => {}) + const secondDispose = mock(async () => {}) + const managerA = { + name: 'a', + getAccountCount: () => 1, + getAccounts: () => [ + { + index: 0, + email: 'a@example.test', + enabled: true, + parts: { refreshToken: 'a-refresh' }, + cachedQuota: undefined, + }, + ], + requestSaveToDisk: mock(() => {}), + dispose: mock(async () => {}), + } + const managerB = { + name: 'b', + getAccountCount: () => 1, + getAccounts: () => [ + { + index: 0, + email: 'b@example.test', + enabled: true, + parts: { refreshToken: 'b-refresh' }, + cachedQuota: undefined, + }, + ], + requestSaveToDisk: mock(() => {}), + dispose: mock(async () => {}), + } + const managers = [managerA, managerB] + const createFetch = mock(() => { + const isFirst = createFetch.mock.calls.length === 1 + return { + fetch: isFirst ? firstFetch : secondFetch, + dispose: isFirst ? firstDispose : secondDispose, + } + }) + const lifecycle = createPluginLifecycle({ + sessionRegistry: { clear: mock(() => {}) }, + shutdownDiskSignatureCache: mock(async () => {}), + clearFetchState: mock(() => {}), + }) + const loader = createAuthLoader({ + client: { + auth: { set: mock(async () => {}) }, + tui: { showToast: mock(async () => {}) }, + } as never, + providerId: 'google', + config: { ...DEFAULT_CONFIG, proactive_token_refresh: false }, + lifecycle, + createFetch, + dependencies: { + loadAccounts: mock(async () => storedAccounts()), + clearAccounts: mock(async () => {}), + loadAccountManager: mock(async () => managers.shift() as never), + }, + }) + const getAuth = mock(async () => ({ + type: 'oauth' as const, + refresh: 'stored-refresh|stored-project|managed-project', + access: 'access', + expires: 100, + })) as unknown as GetAuth + + const firstResult = (await loader(getAuth, { + id: 'g', + name: 'G', + source: 'custom', + env: [], + options: {}, + models: {}, + } as never)) as { + fetch: (input: RequestInfo, init?: RequestInit) => Promise + } + const liveFetch = firstResult.fetch + + await loader.reload(getAuth) + + const midResponse = await liveFetch('https://example.test/mid') + expect(midResponse).toBeInstanceOf(Response) + expect(firstFetch).not.toHaveBeenCalled() + expect(secondFetch).toHaveBeenCalledTimes(1) + }) + + it('exposes load on the handle so contract consumers can read it', () => { + const createFetch = mock(() => ({ + fetch: mock(async () => new Response('ok')), + dispose: mock(async () => {}), + })) + const lifecycle = createPluginLifecycle({ + sessionRegistry: { clear: mock(() => {}) }, + shutdownDiskSignatureCache: mock(async () => {}), + clearFetchState: mock(() => {}), + }) + const loader = createAuthLoader({ + client: { + auth: { set: mock(async () => {}) }, + tui: { showToast: mock(async () => {}) }, + } as never, + providerId: 'google', + config: DEFAULT_CONFIG, + lifecycle, + createFetch, + dependencies: { + loadAccounts: mock(async () => null), + clearAccounts: mock(async () => {}), + }, + }) + expect(typeof loader.load).toBe('function') + expect(loader.load).toBe(loader) + }) + + it('awaits prior runtime dispose before returning from reload()', async () => { + let disposeFinished = false + const firstDispose = mock(async () => { + await new Promise((resolve) => setTimeout(resolve, 10)) + disposeFinished = true + }) + const firstManager = { + name: 'first', + getAccountCount: () => 1, + getAccounts: () => [ + { + index: 0, + email: 'first@example.test', + enabled: true, + parts: { refreshToken: 'first-refresh' }, + cachedQuota: undefined, + }, + ], + requestSaveToDisk: mock(() => {}), + dispose: mock(async () => {}), + } + const secondManager = { + name: 'second', + getAccountCount: () => 1, + getAccounts: () => [ + { + index: 0, + email: 'second@example.test', + enabled: true, + parts: { refreshToken: 'second-refresh' }, + cachedQuota: undefined, + }, + ], + requestSaveToDisk: mock(() => {}), + dispose: mock(async () => {}), + } + const managers = [firstManager, secondManager] + const secondDispose = mock(async () => {}) + const fetchRuntimes = [ + { + fetch: mock(async () => new Response('first')), + dispose: firstDispose, + }, + { + fetch: mock(async () => new Response('second')), + dispose: secondDispose, + }, + ] + const lifecycle = createPluginLifecycle({ + sessionRegistry: { clear: mock(() => {}) }, + shutdownDiskSignatureCache: mock(async () => {}), + clearFetchState: mock(() => {}), + }) + const createFetch = mock(() => fetchRuntimes.shift()!) + const loader = createAuthLoader({ + client: { + auth: { set: mock(async () => {}) }, + tui: { showToast: mock(async () => {}) }, + } as never, + providerId: 'google', + config: { ...DEFAULT_CONFIG, proactive_token_refresh: false }, + lifecycle, + createFetch, + dependencies: { + loadAccounts: mock(async () => storedAccounts()), + clearAccounts: mock(async () => {}), + loadAccountManager: mock(async () => managers.shift() as never), + }, + }) + const getAuth = mock(async () => ({ + type: 'oauth' as const, + refresh: 'stored-refresh|stored-project|managed-project', + access: 'access', + expires: 100, + })) as unknown as GetAuth + + await loader(getAuth, { + id: 'g', + name: 'G', + source: 'custom', + env: [], + options: {}, + models: {}, + } as never) + + const before = disposeFinished + const reloadPromise = loader.reload(getAuth) + // Without awaiting dispose, the reload promise returns before + // the previous runtime's dispose has settled. + expect(disposeFinished).toBe(before) + await reloadPromise + expect(disposeFinished).toBe(true) + }) }) diff --git a/packages/opencode/src/plugin/auth-loader.ts b/packages/opencode/src/plugin/auth-loader.ts index aba8a8c..12d0f08 100644 --- a/packages/opencode/src/plugin/auth-loader.ts +++ b/packages/opencode/src/plugin/auth-loader.ts @@ -117,6 +117,13 @@ export function createAuthLoader({ getLogFilePath: dependencies?.getLogFilePath ?? getLogFilePath, } let fetchRuntime: AuthFetchRuntime | null = null + // Reload chain — `installRuntime` swaps the fetch runtime and + // (previously) discarded the previous runtime's dispose with `void`, + // so two overlapping reloads could interleave teardown. Serialize + // reloads through a single shared promise so each new install waits + // for the previous one to fully settle (including its dispose) + // before tearing down the runtime it depends on. + let reloadChain: Promise = Promise.resolve() lifecycle.register( { @@ -159,9 +166,14 @@ export function createAuthLoader({ await lifecycle.replaceAccountRuntime(accountManager, refreshQueue) refreshQueue?.start() + // Swap the fetch runtime FIRST so the host's captured fetch + // reference (a call-time delegating wrapper below) immediately + // routes through the new interceptor, then await the previous + // runtime's dispose. The swap-then-dispose order guarantees there + // is no fetch gap between the old and new runtimes. const previousRuntime = fetchRuntime fetchRuntime = createFetch({ accountManager, getAuth }) - void previousRuntime?.dispose() + await previousRuntime?.dispose() // Push the freshly materialized account pool into the sidebar so the // TUI's next poll renders the labels / health / cooldown it needs @@ -277,10 +289,14 @@ export function createAuthLoader({ return { apiKey: '', - // `fetchRuntime` is guaranteed non-null by `installRuntime`'s - // `createFetch({ ... })` call above — the assignment replaced - // any prior value. - fetch: fetchRuntime!.fetch, + // Return a stable delegating wrapper that reads `fetchRuntime` + // at CALL time. A direct `fetchRuntime!.fetch` reference would + // capture the current runtime; the host keeps that reference + // across `reload()` calls, so a captured fetch would still route + // through the OLD interceptor after an OAuth add. The wrapper + // matches the host's `LoaderResult.fetch` signature exactly + // (no `this` binding) so behavior is preserved. + fetch: (input, init) => fetchRuntime!.fetch(input, init), } } @@ -303,8 +319,17 @@ export function createAuthLoader({ return runLoader(getAuth, provider) } async function reload(getAuth: GetAuth): Promise { - return reloadRuntime(getAuth) + const next = reloadChain.then(() => reloadRuntime(getAuth)) + reloadChain = next.catch(() => {}) + return next } - const authLoader = Object.assign(authLoaderCallable, { reload }) + // `load` mirrors the authLoaderCallable so the public handle + // (`AuthLoaderHandle.load`) exposes the same entry point the host + // invokes. Without this assignment, contract consumers read + // `.load` as undefined and the documented type would lie. + const authLoader = Object.assign(authLoaderCallable, { + reload, + load: authLoaderCallable, + }) return authLoader as LoadAndInstallRuntime & AuthLoaderHandle } From 37e49528ded06fd6f2eab1a919d65b8d96b6ed2c Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 22:56:02 +0200 Subject: [PATCH 18/25] fix(accounts): preserve per-family active indexes and cooldowns on removal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P1 #2: removeAccount() persisted a hardcoded 0 activeIndex when the captured current token was the one removed. The live AccountManager keeps the cursor at the removed slot, so use the clamped removed index instead of falling back to 0 — restart-time reload now re-elects the same account the live manager kept current. P2 #3: removeAccount() collapsed both per-family active cursors onto the same captured token, silently unelecting whichever family the captured token didn't belong to. Expose getActiveIndexByFamily on the AccountManager + view, capture per-family tokens in the mutator, and persist each family's remapped index independently. P2 #6: the sidebar refresher built dialog rows from the post-mutation command data, which does not carry the running cooldown. A toggle / remove / setCurrent action momentarily cleared any rate-limited account's cooldown. Look up the live account's coolingDownUntil by index and preserve it on the projection. --- packages/core/src/account-manager.ts | 15 +++ .../opencode/src/plugin/command-data.test.ts | 109 ++++++++++++++++++ packages/opencode/src/plugin/command-data.ts | 40 ++++--- packages/opencode/src/plugin/commands.test.ts | 59 ++++++++++ packages/opencode/src/plugin/commands.ts | 68 +++++++---- packages/opencode/src/plugin/index.ts | 5 + 6 files changed, 261 insertions(+), 35 deletions(-) diff --git a/packages/core/src/account-manager.ts b/packages/core/src/account-manager.ts index 04c5ba1..b637289 100644 --- a/packages/core/src/account-manager.ts +++ b/packages/core/src/account-manager.ts @@ -697,6 +697,21 @@ export class AccountManager { return null } + /** + * Numeric active indexes for each model family. Exposed so callers + * that persist `activeIndexByFamily` (e.g. command-data's remove + * path) can capture the live cursor per family without going + * through the account-lookup layer. + */ + getActiveIndexByFamily( + identity?: AccountSessionIdentity, + ): Record { + return { + claude: this.getActiveIndex('claude', identity), + gemini: this.getActiveIndex('gemini', identity), + } + } + markSwitched( account: ManagedAccount, reason: 'rate-limit' | 'initial' | 'rotation', diff --git a/packages/opencode/src/plugin/command-data.test.ts b/packages/opencode/src/plugin/command-data.test.ts index 74f43fb..0a75d8a 100644 --- a/packages/opencode/src/plugin/command-data.test.ts +++ b/packages/opencode/src/plugin/command-data.test.ts @@ -109,6 +109,7 @@ function quotaAccountIdentity(refreshToken: string): string { function makeHarness(options: { accounts: AccountFixture[] activeIndex?: number + geminiActiveIndex?: number refreshResults?: Map now?: () => number /** When true, also wire a storage adapter (defaults to true). */ @@ -207,6 +208,7 @@ function makeHarness(options: { const saveCalls = { count: 0 } const activeIndex = options.activeIndex ?? 0 let liveCurrentIndex = activeIndex + let liveGeminiIndex = options.geminiActiveIndex ?? activeIndex const accountManagerView: CommandDataAccountManagerView = { getAccounts() { @@ -262,6 +264,9 @@ function makeHarness(options: { activeIndex() { return liveCurrentIndex }, + getActiveIndexByFamily(): { claude: number; gemini: number } { + return { claude: liveCurrentIndex, gemini: liveGeminiIndex } + }, setAccountEnabled(index: number, enabled: boolean): boolean { const account = liveView[index] if (!account) return false @@ -273,14 +278,20 @@ function makeHarness(options: { setAccountCurrent(index: number): boolean { if (index < 0 || index >= liveView.length) return false liveCurrentIndex = index + liveGeminiIndex = index return true }, removeAccountByIndex(index: number): boolean { if (index < 0 || index >= liveView.length) return false liveView.splice(index, 1) + if (liveCurrentIndex > index) liveCurrentIndex -= 1 if (liveCurrentIndex >= liveView.length) { liveCurrentIndex = Math.max(0, liveView.length - 1) } + if (liveGeminiIndex > index) liveGeminiIndex -= 1 + if (liveGeminiIndex >= liveView.length) { + liveGeminiIndex = Math.max(0, liveView.length - 1) + } return true }, getRefreshTokenAt(index: number): string | undefined { @@ -507,6 +518,41 @@ describe('createCommandDataService', () => { }) }) + it('silently skips unknown quota keys in a legacy snapshot (tolerant read)', async () => { + // Older Antigravity revisions persisted quota under keys we no + // longer render (`claude`, `gpt-4`, ad-hoc pool names). The + // projection must IGNORE unknown keys rather than crash the + // dialog or surface them as keys without a label. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-a', + label: 'Account A', + cachedQuota: { + // Legacy / unknown keys that the renderer doesn't support. + claude: { remainingFraction: 0.9 }, + 'gpt-4': { remainingFraction: 0.8 }, + // Supported key — must still render. + gemini: { remainingFraction: 0.5 }, + } as unknown as QuotaGroupFixture, + }), + ], + }) + + const rows = await harness.service.listAccounts() + + expect(rows).toHaveLength(1) + expect(rows[0]?.quota).toEqual([ + { + key: 'gemini', + label: 'Gemini', + remainingPercent: 50, + resetAt: undefined, + windows: undefined, + }, + ]) + }) + it('refreshQuota() fetches every account, including an uncached new account, through the shared quota manager', async () => { const baseAccount = ( refreshToken: string, @@ -957,6 +1003,69 @@ describe('createCommandDataService', () => { }) }) + it("removeAccount() persists the live manager's post-removal current when removing the current middle account", async () => { + // Current is refresh-b (index 1). Removing index 1 (refresh-b) + // leaves the live manager's current index pointing at index 1, + // which now holds refresh-c (the account that shifted in). The + // previous implementation persisted 0 because the captured token + // was the one removed and the lookup fell back to "unknown token + // → 0", which would re-elect refresh-a on the next restart even + // though the live manager kept refresh-c current. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha' }), + makeAccountFixture({ refreshToken: 'refresh-b', label: 'Beta' }), + makeAccountFixture({ refreshToken: 'refresh-c', label: 'Gamma' }), + ], + activeIndex: 1, + }) + + await harness.service.removeAccount(1) + + expect(harness.storage.accounts.map((a) => a.refreshToken)).toEqual([ + 'refresh-a', + 'refresh-c', + ]) + // Live manager's current still points at index 1 (now refresh-c). + // Persisted activeIndex must follow the live manager, not zero. + expect(harness.storage.activeIndex).toBe(1) + expect(harness.storage.activeIndexByFamily).toEqual({ + claude: 1, + gemini: 1, + }) + }) + + it('removeAccount() persists per-family active indexes independently when removing a non-current account', async () => { + // Claude is current on refresh-a (index 0); Gemini is current on + // refresh-c (index 2). Removing refresh-b (index 1) — a + // non-current for both families — must keep each family's + // current pinned to its own account: claude still on refresh-a, + // gemini still on refresh-c (which shifted from 2 to 1). + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha' }), + makeAccountFixture({ refreshToken: 'refresh-b', label: 'Beta' }), + makeAccountFixture({ refreshToken: 'refresh-c', label: 'Gamma' }), + ], + activeIndex: 0, + geminiActiveIndex: 2, + }) + + await harness.service.removeAccount(1) + + expect(harness.storage.accounts.map((a) => a.refreshToken)).toEqual([ + 'refresh-a', + 'refresh-c', + ]) + // Both families should track the same numeric slot they had + // before, but for refresh-c, the shift from index 2 → 1 must + // move the persisted gemini index from 2 → 1. + expect(harness.storage.activeIndexByFamily).toEqual({ + claude: 0, + gemini: 1, + }) + }) + it('removeAccount() returns null when the index is out of range', async () => { const harness = makeHarness({ accounts: [ diff --git a/packages/opencode/src/plugin/command-data.ts b/packages/opencode/src/plugin/command-data.ts index bf80ef0..5170e82 100644 --- a/packages/opencode/src/plugin/command-data.ts +++ b/packages/opencode/src/plugin/command-data.ts @@ -308,6 +308,13 @@ export interface CommandDataAccountManagerView { requestSaveToDisk(): void flushSaveToDisk(): Promise activeIndex(): number + /** + * Per-family live cursor. Used by the remove path so the persisted + * `activeIndexByFamily` follows each model's currently-active account + * independently — collapsing both families to a single index would + * silently unelect one family on every restart. + */ + getActiveIndexByFamily(): { claude: number; gemini: number } /** Enable/disable the account at `index`. Returns true when it changed. */ setAccountEnabled(index: number, enabled: boolean): boolean /** Pin `index` as the active account for every family the dialog cares about. */ @@ -642,15 +649,18 @@ export function createCommandDataService( // the storage mutation. The remove action must persist the index that // the same account will occupy AFTER the removal — unconditionally // resetting to 0 made a non-current removal promote whichever account - // shifted into slot 0 to "active" on the next restart. + // shifted into slot 0 to "active" on the next restart. Each model's + // cursor is tracked independently, so collapse-tokens does not work + // here: a Claude-active account and a Gemini-active account can + // point at different rows. const liveCurrentTokens: { claude?: string; gemini?: string } = {} if (action === 'remove') { const liveAccounts = accountManagerView.getAccounts() - const liveCurrentIndex = accountManagerView.activeIndex() - const liveCurrentToken = - liveAccounts[liveCurrentIndex]?.refreshToken ?? undefined - liveCurrentTokens.claude = liveCurrentToken - liveCurrentTokens.gemini = liveCurrentToken + const liveIndexes = accountManagerView.getActiveIndexByFamily() + liveCurrentTokens.claude = + liveAccounts[liveIndexes.claude]?.refreshToken ?? undefined + liveCurrentTokens.gemini = + liveAccounts[liveIndexes.gemini]?.refreshToken ?? undefined } let foundInStorage = false @@ -698,11 +708,11 @@ export function createCommandDataService( // remove: build the post-removal account list, then resolve the // current-account's NEW index in that list. If the removed - // account was the live current, the current token falls out of - // the list entirely and we fall back to index 0 — matching the - // live AccountManager.removeAccount() behavior (which leaves the - // current index pointing at the same numeric slot, now occupied - // by whichever account shifted in). + // account was the live current, the captured token falls out of + // the list — the live AccountManager leaves the numeric current + // cursor at the SAME index (now occupied by whichever account + // shifted in), so the persisted index must follow that numeric + // slot rather than resetting to 0. const nextAccounts = current.accounts.filter( (account) => account.refreshToken !== refreshToken, ) @@ -715,8 +725,12 @@ export function createCommandDataService( const found = nextAccounts.findIndex( (account) => account.refreshToken === liveToken, ) - if (found === -1) return 0 - return found + if (found !== -1) return found + // The captured current token was the one removed. The live + // manager keeps the numeric cursor at the removed slot, so + // mirror that: clamp the removed index to the new account + // range. + return Math.max(0, Math.min(index, nextAccounts.length - 1)) } const legacyClaude = current.activeIndexByFamily?.claude ?? current.activeIndex diff --git a/packages/opencode/src/plugin/commands.test.ts b/packages/opencode/src/plugin/commands.test.ts index 92c39cb..0e0d87e 100644 --- a/packages/opencode/src/plugin/commands.test.ts +++ b/packages/opencode/src/plugin/commands.test.ts @@ -758,6 +758,65 @@ describe('applyCommand', () => { } }) + it('preserves the live coolingDownUntil when an account action refreshes the sidebar', async () => { + // The dialog path builds sidebar rows from the post-mutation + // command data, which does NOT carry the running cooldown. The + // sidebar refresher must look up the live account's timer by + // index so a rate-limited account doesn't momentarily flash as + // available right after the user toggles / removes / sets active. + const sidebarFile = join(dir, 'sidebar-state.json') + const previousSidebarFile = process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = sidebarFile + try { + const getAccounts = mock(() => [ + { + index: 0, + label: 'Rate-limited', + enabled: true, + coolingDownUntil: 1_900_000_000_000, + }, + { + index: 1, + label: 'Available', + enabled: true, + coolingDownUntil: undefined, + }, + ]) + const refresher = createSidebarRefresher(getAccounts) + const dialogRows = [ + { + id: 'acct-0', + index: 0, + label: 'Rate-limited', + enabled: true, + current: true, + quota: [], + }, + { + id: 'acct-1', + index: 1, + label: 'Available', + enabled: true, + current: false, + quota: [], + }, + ] + + await refresher(dialogRows) + + const state = readSidebarState(sidebarFile) + expect(state.accounts).toHaveLength(2) + expect(state.accounts[0]?.cooldownUntil).toBe(1_900_000_000_000) + expect(state.accounts[1]?.cooldownUntil).toBeUndefined() + } finally { + if (previousSidebarFile === undefined) { + delete process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + } else { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + it('starts a quota refresh when the quota dialog opens without delaying its cached payload', async () => { const listAccounts = mock(async () => [ { diff --git a/packages/opencode/src/plugin/commands.ts b/packages/opencode/src/plugin/commands.ts index 4e84844..85cf86a 100644 --- a/packages/opencode/src/plugin/commands.ts +++ b/packages/opencode/src/plugin/commands.ts @@ -754,29 +754,53 @@ export function createSidebarRefresher( }> | null, ): (accounts?: CommandAccountRow[]) => Promise { return async (dialogAccounts) => { + // Index → live coolingDownUntil. The dialog path rebuilds rows + // from the post-mutation command data, which does not carry the + // running cooldown; looking up the live account by index lets the + // sidebar show a rate-limited account as such instead of + // momentarily marking it available. + const liveByIndex = new Map() + try { + const live = getAccounts() + if (live) { + for (const entry of live) + liveByIndex.set(entry.index, entry.coolingDownUntil) + } + } catch { + // The live provider is best-effort; fall through with an empty + // map if the manager is gone. + } const accounts = dialogAccounts - ? dialogAccounts.map((entry) => ({ - index: entry.index, - label: entry.label, - enabled: entry.enabled, - coolingDownUntil: undefined, - cachedQuota: Object.fromEntries( - entry.quota.flatMap((group) => { - if (group.remainingPercent == null) return [] - return [ - [ - group.key, - { - remainingFraction: group.remainingPercent / 100, - ...(group.resetAt === undefined - ? {} - : { resetTime: new Date(group.resetAt).toISOString() }), - }, - ], - ] - }), - ), - })) + ? dialogAccounts.map((entry) => { + const liveCooldown = liveByIndex.get(entry.index) + return { + index: entry.index, + label: entry.label, + enabled: entry.enabled, + // Preserve the live cooldown when present so an apply + // (setCurrent, toggleEnabled, remove) cannot clear a + // rate-limited account's display. The dialog rows do not + // carry the running timer, so the live source is the + // only authoritative feed. + coolingDownUntil: liveCooldown, + cachedQuota: Object.fromEntries( + entry.quota.flatMap((group) => { + if (group.remainingPercent == null) return [] + return [ + [ + group.key, + { + remainingFraction: group.remainingPercent / 100, + ...(group.resetAt === undefined + ? {} + : { resetTime: new Date(group.resetAt).toISOString() }), + }, + ], + ] + }), + ), + } + }) : getAccounts() if (!accounts || accounts.length === 0) return try { diff --git a/packages/opencode/src/plugin/index.ts b/packages/opencode/src/plugin/index.ts index 55bd2e5..43c7d7a 100644 --- a/packages/opencode/src/plugin/index.ts +++ b/packages/opencode/src/plugin/index.ts @@ -252,6 +252,11 @@ export const createAntigravityPlugin = ? (manager.getCurrentAccountForFamily('claude')?.index ?? 0) : 0 }, + getActiveIndexByFamily: () => { + const manager = lifecycle.getAccountManager() + if (!manager) return { claude: 0, gemini: 0 } + return manager.getActiveIndexByFamily() + }, setAccountEnabled: (index, enabled) => { const manager = lifecycle.getAccountManager() if (!manager) return false From cf3cab107d8d52be428cf886df4d5cb3a1f06735 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 22:56:07 +0200 Subject: [PATCH 19/25] fix(quota): endpoint failover, tolerant bucket parsing, concurrent fetches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P2 #7: fetchQuotaSummary used only options.endpoints[0] for every project attempt, so a 5xx on the primary endpoint permanently lost the request. Iterate the endpoint list per project attempt so a 500 or 429 on one endpoint falls through to the next — matches the legacy fetchers' failover convention. P2 #8: aggregateQuotaSummary derived the pool from group.buckets[0] unconditionally, silently dropping the whole group when an unrecognized bucketId led the array. Use the first RECOGNIZED bucket for pool derivation so the recognized windows are kept. P2 #9: modelCount parsed the raw description prose and counted non-comma / non-colon tokens, so the 'Models within this group:' prefix was counted as a model. Strip the prefix before splitting. P2 #10: the windowed summary fetch and the gemini-CLI quota fetch ran sequentially — two 10s timeouts back-to-back, ~20s per account on modal open. Run them concurrently via Promise.all. --- packages/core/src/quota-manager.test.ts | 149 +++++++++++++++++++-- packages/core/src/quota-manager.ts | 69 ++++++++-- packages/opencode/src/plugin/quota.test.ts | 111 +++++++++++++++ packages/opencode/src/plugin/quota.ts | 114 +++++++++++----- 4 files changed, 387 insertions(+), 56 deletions(-) diff --git a/packages/core/src/quota-manager.test.ts b/packages/core/src/quota-manager.test.ts index 34090f4..3eb2358 100644 --- a/packages/core/src/quota-manager.test.ts +++ b/packages/core/src/quota-manager.test.ts @@ -870,18 +870,103 @@ describe('aggregateQuotaSummary', () => { } }) + it('counts models from the description minus the prefix label', () => { + const response: RetrieveUserQuotaSummaryResponse = { + groups: [ + { + displayName: 'Gemini Models', + description: 'Models within this group: Gemini 3.1 Pro, Flash', + buckets: [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly Limit', + window: 'weekly', + resetTime: '2026-01-08T00:00:00Z', + remainingFraction: 0.7, + }, + ], + }, + ], + } + const summary = aggregateQuotaSummary(response) + expect(summary.groups.gemini!.modelCount).toBe(2) + expect(summary.modelCount).toBe(2) + }) + it('back-compat: treats a windows-less QuotaGroupSummary gracefully', () => { // Legacy shapes may have no `windows` array — consumers read - // `remainingFraction`/`resetTime` directly which are the derived values. - const response = loadFixture('free-ruqs.json') + // `remainingFraction`/`resetTime` directly which are the derived + // values. Exercise this through the REAL response shape (a + // retrieveUserQuotaSummary payload that omits `windows` on each + // bucket), not through a windows-shaped fixture that already + // produced a non-empty `windows` array on the way through. + const response: RetrieveUserQuotaSummaryResponse = { + groups: [ + { + displayName: 'Gemini Legacy', + buckets: [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly', + window: 'weekly', + resetTime: '2025-12-29T00:00:00Z', + remainingFraction: 0.3, + }, + ], + }, + ], + } const summary = aggregateQuotaSummary(response) + // The aggregator still surfaces per-window entries when the + // wire shape includes them — that's the legacy path's only + // window. Legacy consumers that read `remainingFraction` + // directly see the single most-constrained value. + expect(summary.groups.gemini).toBeDefined() + expect(summary.groups.gemini!.remainingFraction).toBeCloseTo(0.3, 3) + expect(summary.groups.gemini!.resetTime).toBe('2025-12-29T00:00:00Z') + }) - // `remainingFraction` is the most-constrained (and only) window - expect(summary.groups.gemini!.remainingFraction).toBe( - summary.groups.gemini!.windows![0]!.remainingFraction, - ) - // Consumers that check `if (entry.windows)` skip per-window rendering - // and fall back to a single QuotaRow — this is the tolerant-read path. + it('keeps a pool when the first bucket is unrecognized', () => { + // Older servers may prepend a junk bucket (system noise, a + // non-standard prefix) whose bucketId doesn't match any of our + // pool prefixes. The legacy derivation used group.buckets[0] + // unconditionally, which silently dropped the whole group when + // an unknown prefix led the array. The pool must be derived + // from the first RECOGNIZED bucket, not the first bucket. + const response: RetrieveUserQuotaSummaryResponse = { + groups: [ + { + displayName: 'Gemini Models', + buckets: [ + { + bucketId: 'unknown-prefix-noise', + displayName: 'Unknown', + window: 'weekly', + resetTime: '2026-01-01T00:00:00Z', + remainingFraction: 0, + }, + { + bucketId: 'gemini-weekly', + displayName: 'Weekly Limit', + window: 'weekly', + resetTime: '2026-01-08T00:00:00Z', + remainingFraction: 0.7, + }, + { + bucketId: 'gemini-5h', + displayName: '5h Limit', + window: '5h', + resetTime: '2026-01-01T05:00:00Z', + remainingFraction: 0.85, + }, + ], + }, + ], + } + const summary = aggregateQuotaSummary(response) + expect(summary.groups.gemini).toBeDefined() + expect(summary.groups.gemini!.windows).toHaveLength(2) + expect(summary.groups.gemini!.remainingFraction).toBeCloseTo(0.7, 3) }) }) @@ -1045,4 +1130,52 @@ describe('fetchQuotaSummary', () => { }) expect(receivedProjectId).toBe('managed-from-record') }) + + it('fails over to the next endpoint when the first returns 500', async () => { + // The legacy fetchers iterate the endpoint list per project attempt + // (fetchAvailableModels, fetchGeminiCliQuota). fetchQuotaSummary + // must match: a 500 on the primary endpoint should fall through + // to the next entry in `options.endpoints`. + const summary: RetrieveUserQuotaSummaryResponse = { + groups: [ + { + displayName: 'Recovery', + buckets: [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly', + window: 'weekly', + resetTime: '2026-01-01T00:00:00Z', + remainingFraction: 0.42, + }, + ], + }, + ], + } + const visited: string[] = [] + const fetchVia = async (url: string): Promise => { + visited.push(url) + if (url.startsWith('https://failover-a.test')) { + return new Response('boom', { status: 500 }) + } + return new Response(JSON.stringify(summary), { status: 200 }) + } + const result = await fetchQuotaSummary({ + accessToken: 'tok', + managedProjectId: 'mp', + endpoints: [ + 'https://failover-a.test', + 'https://failover-b.test', + ] as const, + fetchVia: fetchVia as any, + }) + expect(visited).toHaveLength(2) + expect(visited[0]).toBe( + 'https://failover-a.test/v1internal:retrieveUserQuotaSummary', + ) + expect(visited[1]).toBe( + 'https://failover-b.test/v1internal:retrieveUserQuotaSummary', + ) + expect(result.summary.groups[0]!.buckets[0]!.remainingFraction).toBe(0.42) + }) }) diff --git a/packages/core/src/quota-manager.ts b/packages/core/src/quota-manager.ts index 821a854..366b211 100644 --- a/packages/core/src/quota-manager.ts +++ b/packages/core/src/quota-manager.ts @@ -610,6 +610,27 @@ function poolForBucketId(bucketId: string): QuotaGroup | null { return null } +/** + * Count models listed in a group's description. The description is + * typically a comma-separated list of model names ("Claude Sonnet 4.6, + * Gemini 3.1 Pro, Flash") often prefixed with a label like + * "Models within this group:". Strip the prefix before splitting so + * the label itself is not counted as a model. Returns 0 for shapes + * that don't match a comma-separated list. + */ +function parseDescriptionModelCount(description: string): number { + const prefixMatch = description.match(/^[^:]+:\s*/) + const payload = prefixMatch + ? description.slice(prefixMatch[0].length) + : description + if (!payload) return 0 + const entries = payload + .split(',') + .map((entry) => entry.trim()) + .filter((entry) => entry.length > 0) + return entries.length +} + /** * Aggregate a retrieveUserQuotaSummary response into a QuotaSummary. * @@ -643,13 +664,24 @@ export function aggregateQuotaSummary( }) const constrained = mostConstrainedWindow(windows) - const firstBucket = group.buckets[0] - if (!firstBucket) continue - const pool = poolForBucketId(firstBucket.bucketId) + // Pick the first RECOGNIZED bucket for pool derivation. Older + // servers may prepend a junk bucket (system noise, a non-standard + // prefix like `claude-3p-`) whose bucketId doesn't match any of + // our pool prefixes; the legacy poolForBucketId derivation used + // group.buckets[0] unconditionally, which silently dropped the + // whole group when an unknown prefix led the array. + const recognizedBucket = group.buckets.find((bucket) => + poolForBucketId(bucket.bucketId), + ) + if (!recognizedBucket) continue + const pool = poolForBucketId(recognizedBucket.bucketId) if (!pool || !constrained) continue + // The description is a list of model names comma-separated, often + // prefixed with a label like "Models within this group:". Strip + // the prefix before splitting so it doesn't count as a model. const modelCount = group.description - ? (group.description.match(/[^,:]+/g)?.length ?? 0) + ? parseDescriptionModelCount(group.description) : 0 groups[pool] = { @@ -690,9 +722,12 @@ export interface FetchQuotaSummaryResult { * Fetch the windowed quota summary via `retrieveUserQuotaSummary`. * * Uses the same transport/UA/timeout conventions as `fetchAvailableModels`. - * On 403 with the managedProjectId, retries with the regular projectId. - * If that also 403s, falls back to `fetchAvailableModels` so quota never - * goes dark. On missing managedProjectId, tries projectId first. + * On a 429 or 5xx against one endpoint, falls through to the next entry + * in `options.endpoints` (matching the legacy fetchers' failover + * convention). On 403 with the managedProjectId, retries with the + * regular projectId. If that also 403s, falls back to + * `fetchAvailableModels` so quota never goes dark. On missing + * managedProjectId, tries projectId first. */ export async function fetchQuotaSummary( options: FetchQuotaSummaryOptions, @@ -702,12 +737,12 @@ export async function fetchQuotaSummary( const transport = options.fetchVia ?? defaultTransport const errors: string[] = [] - const endpoint = options.endpoints[0] - if (!endpoint) { + if (options.endpoints.length === 0) { throw new Error('No endpoints configured for fetchQuotaSummary') } const tryBody = async ( + endpoint: string, projectId: string, ): Promise => { const body = { project: projectId } @@ -744,6 +779,8 @@ export async function fetchQuotaSummary( errors.push( `retrieveUserQuotaSummary ${status} at ${endpoint}${message ? `: ${message.trim().slice(0, 200)}` : ''}`, ) + // Endpoint failover: the caller may have multiple endpoints; + // continue to the next one with the same project ID. return null } @@ -761,10 +798,14 @@ export async function fetchQuotaSummary( } // Try managedProjectId first, then projectId, then legacy fallback. + // For each project, iterate the endpoint list so a 500/429 on the + // primary endpoint falls through to the next entry. const primary = options.managedProjectId ?? options.projectId if (primary) { - const result = await tryBody(primary) - if (result) return { summary: result } + for (const endpoint of options.endpoints) { + const result = await tryBody(endpoint, primary) + if (result) return { summary: result } + } } // If primary failed with 403 and we used managedProjectId, @@ -776,8 +817,10 @@ export async function fetchQuotaSummary( ? options.projectId : undefined if (fallbackId) { - const result = await tryBody(fallbackId) - if (result) return { summary: result } + for (const endpoint of options.endpoints) { + const result = await tryBody(endpoint, fallbackId) + if (result) return { summary: result } + } } // Give up — the caller should fall back to fetchAvailableModels. diff --git a/packages/opencode/src/plugin/quota.test.ts b/packages/opencode/src/plugin/quota.test.ts index e142ecb..82be844 100644 --- a/packages/opencode/src/plugin/quota.test.ts +++ b/packages/opencode/src/plugin/quota.test.ts @@ -176,6 +176,117 @@ describe('pushSidebarQuotaSnapshot', () => { expect(state.accounts).toEqual([]) }) + it('runs the windowed summary and gemini-cli quota fetch concurrently', async () => { + // The summary fetch and the gemini-CLI quota fetch previously + // ran sequentially — two 10s timeouts back-to-back. Run them + // concurrently instead. We assert by recording the sequence + // numbers of each fetch via a gate so the test is + // deterministic across runtimes. + const summarySeq: { seq: number } = { seq: 0 } + const cliSeq: { seq: number } = { seq: 0 } + let releaseSummary!: () => void + let releaseCli!: () => void + const summaryGate = new Promise((resolve) => { + releaseSummary = resolve + }) + const cliGate = new Promise((resolve) => { + releaseCli = resolve + }) + + const fetchSpy = spyOn(globalThis, 'fetch').mockImplementation((async ( + input: unknown, + ) => { + const url = + typeof input === 'string' + ? input + : input instanceof URL + ? input.toString() + : (input as Request).url + if (url.includes('retrieveUserQuotaSummary')) { + summarySeq.seq += 1 + await summaryGate + summarySeq.seq += 1 + return new Response( + JSON.stringify({ + groups: [ + { + displayName: 'Gemini Models', + buckets: [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly', + window: 'weekly', + resetTime: '2026-01-08T00:00:00Z', + remainingFraction: 0.7, + }, + ], + }, + ], + }), + { status: 200 }, + ) + } + if (url.includes('retrieveUserQuota')) { + cliSeq.seq += 1 + await cliGate + cliSeq.seq += 1 + return new Response(JSON.stringify({ buckets: [] }), { status: 200 }) + } + // Token refresh + anything else: return a pliable JSON + // response so the rest of the quota pipeline can carry on. + return new Response( + JSON.stringify({ + access_token: 'access-token', + expires_in: 3600, + }), + { status: 200 }, + ) + }) as unknown as typeof fetch) + + const client = { + auth: { set: mock(async () => {}) }, + } as unknown as PluginClient + // Use a UNIQUE refresh token per test run so the QuotaManager's + // singleton cache / backoff state from earlier tests in the + // same run can't skip the fetch behind our spy. + const account: AccountMetadataV3 = { + refreshToken: `concurrent-fetch-${Date.now()}-${Math.random()}`, + managedProjectId: 'managed-project', + projectId: 'project-id', + addedAt: 0, + lastUsed: 0, + } + const manager = createOpenCodeQuotaManager(client, 'google') + + try { + const refresh = manager.refreshAccounts([account], { + indexFor: () => 0, + force: true, + }) + // Yield until both fetches have started (seq=1). + const deadline = Date.now() + 5_000 + while ((summarySeq.seq < 1 || cliSeq.seq < 1) && Date.now() < deadline) { + await new Promise((resolve) => setTimeout(resolve, 1)) + } + // Both fetches started before either finished. If they ran + // sequentially, the cli fetch would not have started yet. + expect(summarySeq.seq).toBe(1) + expect(cliSeq.seq).toBe(1) + releaseSummary() + releaseCli() + await refresh + + // Both fetches completed (seq=2). + expect(summarySeq.seq).toBe(2) + expect(cliSeq.seq).toBe(2) + } finally { + fetchSpy.mockRestore() + // Yield once more so the spy is fully torn down before the + // next test's bun event loop watches see the partial state. + await new Promise((resolve) => setImmediate(resolve)) + } + }) + it('fences the real quota wrapper sidebar enqueue before the lifecycle drain', async () => { const events: string[] = [] let releaseFetch!: () => void diff --git a/packages/opencode/src/plugin/quota.ts b/packages/opencode/src/plugin/quota.ts index f98489e..bc913fb 100644 --- a/packages/opencode/src/plugin/quota.ts +++ b/packages/opencode/src/plugin/quota.ts @@ -319,6 +319,43 @@ export async function pushSidebarQuotaSnapshot( } } +/** + * Legacy fallback: fetchAvailableModels → aggregateQuota. Used when + * `fetchQuotaSummary` rejects (network, 403, etc.). Extracted so the + * concurrent fetch path can reuse the same fallback logic without + * duplicating the catch chain. + */ +async function fetchLegacyModelsFallback(options: { + accessToken: string + projectId: string + fetchVia?: QuotaFetch +}): Promise { + try { + const modelsResponse = await fetchAvailableModels({ + accessToken: options.accessToken, + projectId: options.projectId, + endpoints: ANTIGRAVITY_ENDPOINT_FALLBACKS, + userAgent: buildAntigravityHarnessUserAgent(), + timeoutMs: 10_000, + ...(options.fetchVia ? { fetchVia: options.fetchVia } : {}), + }) + if (modelsResponse.models) { + return aggregateQuota(modelsResponse.models) + } + return { + groups: {}, + modelCount: 0, + error: 'Failed to fetch Antigravity quota (legacy fallback)', + } + } catch { + return { + groups: {}, + modelCount: 0, + error: 'Failed to fetch Antigravity quota', + } + } +} + function makeFetchAccountQuota( client: PluginClient | undefined, providerId: string, @@ -383,62 +420,69 @@ function makeFetchAccountQuota( const managedProjectId = authParts.managedProjectId ?? account.managedProjectId - // Primary: retrieveUserQuotaSummary (windowed quota). - // Fall back to fetchAvailableModels on any failure (403, network, etc.) - // so quota never goes dark. - try { - const summaryResult = await fetchQuotaSummary({ - accessToken: auth.access ?? '', - managedProjectId, - projectId: projectContext.effectiveProjectId, - endpoints: ANTIGRAVITY_ENDPOINT_FALLBACKS, - userAgent: buildAntigravityHarnessUserAgent(), - timeoutMs: 10_000, - ...(fetchVia ? { fetchVia } : {}), - }) - quotaResult = aggregateQuotaSummary(summaryResult.summary) - if (summaryResult.fellBackToLegacy) { - fellBackToLegacy = true - } - } catch { - // Fall back to fetchAvailableModels-based path. - fellBackToLegacy = true + // Two independent payload contracts: the windowed summary + // (with legacy fallback) and the gemini-CLI quota. They share + // access + project but target different endpoints, so the + // two 10s timeouts ran back-to-back for ~20s per account on + // modal open. Run them concurrently; either rejection is + // handled by its own branch and the result is still merged. + const fetchSummaryPayload = (async (): Promise<{ + result: QuotaSummary + fellBackToLegacy: boolean + }> => { try { - const modelsResponse = await fetchAvailableModels({ + const summaryResult = await fetchQuotaSummary({ accessToken: auth.access ?? '', + managedProjectId, projectId: projectContext.effectiveProjectId, endpoints: ANTIGRAVITY_ENDPOINT_FALLBACKS, userAgent: buildAntigravityHarnessUserAgent(), timeoutMs: 10_000, ...(fetchVia ? { fetchVia } : {}), }) - if (modelsResponse.models) { - quotaResult = aggregateQuota(modelsResponse.models) - } else { - quotaResult = { - groups: {}, - modelCount: 0, - error: 'Failed to fetch Antigravity quota (legacy fallback)', - } + return { + result: aggregateQuotaSummary(summaryResult.summary), + fellBackToLegacy: summaryResult.fellBackToLegacy ?? false, } } catch { - quotaResult = { - groups: {}, - modelCount: 0, - error: 'Failed to fetch Antigravity quota', + return { + result: await fetchLegacyModelsFallback({ + accessToken: auth.access ?? '', + projectId: projectContext.effectiveProjectId, + fetchVia, + }), + fellBackToLegacy: true, } } - } + })() - const geminiCliResponse = await fetchGeminiCliQuota({ + const fetchGeminiCliPayload = fetchGeminiCliQuota({ accessToken: auth.access ?? '', projectId: projectContext.effectiveProjectId, endpoints: ANTIGRAVITY_ENDPOINT_FALLBACKS, userAgent: buildGeminiCliUserAgent(), timeoutMs: 10_000, ...(fetchVia ? { fetchVia } : {}), + }).catch((error: unknown) => { + // The previous sequence (`await fetchGeminiCliQuota`) let + // any throw bubble to the outer try/catch and produced an + // undefined buckets downstream. Preserve that semantics so + // the parallel path is a drop-in replacement. + log.debug('fetchGeminiCliQuota failed', { + error: error instanceof Error ? error.message : String(error), + }) + return { buckets: undefined } as Awaited< + ReturnType + > }) + const [summary, geminiCliResponse] = await Promise.all([ + fetchSummaryPayload, + fetchGeminiCliPayload, + ]) + quotaResult = summary.result + fellBackToLegacy = summary.fellBackToLegacy + const geminiCliQuotaResult = aggregateGeminiCliQuota(geminiCliResponse) const annotated: GeminiCliQuotaSummary = geminiCliResponse.buckets === undefined || From 0aa32481b9a8c6cf151e7ca7ceb8d831b603c3c3 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 22:56:28 +0200 Subject: [PATCH 20/25] test: strengthen oauth/removal/back-compat coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P3 #11: the e2e mock's quotaSummaryWindow fixture dropped fixture headers before writing the body. Apply them via the shared applyHeaders seam so cache-bust and trace-header tests have a path. P3 #12: the quota-manager 'back-compat' test loaded a windows-shaped fixture rather than exercising a windows-less legacy shape through the real read path. Rewire it through a hand-rolled RetrieveUserQuotaSummaryResponse that omits per-window data. P3 #13: account-manager 'remove-during-refresh' test never attempted the quota write, so the expectedRefreshToken guard wasn't exercised. Call updateQuotaCache(0, …, refreshTokenForA) after the removal so the assertion covers the actual cross-account guard. P3 #14: the sidebar-state 'half-missing stamp' test covered only the cachedQuotaAccountId-only case. Add the symmetric currentQuotaAccountId-only case so legacy snapshots missing the persisted stamp don't silently drop quota. P3 #15: account-command-oauth failure messages said 'Please try again' even though the pending entry was consumed and the same code could not be redeemed. Update the messages to 'start a new OAuth flow' so the operator knows the OAuth session is gone. P3 #16: tighten the account-command-oauth test assertions — pin the failed-exchange response, pin the full persisted success object (not just label), assert 'Work account' renders in the add-oauth-finish dialog, and restore an explicit unknown-key legacy fixture for the renderer. The tui-compiled mirror of command-data.ts is regenerated by the build:tui gate so the bundled TUI payload matches the source. --- packages/core/src/account-manager.test.ts | 12 +++++- .../e2e-tests/src/mock-antigravity-server.ts | 21 +++++++--- .../src/plugin/account-command-oauth.test.ts | 29 ++++++++++---- .../src/plugin/account-command-oauth.ts | 10 ++--- packages/opencode/src/sidebar-state.test.ts | 14 +++++++ .../src/tui-compiled/plugin/command-data.ts | 40 +++++++++++++------ .../opencode/src/tui/command-dialogs.test.tsx | 8 ++++ 7 files changed, 102 insertions(+), 32 deletions(-) diff --git a/packages/core/src/account-manager.test.ts b/packages/core/src/account-manager.test.ts index 55b060d..64b5edf 100644 --- a/packages/core/src/account-manager.test.ts +++ b/packages/core/src/account-manager.test.ts @@ -250,12 +250,20 @@ describe('core AccountManager', () => { expect(manager.getAccounts()[0]?.parts.refreshToken).toBe('r2') // The async refresh finally resolves. The caller re-resolves the - // live index for `r1` (which is now `-1`) and skips the write — the - // AccountManager's existing expectedRefreshToken guard enforces that. + // live index for `r1` (which is now `-1`) and the quota write is + // then attempted via `updateQuotaCache` at index 0 with the + // captured `expectedRefreshToken`. The guard MUST drop the write + // because the captured token no longer matches the account at + // index 0 — B would otherwise receive A's quota percentages. const liveIndex = manager .getAccounts() .findIndex((entry) => entry.parts.refreshToken === refreshTokenForA) expect(liveIndex).toBe(-1) + manager.updateQuotaCache( + 0, + { gemini: { remainingFraction: 0.42, modelCount: 1 } }, + refreshTokenForA, + ) // No quota should have landed on whichever account shifted into // index 0. expect(manager.getAccounts()[0]?.cachedQuota).toBeUndefined() diff --git a/packages/e2e-tests/src/mock-antigravity-server.ts b/packages/e2e-tests/src/mock-antigravity-server.ts index d987448..76ce1d5 100644 --- a/packages/e2e-tests/src/mock-antigravity-server.ts +++ b/packages/e2e-tests/src/mock-antigravity-server.ts @@ -361,11 +361,22 @@ export async function startMockAntigravityServer(): Promise { return } } - sendJson(response, fixture.status ?? 200, { - groups: fixture.groups, - description: - 'Within each group, models share a weekly limit and a 5-hour limit.', - }) + // Apply the fixture's custom headers (e.g. for cache-bust + // trace assertions) before writing the JSON body — mirrors + // every other fixture's header-merge convention. + applyHeaders( + response, + { 'content-type': 'application/json' }, + fixture.headers, + ) + response.writeHead(fixture.status ?? 200) + response.end( + JSON.stringify({ + groups: fixture.groups, + description: + 'Within each group, models share a weekly limit and a 5-hour limit.', + }), + ) return } case 'generateContent': { diff --git a/packages/opencode/src/plugin/account-command-oauth.test.ts b/packages/opencode/src/plugin/account-command-oauth.test.ts index 95b852d..31ad478 100644 --- a/packages/opencode/src/plugin/account-command-oauth.test.ts +++ b/packages/opencode/src/plugin/account-command-oauth.test.ts @@ -21,7 +21,7 @@ describe('createAccountCommandOAuthService', () => { })) const exchange = mock(async () => ({ type: 'failed' as const, - error: 'unused', + error: 'invalid_grant', })) const service = createAccountCommandOAuthService({ authorize, @@ -31,12 +31,18 @@ describe('createAccountCommandOAuthService', () => { }) const started = await service.start('session-1') - await service.finish('session-1', 'callback-code') + const finished = await service.finish('session-1', 'callback-code') expect(authorize).toHaveBeenCalledTimes(1) expect(started.url).toContain('accounts.google.test') expect(started.accounts).toEqual(initialRows) expect(exchange).toHaveBeenCalledWith('callback-code', 'signed-state') + // Failed-exchange response: the pending entry is consumed, the + // operator must start a new OAuth flow to retry. + expect(finished.text).toBe( + 'OAuth authentication failed. Please start a new OAuth flow and try again.', + ) + expect(finished.accounts).toEqual(initialRows) }) it('exchanges, persists, and returns updated label-only account rows', async () => { @@ -78,9 +84,18 @@ describe('createAccountCommandOAuthService', () => { 'Work account', ) - expect(persisted).toHaveBeenCalledWith( - expect.objectContaining({ label: 'Work account' }), - ) + // The persisted success object is the FULL exchanged result with + // the operator-provided label override; we pin every field so a + // regression that drops e.g. projectId or email is caught. + expect(persisted).toHaveBeenCalledWith({ + type: 'success', + refresh: 'refresh-token|project', + access: 'access-token', + expires: 123, + email: 'private@example.test', + label: 'Work account', + projectId: 'project', + }) expect(finished).toEqual({ text: 'OAuth account added.', accounts: updatedRows, @@ -241,7 +256,7 @@ describe('createAccountCommandOAuthService', () => { expect(persist).toHaveBeenCalledTimes(1) expect(finished.text).toBe( - 'OAuth account could not be saved to disk. Please try again.', + 'OAuth account could not be saved to disk. Please start a new OAuth flow.', ) }) @@ -266,7 +281,7 @@ describe('createAccountCommandOAuthService', () => { const finished = await service.finish('session-1', 'callback-code') expect(finished.text).toBe( - 'OAuth exchange failed due to a network error. Please try again.', + 'OAuth exchange failed due to a network error. Please start a new OAuth flow.', ) expect(persist).not.toHaveBeenCalled() }) diff --git a/packages/opencode/src/plugin/account-command-oauth.ts b/packages/opencode/src/plugin/account-command-oauth.ts index 095e7c4..846870e 100644 --- a/packages/opencode/src/plugin/account-command-oauth.ts +++ b/packages/opencode/src/plugin/account-command-oauth.ts @@ -126,13 +126,13 @@ export function createAccountCommandOAuthService( callback = parseOAuthCallbackInput(callbackInput, pending.state) } catch { return { - text: 'OAuth authentication failed: could not parse the callback. Please try again.', + text: 'OAuth authentication failed: could not parse the callback. Please start a new OAuth flow.', accounts: await options.listAccounts(), } } if ('error' in callback) { return { - text: `OAuth authentication failed: ${callback.error}`, + text: `OAuth authentication failed: ${callback.error}. Please start a new OAuth flow.`, accounts: await options.listAccounts(), } } @@ -145,13 +145,13 @@ export function createAccountCommandOAuthService( result = await options.exchange(callback.code, callback.state) } catch { return { - text: 'OAuth exchange failed due to a network error. Please try again.', + text: 'OAuth exchange failed due to a network error. Please start a new OAuth flow.', accounts: await options.listAccounts(), } } if (result.type === 'failed') { return { - text: 'OAuth authentication failed. Please check the code and try again.', + text: 'OAuth authentication failed. Please start a new OAuth flow and try again.', accounts: await options.listAccounts(), } } @@ -167,7 +167,7 @@ export function createAccountCommandOAuthService( await options.persist(persisted) } catch { return { - text: 'OAuth account could not be saved to disk. Please try again.', + text: 'OAuth account could not be saved to disk. Please start a new OAuth flow.', accounts: await options.listAccounts(), } } diff --git a/packages/opencode/src/sidebar-state.test.ts b/packages/opencode/src/sidebar-state.test.ts index a664ea2..cbf24d1 100644 --- a/packages/opencode/src/sidebar-state.test.ts +++ b/packages/opencode/src/sidebar-state.test.ts @@ -553,6 +553,20 @@ describe('redaction', () => { }) expect(redacted.quota.gemini?.remainingPercent).toBe(30) }) + + it('preserves cachedQuota when only the current quota account id is provided (fail open for legacy)', () => { + // Symmetric half-missing case: the live snapshot has a stamp + // (provider added currentQuotaAccountId) but the persisted cache + // row does not yet (legacy). The projection must not silently + // drop the quota; the absence of the persisted stamp alone is + // not enough to mark the cache as stale. + const redacted = redactAccountForSidebar({ + index: 0, + cachedQuota: { 'non-gemini': { remainingFraction: 0.6 } }, + currentQuotaAccountId: 'b'.repeat(16), + }) + expect(redacted.quota['non-gemini']?.remainingPercent).toBe(60) + }) }) describe('windows rework — producer seam tests', () => { diff --git a/packages/opencode/src/tui-compiled/plugin/command-data.ts b/packages/opencode/src/tui-compiled/plugin/command-data.ts index bf80ef0..5170e82 100644 --- a/packages/opencode/src/tui-compiled/plugin/command-data.ts +++ b/packages/opencode/src/tui-compiled/plugin/command-data.ts @@ -308,6 +308,13 @@ export interface CommandDataAccountManagerView { requestSaveToDisk(): void flushSaveToDisk(): Promise activeIndex(): number + /** + * Per-family live cursor. Used by the remove path so the persisted + * `activeIndexByFamily` follows each model's currently-active account + * independently — collapsing both families to a single index would + * silently unelect one family on every restart. + */ + getActiveIndexByFamily(): { claude: number; gemini: number } /** Enable/disable the account at `index`. Returns true when it changed. */ setAccountEnabled(index: number, enabled: boolean): boolean /** Pin `index` as the active account for every family the dialog cares about. */ @@ -642,15 +649,18 @@ export function createCommandDataService( // the storage mutation. The remove action must persist the index that // the same account will occupy AFTER the removal — unconditionally // resetting to 0 made a non-current removal promote whichever account - // shifted into slot 0 to "active" on the next restart. + // shifted into slot 0 to "active" on the next restart. Each model's + // cursor is tracked independently, so collapse-tokens does not work + // here: a Claude-active account and a Gemini-active account can + // point at different rows. const liveCurrentTokens: { claude?: string; gemini?: string } = {} if (action === 'remove') { const liveAccounts = accountManagerView.getAccounts() - const liveCurrentIndex = accountManagerView.activeIndex() - const liveCurrentToken = - liveAccounts[liveCurrentIndex]?.refreshToken ?? undefined - liveCurrentTokens.claude = liveCurrentToken - liveCurrentTokens.gemini = liveCurrentToken + const liveIndexes = accountManagerView.getActiveIndexByFamily() + liveCurrentTokens.claude = + liveAccounts[liveIndexes.claude]?.refreshToken ?? undefined + liveCurrentTokens.gemini = + liveAccounts[liveIndexes.gemini]?.refreshToken ?? undefined } let foundInStorage = false @@ -698,11 +708,11 @@ export function createCommandDataService( // remove: build the post-removal account list, then resolve the // current-account's NEW index in that list. If the removed - // account was the live current, the current token falls out of - // the list entirely and we fall back to index 0 — matching the - // live AccountManager.removeAccount() behavior (which leaves the - // current index pointing at the same numeric slot, now occupied - // by whichever account shifted in). + // account was the live current, the captured token falls out of + // the list — the live AccountManager leaves the numeric current + // cursor at the SAME index (now occupied by whichever account + // shifted in), so the persisted index must follow that numeric + // slot rather than resetting to 0. const nextAccounts = current.accounts.filter( (account) => account.refreshToken !== refreshToken, ) @@ -715,8 +725,12 @@ export function createCommandDataService( const found = nextAccounts.findIndex( (account) => account.refreshToken === liveToken, ) - if (found === -1) return 0 - return found + if (found !== -1) return found + // The captured current token was the one removed. The live + // manager keeps the numeric cursor at the removed slot, so + // mirror that: clamp the removed index to the new account + // range. + return Math.max(0, Math.min(index, nextAccounts.length - 1)) } const legacyClaude = current.activeIndexByFamily?.claude ?? current.activeIndex diff --git a/packages/opencode/src/tui/command-dialogs.test.tsx b/packages/opencode/src/tui/command-dialogs.test.tsx index 3615254..fb25e5e 100644 --- a/packages/opencode/src/tui/command-dialogs.test.tsx +++ b/packages/opencode/src/tui/command-dialogs.test.tsx @@ -1008,6 +1008,14 @@ describe('openCommandDialog (imperative dispatcher)', () => { 'add-oauth-finish callback-code --label Work account', ) expect(applyMock.mock.calls.at(1)?.[2]?.timeoutMs).toBe(120_000) + // The post-finish dialog must surface the newly added account + // label so the user sees the result of the OAuth flow without + // reopening the dialog. + await renderDialog(localFake) + const postFinishTitles = (localFake.capturedSelectProps?.options ?? []).map( + (option) => option.title, + ) + expect(postFinishTitles).toContain('Work account') }) it('antigravity-account opens a row subdialog with toggle/current/remove/back', async () => { From 7d318878fe0605b1dc00902dc499c9fee96de2d7 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 22:59:53 +0200 Subject: [PATCH 21/25] fix(test): pin fences test events for concurrent quota fetches The fences test asserts the EXACT sequence of fetch starts so a regression that bypasses the lifecycle drain guard surfaces. With the new concurrent quota path (summary + gemini-cli run in parallel), each refreshAccounts call now performs 2 fetches instead of 1, so the expected list grows by 1. --- packages/opencode/src/plugin/quota.test.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/opencode/src/plugin/quota.test.ts b/packages/opencode/src/plugin/quota.test.ts index 82be844..5225605 100644 --- a/packages/opencode/src/plugin/quota.test.ts +++ b/packages/opencode/src/plugin/quota.test.ts @@ -370,6 +370,7 @@ describe('pushSidebarQuotaSnapshot', () => { 'fetch:start', 'fetch:start', 'fetch:start', + 'fetch:start', 'sidebar:write-start', 'lifecycle:drain', 'drain:sees-sidebar-write', From ffeffde35a799c5620c0239350b9d2ebef2a4781 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Fri, 24 Jul 2026 23:32:31 +0200 Subject: [PATCH 22/25] test: make windowed-quota and record-fallback regressions effective --- package.json | 2 +- packages/core/src/account-manager.test.ts | 11 ++- packages/core/src/quota-manager.test.ts | 14 ++- .../src/mock-antigravity-server.test.ts | 90 +++++++++++++++++++ .../e2e-tests/src/mock-antigravity-server.ts | 2 +- 5 files changed, 106 insertions(+), 13 deletions(-) diff --git a/package.json b/package.json index 079ace6..2282d0b 100644 --- a/package.json +++ b/package.json @@ -9,7 +9,7 @@ "build": "bun run --cwd packages/core build && bun run --cwd packages/opencode build && bun run --cwd packages/pi build", "typecheck": "bun run --cwd packages/core build && bun run --cwd packages/opencode typecheck && bun run --cwd packages/pi typecheck && tsc -p tsconfig.scripts.json", "test": "bun run --cwd packages/core build && bun test --isolate packages/core/src packages/opencode/src packages/pi/src test/", - "test:e2e": "bun test --isolate ./packages/e2e-tests/src/plugin-flow.e2e.test.ts ./packages/e2e-tests/src/cli-flow.e2e.test.ts ./packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts ./packages/e2e-tests/src/fetch-guard.test.ts", + "test:e2e": "bun test --isolate ./packages/e2e-tests/src/plugin-flow.e2e.test.ts ./packages/e2e-tests/src/cli-flow.e2e.test.ts ./packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts ./packages/e2e-tests/src/fetch-guard.test.ts ./packages/e2e-tests/src/mock-antigravity-server.test.ts", "test:e2e:models": "bun run --cwd packages/opencode test:e2e:models", "test:e2e:regression": "bun run --cwd packages/opencode test:e2e:regression", "prepublishOnly": "bun run build", diff --git a/packages/core/src/account-manager.test.ts b/packages/core/src/account-manager.test.ts index 64b5edf..62f412f 100644 --- a/packages/core/src/account-manager.test.ts +++ b/packages/core/src/account-manager.test.ts @@ -304,7 +304,7 @@ describe('AccountManager instance dependencies', () => { }) describe('managedProjectId projection', () => { - it('getAccountsForQuotaCheck returns managedProjectId from the record when refresh token is bare', () => { + it('getAccountsForQuotaCheck falls back to record managedProjectId when parts lack it', () => { const stored: AccountStorageV4 = { version: 4, accounts: [ @@ -323,6 +323,11 @@ describe('managedProjectId projection', () => { store: createStore(stored).store, now: () => 1_000, }) + // Simulate a bare-token rotation that strips managedProjectId from + // parts — the record-level field is the only remaining source. + const allAccounts = manager.getAccounts() + allAccounts[0]!.parts.managedProjectId = undefined + const accounts = manager.getAccountsForQuotaCheck() expect(accounts).toHaveLength(1) expect(accounts[0]!.projectId).toBe('my-project') @@ -349,9 +354,9 @@ describe('managedProjectId projection', () => { store, now: () => 1_000, }) - // Trigger a save — buildStorageSnapshot must preserve managedProjectId. + // Trigger a save — dispose clears the debounce and forces it immediately. manager.requestSaveToDisk() - await manager.flushSaveToDisk() + await manager.dispose() const saved = state() expect(saved?.accounts[0]?.managedProjectId).toBe('my-managed-project') diff --git a/packages/core/src/quota-manager.test.ts b/packages/core/src/quota-manager.test.ts index 3eb2358..13213c4 100644 --- a/packages/core/src/quota-manager.test.ts +++ b/packages/core/src/quota-manager.test.ts @@ -1106,11 +1106,10 @@ describe('fetchQuotaSummary', () => { expect(receivedProjectId).toBe('regular-proj') }) - it('uses managedProjectId from options when present (the record fallback path)', async () => { - // Regression: bare refresh tokens have no packed managedProjectId. - // The caller must supply it from the account record. This test - // asserts that fetchQuotaSummary posts the options.managedProjectId - // when it is provided. + it('uses managedProjectId when only managedProjectId is provided (managed-only path)', async () => { + // Distinct from the precedence test above: only managedProjectId is + // supplied — no projectId fallback. The managed-only input path must + // still post `project` with the managed project ID. let receivedProjectId = '' const summary: RetrieveUserQuotaSummaryResponse = { groups: [] } const fetchVia = async ( @@ -1123,12 +1122,11 @@ describe('fetchQuotaSummary', () => { } await fetchQuotaSummary({ accessToken: 'tok', - managedProjectId: 'managed-from-record', - projectId: 'regular-proj', + managedProjectId: 'managed-only', endpoints: ENDPOINTS, fetchVia: fetchVia as any, }) - expect(receivedProjectId).toBe('managed-from-record') + expect(receivedProjectId).toBe('managed-only') }) it('fails over to the next endpoint when the first returns 500', async () => { diff --git a/packages/e2e-tests/src/mock-antigravity-server.test.ts b/packages/e2e-tests/src/mock-antigravity-server.test.ts index c502906..6330f90 100644 --- a/packages/e2e-tests/src/mock-antigravity-server.test.ts +++ b/packages/e2e-tests/src/mock-antigravity-server.test.ts @@ -164,3 +164,93 @@ describe('mock antigravity server', () => { await pending }) }) + +// =========================================================================== +// quotaSummaryWindow — managedProjectId 403 gate +// =========================================================================== + +describe('quotaSummaryWindow managedProjectId 403 gate', () => { + it('returns 200 when project matches managedProjectId', async () => { + await withServer(async (server) => { + server.enqueue({ + kind: 'quotaSummaryWindow', + managedProjectId: 'managed-rpc', + groups: [ + { + displayName: 'Gemini Models', + buckets: [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly Limit', + window: 'weekly', + resetTime: '2026-07-31T00:00:00Z', + remainingFraction: 0.7, + }, + ], + }, + ], + }) + const response = await fetch( + `${server.baseUrl}/v1internal:retrieveUserQuotaSummary`, + { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ project: 'managed-rpc' }), + }, + ) + expect(response.status).toBe(200) + const body = (await response.json()) as { + groups: Array<{ + displayName: string + buckets: Array<{ bucketId: string; remainingFraction: number }> + }> + } + expect(body.groups).toHaveLength(1) + expect(body.groups[0]!.buckets[0]!.bucketId).toBe('gemini-weekly') + expect(body.groups[0]!.buckets[0]!.remainingFraction).toBe(0.7) + }) + }) + + it('returns 403 when project field is omitted', async () => { + await withServer(async (server) => { + server.enqueue({ + kind: 'quotaSummaryWindow', + managedProjectId: 'managed-rpc', + groups: [], + }) + const response = await fetch( + `${server.baseUrl}/v1internal:retrieveUserQuotaSummary`, + { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({}), + }, + ) + expect(response.status).toBe(403) + const body = (await response.json()) as { + error: { code: number; message: string } + } + expect(body.error.code).toBe(7) + expect(body.error.message).toBe('PERMISSION_DENIED') + }) + }) + + it('returns 403 when project does not match managedProjectId', async () => { + await withServer(async (server) => { + server.enqueue({ + kind: 'quotaSummaryWindow', + managedProjectId: 'managed-rpc', + groups: [], + }) + const response = await fetch( + `${server.baseUrl}/v1internal:retrieveUserQuotaSummary`, + { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ project: 'wrong-project' }), + }, + ) + expect(response.status).toBe(403) + }) + }) +}) diff --git a/packages/e2e-tests/src/mock-antigravity-server.ts b/packages/e2e-tests/src/mock-antigravity-server.ts index 76ce1d5..bc433df 100644 --- a/packages/e2e-tests/src/mock-antigravity-server.ts +++ b/packages/e2e-tests/src/mock-antigravity-server.ts @@ -350,7 +350,7 @@ export async function startMockAntigravityServer(): Promise { } catch { /* malformed — let the 200 path handle it */ } - if (postedProject && postedProject !== fixture.managedProjectId) { + if (!postedProject || postedProject !== fixture.managedProjectId) { sendJson(response, 403, { error: { code: 7, From 560219357ad2c7beaee26cb1a540c2617a6e8454 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Sat, 25 Jul 2026 08:00:41 +0200 Subject: [PATCH 23/25] fix(accounts): align removal index semantics with the core manager - persist 0 (not -1) when removing the current last account, matching the core's buildStorageSnapshot clamp; auth-doctor treats negative activeIndex as corruption - remove dead activeIndex() from CommandDataAccountManagerView interface and all adapters (replaced by getActiveIndexByFamily) - carry coolingDownUntil on CommandAccountRow so the sidebar refresher uses the row's own cooldown instead of matching by unstable numeric index - distinguish malformed quota-summary JSON bodies (400 INVALID_ARGUMENT) from the missing-project 403 PERMISSION_DENIED gate in the mock server --- .../src/mock-antigravity-server.test.ts | 24 +++++++ .../e2e-tests/src/mock-antigravity-server.ts | 14 +++- .../opencode/src/plugin/command-data.test.ts | 66 +++++++++++++++++- packages/opencode/src/plugin/command-data.ts | 24 ++++--- packages/opencode/src/plugin/commands.test.ts | 69 +++++++++++++++++++ packages/opencode/src/plugin/commands.ts | 22 +++--- packages/opencode/src/plugin/index.ts | 7 +- .../src/tui-compiled/plugin/command-data.ts | 24 ++++--- 8 files changed, 212 insertions(+), 38 deletions(-) diff --git a/packages/e2e-tests/src/mock-antigravity-server.test.ts b/packages/e2e-tests/src/mock-antigravity-server.test.ts index 6330f90..1fa429d 100644 --- a/packages/e2e-tests/src/mock-antigravity-server.test.ts +++ b/packages/e2e-tests/src/mock-antigravity-server.test.ts @@ -253,4 +253,28 @@ describe('quotaSummaryWindow managedProjectId 403 gate', () => { expect(response.status).toBe(403) }) }) + + it('returns 400 when the request body is unparseable JSON', async () => { + await withServer(async (server) => { + server.enqueue({ + kind: 'quotaSummaryWindow', + managedProjectId: 'managed-rpc', + groups: [], + }) + const response = await fetch( + `${server.baseUrl}/v1internal:retrieveUserQuotaSummary`, + { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: 'not-json', + }, + ) + expect(response.status).toBe(400) + const body = (await response.json()) as { + error: { code: number; message: string } + } + expect(body.error.code).toBe(3) + expect(body.error.message).toBe('INVALID_ARGUMENT') + }) + }) }) diff --git a/packages/e2e-tests/src/mock-antigravity-server.ts b/packages/e2e-tests/src/mock-antigravity-server.ts index bc433df..e84d3dd 100644 --- a/packages/e2e-tests/src/mock-antigravity-server.ts +++ b/packages/e2e-tests/src/mock-antigravity-server.ts @@ -343,12 +343,22 @@ export async function startMockAntigravityServer(): Promise { // on the fixture and the posted project does not match, return // PERMISSION_DENIED so the test exercises the fallback path. if (fixture.managedProjectId) { - let postedProject = '' + let postedProject: string | undefined try { postedProject = (JSON.parse(reqBody) as { project?: string }).project ?? '' } catch { - /* malformed — let the 200 path handle it */ + // Detect parse failure before the missing-project 403 gate + // so a serialisation bug does not masquerade as an auth + // fallback. + sendJson(response, 400, { + error: { + code: 3, + message: 'INVALID_ARGUMENT', + status: 'INVALID_ARGUMENT', + }, + }) + return } if (!postedProject || postedProject !== fixture.managedProjectId) { sendJson(response, 403, { diff --git a/packages/opencode/src/plugin/command-data.test.ts b/packages/opencode/src/plugin/command-data.test.ts index 0a75d8a..453607c 100644 --- a/packages/opencode/src/plugin/command-data.test.ts +++ b/packages/opencode/src/plugin/command-data.test.ts @@ -73,6 +73,7 @@ interface AccountFixture { cachedQuotaUpdatedAt?: number cachedQuotaAccountId?: string accountIneligible?: boolean + coolingDownUntil?: number } function makeAccountFixture( @@ -224,6 +225,7 @@ function makeHarness(options: { cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, cachedQuotaAccountId: entry.cachedQuotaAccountId, accountIneligible: entry.accountIneligible, + coolingDownUntil: entry.coolingDownUntil, })) }, getAccountsForQuotaCheck() { @@ -261,9 +263,6 @@ function makeHarness(options: { // storage adapter already mirrors the in-memory state on every // mutation, so a flush is a no-op for test purposes. }, - activeIndex() { - return liveCurrentIndex - }, getActiveIndexByFamily(): { claude: number; gemini: number } { return { claude: liveCurrentIndex, gemini: liveGeminiIndex } }, @@ -1035,6 +1034,67 @@ describe('createCommandDataService', () => { }) }) + it('removeAccount() keeps the cursor at the removed first slot when the current was the first account', async () => { + // Current is refresh-a (index 0). Removing index 0 (refresh-a) + // leaves index 0 still valid — it now points at refresh-b, the + // account that shifted in. The persisted activeIndex must mirror + // AccountManager.removeAccount(), which keeps the cursor at the + // removed slot when it still exists. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha' }), + makeAccountFixture({ refreshToken: 'refresh-b', label: 'Beta' }), + makeAccountFixture({ refreshToken: 'refresh-c', label: 'Gamma' }), + ], + activeIndex: 0, + }) + + await harness.service.removeAccount(0) + + expect(harness.storage.accounts.map((a) => a.refreshToken)).toEqual([ + 'refresh-b', + 'refresh-c', + ]) + // Slot 0 still exists in the new array → kept at 0. + expect(harness.storage.activeIndex).toBe(0) + expect(harness.storage.activeIndexByFamily).toEqual({ + claude: 0, + gemini: 0, + }) + }) + + it('removeAccount() persists index 0 when the current last account is removed, matching the core snapshot', async () => { + // Current is refresh-c (index 2, last account). Removing index 2 + // (the last slot) leaves no slot at that position in the new + // array. The in-memory AccountManager sets the sentinel to -1, + // but buildStorageSnapshot clamps it to 0 for disk (cf. + // account-manager.ts:1652). The persisted layer must follow that + // contract: persisted activeIndex must be >= 0 because + // auth-doctor.ts:149-156 treats a negative value as corruption. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha' }), + makeAccountFixture({ refreshToken: 'refresh-b', label: 'Beta' }), + makeAccountFixture({ refreshToken: 'refresh-c', label: 'Gamma' }), + ], + activeIndex: 2, + }) + + await harness.service.removeAccount(2) + + expect(harness.storage.accounts.map((a) => a.refreshToken)).toEqual([ + 'refresh-a', + 'refresh-b', + ]) + // The removed account was the last slot — no slot at that + // position remains. Persisted contract is non-negative. + expect(harness.storage.activeIndex).toBe(0) + expect(harness.storage.activeIndexByFamily).toEqual({ + claude: 0, + gemini: 0, + }) + }) + it('removeAccount() persists per-family active indexes independently when removing a non-current account', async () => { // Claude is current on refresh-a (index 0); Gemini is current on // refresh-c (index 2). Removing refresh-b (index 1) — a diff --git a/packages/opencode/src/plugin/command-data.ts b/packages/opencode/src/plugin/command-data.ts index 5170e82..e517ca4 100644 --- a/packages/opencode/src/plugin/command-data.ts +++ b/packages/opencode/src/plugin/command-data.ts @@ -137,6 +137,12 @@ export interface CommandAccountRow { enabled: boolean /** `true` when this row matches the harness-active account. */ current: boolean + /** + * Absolute timestamp (ms) until which this account is rate-limited. + * Carried on the row so the sidebar refresher can match cooldowns + * without relying on unstable numeric indices. + */ + coolingDownUntil?: number quota: Array<{ key: 'gemini' | 'non-gemini' label: string @@ -179,6 +185,7 @@ interface LiveAccountSnapshot { cachedQuotaUpdatedAt?: number cachedQuotaAccountId?: string accountIneligible?: boolean + coolingDownUntil?: number } function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { @@ -238,6 +245,7 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { label, enabled: entry.enabled, current: entry.active, + coolingDownUntil: entry.coolingDownUntil, quota, } } @@ -292,8 +300,6 @@ export function projectCommandAccountRows( * dialog-triggered mutation lands on disk before the dialog's response * returns. Without it, the dialog could toast "Account removed" while * the file still carries the old pool. - * - `activeIndex()` returns the position the harness considers active; - * the service uses it to mark the matching row as `current: true`. */ export interface CommandDataAccountManagerView { getAccounts(): LiveAccountSnapshot[] @@ -307,7 +313,6 @@ export interface CommandDataAccountManagerView { ): void requestSaveToDisk(): void flushSaveToDisk(): Promise - activeIndex(): number /** * Per-family live cursor. Used by the remove path so the persisted * `activeIndexByFamily` follows each model's currently-active account @@ -726,11 +731,14 @@ export function createCommandDataService( (account) => account.refreshToken === liveToken, ) if (found !== -1) return found - // The captured current token was the one removed. The live - // manager keeps the numeric cursor at the removed slot, so - // mirror that: clamp the removed index to the new account - // range. - return Math.max(0, Math.min(index, nextAccounts.length - 1)) + // The captured current token was the one removed. Mirror the + // core's in-memory semantics: keep the cursor at the removed + // slot when it still exists in the new array; when the removed + // slot was the last position and no longer exists, persist 0 + // (the core's buildStorageSnapshot clamps negative sentinels + // to 0, and auth-doctor treats a negative activeIndex as + // corruption — cf. auth-doctor.ts:149-156). + return index < nextAccounts.length ? index : 0 } const legacyClaude = current.activeIndexByFamily?.claude ?? current.activeIndex diff --git a/packages/opencode/src/plugin/commands.test.ts b/packages/opencode/src/plugin/commands.test.ts index 0e0d87e..22fcee1 100644 --- a/packages/opencode/src/plugin/commands.test.ts +++ b/packages/opencode/src/plugin/commands.test.ts @@ -7,6 +7,7 @@ import { AccountStorageUnreadableError } from '@cortexkit/antigravity-auth-core' import type { CommandModalName } from '../rpc/protocol' import { readSidebarState } from '../sidebar-state' import { registerAntigravityCommands } from './catalog' +import type { CommandAccountRow } from './command-data' import { applyCommand, buildDialogPayload, @@ -817,6 +818,74 @@ describe('applyCommand', () => { } }) + it('uses the dialog row cooldown over the live-by-index fallback when the row carries it', async () => { + // Live accounts have been renumbered between projection and + // refresh (e.g. a concurrent reload). The dialog row projected + // before the reorder carries its own cooldown, which must survive + // through the refresher — the live-by-index lookup at the row's + // stale index would give the wrong account's cooldown. + const sidebarFile = join(dir, 'sidebar-reorder.json') + const previousSidebarFile = process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = sidebarFile + try { + // Live accounts AFTER reorder: index 0 now has cooldown + // 1_800_000_000_000 (the wrong account's timer). + const getAccounts = mock(() => [ + { + index: 0, + label: 'Shifted account', + enabled: true, + coolingDownUntil: 1_800_000_000_000, + }, + { + index: 1, + label: 'Unlimited', + enabled: true, + coolingDownUntil: undefined, + }, + ]) + const refresher = createSidebarRefresher(getAccounts) + // Dialog rows projected BEFORE the reorder: index 0 already + // carries the original cooldown from the projection. + const dialogRows: CommandAccountRow[] = [ + { + id: 'acct-0', + index: 0, + label: 'Rate-limited', + enabled: true, + current: true, + coolingDownUntil: 1_900_000_000_000, + quota: [], + }, + { + id: 'acct-1', + index: 1, + label: 'Available', + enabled: true, + current: false, + coolingDownUntil: undefined, + quota: [], + }, + ] + + await refresher(dialogRows) + + const state = readSidebarState(sidebarFile) + expect(state.accounts).toHaveLength(2) + // The row's own cooldown (1_900_000_000_000) must be used, not + // the live-by-index fallback (1_800_000_000_000) which belongs + // to the account that shifted into slot 0 after the reorder. + expect(state.accounts[0]?.cooldownUntil).toBe(1_900_000_000_000) + expect(state.accounts[1]?.cooldownUntil).toBeUndefined() + } finally { + if (previousSidebarFile === undefined) { + delete process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + } else { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + it('starts a quota refresh when the quota dialog opens without delaying its cached payload', async () => { const listAccounts = mock(async () => [ { diff --git a/packages/opencode/src/plugin/commands.ts b/packages/opencode/src/plugin/commands.ts index 85cf86a..de8916a 100644 --- a/packages/opencode/src/plugin/commands.ts +++ b/packages/opencode/src/plugin/commands.ts @@ -754,11 +754,11 @@ export function createSidebarRefresher( }> | null, ): (accounts?: CommandAccountRow[]) => Promise { return async (dialogAccounts) => { - // Index → live coolingDownUntil. The dialog path rebuilds rows - // from the post-mutation command data, which does not carry the - // running cooldown; looking up the live account by index lets the - // sidebar show a rate-limited account as such instead of - // momentarily marking it available. + // Live cooldown lookup keyed by numeric index — a best-effort + // fallback for rows projected before coolingDownUntil was + // added to CommandAccountRow. Rows now carry their own + // cooldown from projection time; the live map only fills in + // when the row value is absent. const liveByIndex = new Map() try { const live = getAccounts() @@ -772,16 +772,16 @@ export function createSidebarRefresher( } const accounts = dialogAccounts ? dialogAccounts.map((entry) => { - const liveCooldown = liveByIndex.get(entry.index) + // Prefer the row's own cooldown — it was projected from the + // live view at the last projection and is stable regardless + // of concurrent reloads. Fall back to the live-by-index map + // for rows from older projections that lack the field. + const liveCooldown = + entry.coolingDownUntil ?? liveByIndex.get(entry.index) return { index: entry.index, label: entry.label, enabled: entry.enabled, - // Preserve the live cooldown when present so an apply - // (setCurrent, toggleEnabled, remove) cannot clear a - // rate-limited account's display. The dialog rows do not - // carry the running timer, so the live source is the - // only authoritative feed. coolingDownUntil: liveCooldown, cachedQuota: Object.fromEntries( entry.quota.flatMap((group) => { diff --git a/packages/opencode/src/plugin/index.ts b/packages/opencode/src/plugin/index.ts index 43c7d7a..d1f7aeb 100644 --- a/packages/opencode/src/plugin/index.ts +++ b/packages/opencode/src/plugin/index.ts @@ -229,6 +229,7 @@ export const createAntigravityPlugin = cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, cachedQuotaAccountId: entry.cachedQuotaAccountId, accountIneligible: entry.accountIneligible, + coolingDownUntil: entry.coolingDownUntil, })) }, getAccountsForQuotaCheck: () => { @@ -246,12 +247,6 @@ export const createAntigravityPlugin = flushSaveToDisk: async () => { await lifecycle.getAccountManager()?.flushSaveToDisk() }, - activeIndex: () => { - const manager = lifecycle.getAccountManager() - return manager - ? (manager.getCurrentAccountForFamily('claude')?.index ?? 0) - : 0 - }, getActiveIndexByFamily: () => { const manager = lifecycle.getAccountManager() if (!manager) return { claude: 0, gemini: 0 } diff --git a/packages/opencode/src/tui-compiled/plugin/command-data.ts b/packages/opencode/src/tui-compiled/plugin/command-data.ts index 5170e82..e517ca4 100644 --- a/packages/opencode/src/tui-compiled/plugin/command-data.ts +++ b/packages/opencode/src/tui-compiled/plugin/command-data.ts @@ -137,6 +137,12 @@ export interface CommandAccountRow { enabled: boolean /** `true` when this row matches the harness-active account. */ current: boolean + /** + * Absolute timestamp (ms) until which this account is rate-limited. + * Carried on the row so the sidebar refresher can match cooldowns + * without relying on unstable numeric indices. + */ + coolingDownUntil?: number quota: Array<{ key: 'gemini' | 'non-gemini' label: string @@ -179,6 +185,7 @@ interface LiveAccountSnapshot { cachedQuotaUpdatedAt?: number cachedQuotaAccountId?: string accountIneligible?: boolean + coolingDownUntil?: number } function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { @@ -238,6 +245,7 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { label, enabled: entry.enabled, current: entry.active, + coolingDownUntil: entry.coolingDownUntil, quota, } } @@ -292,8 +300,6 @@ export function projectCommandAccountRows( * dialog-triggered mutation lands on disk before the dialog's response * returns. Without it, the dialog could toast "Account removed" while * the file still carries the old pool. - * - `activeIndex()` returns the position the harness considers active; - * the service uses it to mark the matching row as `current: true`. */ export interface CommandDataAccountManagerView { getAccounts(): LiveAccountSnapshot[] @@ -307,7 +313,6 @@ export interface CommandDataAccountManagerView { ): void requestSaveToDisk(): void flushSaveToDisk(): Promise - activeIndex(): number /** * Per-family live cursor. Used by the remove path so the persisted * `activeIndexByFamily` follows each model's currently-active account @@ -726,11 +731,14 @@ export function createCommandDataService( (account) => account.refreshToken === liveToken, ) if (found !== -1) return found - // The captured current token was the one removed. The live - // manager keeps the numeric cursor at the removed slot, so - // mirror that: clamp the removed index to the new account - // range. - return Math.max(0, Math.min(index, nextAccounts.length - 1)) + // The captured current token was the one removed. Mirror the + // core's in-memory semantics: keep the cursor at the removed + // slot when it still exists in the new array; when the removed + // slot was the last position and no longer exists, persist 0 + // (the core's buildStorageSnapshot clamps negative sentinels + // to 0, and auth-doctor treats a negative activeIndex as + // corruption — cf. auth-doctor.ts:149-156). + return index < nextAccounts.length ? index : 0 } const legacyClaude = current.activeIndexByFamily?.claude ?? current.activeIndex From bfd519a297e7f002da71970f3626d58b424f14c8 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Sat, 25 Jul 2026 08:30:47 +0200 Subject: [PATCH 24/25] fix(accounts): distinguish an absent cooldown field from an empty one --- packages/opencode/src/plugin/commands.test.ts | 118 ++++++++++++++++++ packages/opencode/src/plugin/commands.ts | 7 +- 2 files changed, 123 insertions(+), 2 deletions(-) diff --git a/packages/opencode/src/plugin/commands.test.ts b/packages/opencode/src/plugin/commands.test.ts index 22fcee1..81ca4b3 100644 --- a/packages/opencode/src/plugin/commands.test.ts +++ b/packages/opencode/src/plugin/commands.test.ts @@ -886,6 +886,124 @@ describe('applyCommand', () => { } }) + it('does not fall through to the live-by-index map when the row carries its own undefined cooldown', async () => { + // After a concurrent reorder, the index that belonged to a + // no-cooldown account now holds a DIFFERENT account WITH a + // cooldown. The dialog row (projected before the reorder) + // carries coolingDownUntil: undefined (present = genuinely no + // cooldown) — this must NOT fall through to the index map and + // inherit the wrong account's timer. + const sidebarFile = join(dir, 'sidebar-misattrib.json') + const previousSidebarFile = process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = sidebarFile + try { + // Live accounts AFTER reorder: the account that shifted into + // slot 0 has a cooldown. + const getAccounts = mock(() => [ + { + index: 0, + label: 'Shifted-in account', + enabled: true, + coolingDownUntil: 1_700_000_000_000, + }, + { + index: 1, + label: 'Second account', + enabled: true, + coolingDownUntil: undefined, + }, + ]) + const refresher = createSidebarRefresher(getAccounts) + // Dialog rows projected BEFORE the reorder. Index 0 was the + // no-cooldown account and the row carries undefined (present). + const dialogRows: CommandAccountRow[] = [ + { + id: 'acct-0', + index: 0, + label: 'No-cooldown account', + enabled: true, + current: true, + coolingDownUntil: undefined, + quota: [], + }, + ] + + await refresher(dialogRows) + + const state = readSidebarState(sidebarFile) + expect(state.accounts).toHaveLength(1) + // Must NOT inherit 1_700_000_000_000 from the shifted-in + // account at live index 0. + expect(state.accounts[0]?.cooldownUntil).toBeUndefined() + } finally { + if (previousSidebarFile === undefined) { + delete process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + } else { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + + it('falls back to the live-by-index cooldown when the row lacks the property entirely', async () => { + // Legacy projection rows (pre-coolingDownUntil field) do not + // carry the key at all. The refresher must still read the live + // timer from the index map so a rate-limited account's cooldown + // survives through the sidebar update. + const sidebarFile = join(dir, 'sidebar-legacy.json') + const previousSidebarFile = process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = sidebarFile + try { + const getAccounts = mock(() => [ + { + index: 0, + label: 'Rate-limited', + enabled: true, + coolingDownUntil: 1_800_000_000_000, + }, + { + index: 1, + label: 'Available', + enabled: true, + coolingDownUntil: undefined, + }, + ]) + const refresher = createSidebarRefresher(getAccounts) + // Legacy dialog rows that DO NOT have coolingDownUntil at + // all (the property is absent, not undefined). + const dialogRows = [ + { + id: 'acct-0', + index: 0, + label: 'Rate-limited', + enabled: true, + current: true, + quota: [], + }, + { + id: 'acct-1', + index: 1, + label: 'Available', + enabled: true, + current: false, + quota: [], + }, + ] as CommandAccountRow[] + + await refresher(dialogRows) + + const state = readSidebarState(sidebarFile) + expect(state.accounts).toHaveLength(2) + expect(state.accounts[0]?.cooldownUntil).toBe(1_800_000_000_000) + expect(state.accounts[1]?.cooldownUntil).toBeUndefined() + } finally { + if (previousSidebarFile === undefined) { + delete process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + } else { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + it('starts a quota refresh when the quota dialog opens without delaying its cached payload', async () => { const listAccounts = mock(async () => [ { diff --git a/packages/opencode/src/plugin/commands.ts b/packages/opencode/src/plugin/commands.ts index de8916a..5f1cd3c 100644 --- a/packages/opencode/src/plugin/commands.ts +++ b/packages/opencode/src/plugin/commands.ts @@ -775,9 +775,12 @@ export function createSidebarRefresher( // Prefer the row's own cooldown — it was projected from the // live view at the last projection and is stable regardless // of concurrent reloads. Fall back to the live-by-index map - // for rows from older projections that lack the field. + // only for legacy rows that lack the field entirely (?? can't + // distinguish absent from present-but-undefined). const liveCooldown = - entry.coolingDownUntil ?? liveByIndex.get(entry.index) + 'coolingDownUntil' in entry + ? entry.coolingDownUntil + : liveByIndex.get(entry.index) return { index: entry.index, label: entry.label, From 5fddd9efcda444e53eae480e0eb91a44aedc6456 Mon Sep 17 00:00:00 2001 From: PR8 Fix Pass Date: Sat, 25 Jul 2026 21:40:21 +0200 Subject: [PATCH 25/25] feat(quota): background quota refresh for idle accounts Antigravity inference responses carry no quota headers, so quota is poll-only: today it refreshes only when a request finds a stale cache or when the quota dialog is opened. An account left idle keeps whatever numbers it had when it was last used, indefinitely. Adds a 5-minute jittered background refresh, bounded in three layers: the sidebar 'checkedAt' freshness gate (the cross-process bound -- N processes on one machine converge on ~one poll per interval), an advisory fenced file lock (an optimization, not the correctness mechanism), and fail-closed skip when the lock throws. There is deliberately no fallback mutex: a fallback for a failed lock is a second, worse lock. The poll reuses the existing refresh path end to end and publishes one sidebar snapshot per tick, with identity stamps and the live active-account flag intact. Also in this range, from review: - transport failures inside fetchGeminiCliQuota were swallowed by a blank catch and reported as a permanent-looking 'no CLI quota available'; they now propagate so a socket hang is distinguishable from an unconfigured CLI (a 403 still returns empty buckets) - pre-upgrade cachedQuota keys are normalized at read time -- killswitch and soft-quota read the same map, so a configured threshold sat inert until the first refresh - quota windows render shortest-first, sorted by duration - the sidebar publishes the real health score and active-account flag instead of constants - the quota dialog persists what it refreshes instead of discarding it - captured plan tier is published for external readers of the state file --- packages/core/src/account-manager.test.ts | 159 ++ packages/core/src/account-manager.ts | 86 +- packages/core/src/account-types.ts | 11 + packages/core/src/auth-types.ts | 6 + packages/core/src/project.test.ts | 49 + packages/core/src/project.ts | 42 +- packages/core/src/quota-manager.test.ts | 135 +- packages/core/src/quota-manager.ts | 93 +- packages/core/src/quota-types.ts | 64 +- packages/e2e-tests/src/harness.ts | 5 + .../e2e-tests/src/plugin-flow.e2e.test.ts | 88 + .../opencode/assets/antigravity.schema.json | 10 + .../src/plugin/account-command-oauth.ts | 25 +- .../opencode/src/plugin/auth-loader.test.ts | 7 + packages/opencode/src/plugin/auth-loader.ts | 40 +- .../plugin/background-quota-refresh.test.ts | 1981 +++++++++++++++++ .../src/plugin/background-quota-refresh.ts | 475 ++++ .../opencode/src/plugin/command-data.test.ts | 256 ++- packages/opencode/src/plugin/command-data.ts | 216 +- packages/opencode/src/plugin/commands.test.ts | 176 +- packages/opencode/src/plugin/commands.ts | 44 +- packages/opencode/src/plugin/config/schema.ts | 27 + packages/opencode/src/plugin/debug.test.ts | 17 +- packages/opencode/src/plugin/debug.ts | 24 + packages/opencode/src/plugin/index.ts | 88 +- .../opencode/src/plugin/lifecycle.test.ts | 111 +- packages/opencode/src/plugin/quota.test.ts | 246 +- packages/opencode/src/plugin/quota.ts | 132 +- packages/opencode/src/sidebar-state.test.ts | 271 ++- packages/opencode/src/sidebar-state.ts | 176 +- .../src/tui-compiled/plugin/command-data.ts | 216 +- .../src/tui-compiled/sidebar-state.ts | 176 +- packages/opencode/src/tui-compiled/tui.tsx | 9 +- .../opencode/src/tui-windows-frames.test.tsx | 33 +- packages/opencode/src/tui.test.tsx | 16 +- packages/opencode/src/tui.tsx | 15 +- test/environment.test.ts | 232 ++ test/setup.ts | 23 +- 38 files changed, 5483 insertions(+), 297 deletions(-) create mode 100644 packages/opencode/src/plugin/background-quota-refresh.test.ts create mode 100644 packages/opencode/src/plugin/background-quota-refresh.ts diff --git a/packages/core/src/account-manager.test.ts b/packages/core/src/account-manager.test.ts index 62f412f..3146ed9 100644 --- a/packages/core/src/account-manager.test.ts +++ b/packages/core/src/account-manager.test.ts @@ -54,6 +54,49 @@ describe('core AccountManager', () => { ).toEqual(['r1', 'r2', 'r3']) }) + it('normalizes persisted legacy quota keys for soft-quota and proactive-rotation reads', () => { + const now = 1_700_000_000_000 + const legacy: AccountStorageV4 = { + version: 4, + accounts: [ + { + refreshToken: 'legacy-token', + addedAt: 1, + lastUsed: 0, + cachedQuota: { + claude: { remainingFraction: 0.4, modelCount: 1 }, + }, + cachedQuotaUpdatedAt: now, + }, + { + refreshToken: 'other-token', + addedAt: 1, + lastUsed: 0, + }, + ], + activeIndex: 0, + } + const memory = createStore(legacy) + const manager = new AccountManager(undefined, legacy, { + store: memory.store, + now: () => now, + }) + const account = manager.getAccounts()[0]! + + expect( + manager.isAccountOverSoftQuota( + account, + 'claude', + 50, + 60_000, + 'claude-sonnet', + ), + ).toBe(true) + expect( + manager.shouldProactivelyRotate('claude', 'claude-sonnet', 50, 60_000), + ).toBe(true) + }) + it.each([ 'sticky', 'round-robin', @@ -69,6 +112,83 @@ describe('core AccountManager', () => { ).not.toBeNull() }) + it('hybrid skips the active Gemini account limited on the antigravity header style', () => { + const now = 1_700_000_000_000 + const hybridStored: AccountStorageV4 = { + version: 4, + accounts: [ + { refreshToken: 'r1', projectId: 'p1', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r2', projectId: 'p2', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r3', projectId: 'p3', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r4', projectId: 'p4', addedAt: 1, lastUsed: 0 }, + ], + activeIndex: 1, + activeIndexByFamily: { gemini: 1 }, + } + const memory = createStore(hybridStored) + const manager = new AccountManager(undefined, hybridStored, { + store: memory.store, + now: () => now, + random: () => 0.5, + }) + const limited = manager.getAccounts()[1]! + manager.markRateLimitedWithReason( + limited, + 'gemini', + 'antigravity', + 'antigravity-gemini-3.6-flash', + 'RATE_LIMIT_EXCEEDED', + ) + + const selected = manager.getCurrentOrNextForFamily( + 'gemini', + 'antigravity-gemini-3.6-flash', + 'hybrid', + 'antigravity', + ) + + expect(selected?.index).toBe(0) + }) + + it('hybrid returns null when every Gemini account is limited on the antigravity header style', () => { + const now = 1_700_000_000_000 + const hybridStored: AccountStorageV4 = { + version: 4, + accounts: [ + { refreshToken: 'r1', projectId: 'p1', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r2', projectId: 'p2', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r3', projectId: 'p3', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r4', projectId: 'p4', addedAt: 1, lastUsed: 0 }, + ], + activeIndex: 1, + activeIndexByFamily: { gemini: 1 }, + } + const memory = createStore(hybridStored) + const manager = new AccountManager(undefined, hybridStored, { + store: memory.store, + now: () => now, + random: () => 0.5, + }) + for (const account of manager.getAccounts()) { + manager.markRateLimitedWithReason( + account, + 'gemini', + 'antigravity', + 'antigravity-gemini-3.6-flash', + 'RATE_LIMIT_EXCEEDED', + ) + } + + const selected = manager.getCurrentOrNextForFamily( + 'gemini', + 'antigravity-gemini-3.6-flash', + 'hybrid', + 'antigravity', + ) + + expect(selected).toBeNull() + }) + it('tracks model-specific limits independently', () => { let now = 1_000 const memory = createStore(stored) @@ -219,6 +339,45 @@ describe('core AccountManager', () => { ) }) + it('persists and restores the captured tier schema marker across save→loadFromDisk', async () => { + const seeded = { + version: 4, + accounts: [ + { + refreshToken: 'r1', + projectId: 'p1', + addedAt: 1, + lastUsed: 0, + capturedTierId: 'free-tier', + capturedTierAt: 1_700_000_000_000, + capturedTierSchemaVersion: 1, + }, + ], + activeIndex: 0, + } as AccountStorageV4 & { + accounts: Array<{ capturedTierSchemaVersion?: number }> + } + const memory = createStore(seeded) + const manager = new AccountManager(undefined, seeded, { + store: memory.store, + }) + + await manager.saveToDiskReplace() + + expect(memory.state()?.accounts[0]).toMatchObject({ + capturedTierSchemaVersion: 1, + }) + const reloaded = new AccountManager( + undefined, + memory.state() ?? undefined, + { store: memory.store }, + ) + expect( + (reloaded.getAccounts()[0] as { capturedTierSchemaVersion?: number }) + ?.capturedTierSchemaVersion, + ).toBe(1) + }) + it('drops the quota write when the refresh token captured at refresh time is gone (remove-during-refresh race)', () => { // Race: an async quota refresh is in flight for account A while the // user removes account A from the pool. When the refresh resolves, diff --git a/packages/core/src/account-manager.ts b/packages/core/src/account-manager.ts index b637289..9be86b4 100644 --- a/packages/core/src/account-manager.ts +++ b/packages/core/src/account-manager.ts @@ -20,7 +20,11 @@ import { updateFingerprintVersion, } from './fingerprint.ts' import { getQuotaGroupForModel } from './model-registry.ts' -import type { QuotaGroup, QuotaGroupSummary } from './quota-types.ts' +import { + normalizeLegacyCachedQuota, + type QuotaGroup, + type QuotaGroupSummary, +} from './quota-types.ts' import { type AccountWithMetrics, getHealthTracker, @@ -100,6 +104,17 @@ export interface ManagedAccount { /** Opaque identity of the refresh token that produced `cachedQuota`. */ cachedQuotaAccountId?: string cachedQuotaUpdatedAt?: number + /** + * Captured plan tier ID from the most recent `loadCodeAssist` response. + * Raw upstream string (e.g. `"free-tier"`) — never normalised. + */ + capturedTierId?: string + /** Raw paid-tier ID from the most recent `loadCodeAssist` response. */ + capturedPaidTierId?: string + /** Epoch ms when `capturedTierId` was last recorded. */ + capturedTierAt?: number + /** Schema version of the most recent tier capture, even when paid tier is absent. */ + capturedTierSchemaVersion?: number verificationRequired?: boolean verificationRequiredAt?: number verificationRequiredReason?: string @@ -415,14 +430,16 @@ export class AccountManager { touchedForQuota: {}, fingerprint: acc.fingerprint ?? generateFingerprint(), fingerprintHistory: acc.fingerprintHistory ?? [], - cachedQuota: acc.cachedQuota as - | Partial> - | undefined, + cachedQuota: normalizeLegacyCachedQuota(acc.cachedQuota), // Restore the opaque identity stamp alongside the quota so the // post-load projection can detect a stale snapshot captured // for a different account after an index shift. cachedQuotaAccountId: acc.cachedQuotaAccountId, cachedQuotaUpdatedAt: acc.cachedQuotaUpdatedAt, + capturedTierId: acc.capturedTierId, + capturedPaidTierId: acc.capturedPaidTierId, + capturedTierAt: acc.capturedTierAt, + capturedTierSchemaVersion: acc.capturedTierSchemaVersion, dailyRequestCounts: acc.dailyRequestCounts, verificationRequired: acc.verificationRequired, verificationRequiredAt: acc.verificationRequiredAt, @@ -830,7 +847,13 @@ export class AccountManager { lastUsed: acc.lastUsed, healthScore: healthTracker.getScore(acc.index), isRateLimited: - isRateLimitedForFamily(acc, family, this.now, model) || + isRateLimitedForHeaderStyle( + acc, + family, + headerStyle, + this.now, + model, + ) || isOverSoftQuotaThreshold( acc, family, @@ -1680,6 +1703,10 @@ export class AccountManager { // for a different account after an index shift. cachedQuotaAccountId: a.cachedQuotaAccountId, cachedQuotaUpdatedAt: a.cachedQuotaUpdatedAt, + capturedTierId: a.capturedTierId, + capturedPaidTierId: a.capturedPaidTierId, + capturedTierAt: a.capturedTierAt, + capturedTierSchemaVersion: a.capturedTierSchemaVersion, dailyRequestCounts: a.dailyRequestCounts, verificationRequired: a.verificationRequired, verificationRequiredAt: a.verificationRequiredAt, @@ -1905,6 +1932,55 @@ export class AccountManager { account.cachedQuotaUpdatedAt = this.now() } + /** + * Apply a subset of fields from a quota-fetch `updatedAccount` result onto + * the live in-memory record for the given index. Only patches fields that + * are present and non-empty in `patch` to avoid overwriting valid state + * with stale or missing values. + * + * Identity guard: if `expectedRefreshToken` is provided and the account at + * `accountIndex` no longer carries that token (concurrent reorder/replace), + * the patch is silently dropped. + * + * `managedProjectId` is intentionally absent from the patch type: the only + * caller (`BackgroundQuotaRefresh`) routes through `PollerAccountView` which + * exposes only `capturedTierId`/`capturedTierAt`; project-context updates + * happen via `ensureProjectContext`, not through this method. + */ + applyUpdatedAccount( + accountIndex: number, + patch: Partial< + Pick< + AccountMetadataV3, + | 'capturedTierId' + | 'capturedPaidTierId' + | 'capturedTierAt' + | 'capturedTierSchemaVersion' + > + >, + expectedRefreshToken?: string, + ): void { + const account = this.accounts[accountIndex] + if ( + !account || + (expectedRefreshToken !== undefined && + account.parts.refreshToken !== expectedRefreshToken) + ) + return + if (patch.capturedTierId !== undefined) { + account.capturedTierId = patch.capturedTierId + } + if (patch.capturedPaidTierId !== undefined) { + account.capturedPaidTierId = patch.capturedPaidTierId + } + if (patch.capturedTierAt !== undefined) { + account.capturedTierAt = patch.capturedTierAt + } + if (patch.capturedTierSchemaVersion !== undefined) { + account.capturedTierSchemaVersion = patch.capturedTierSchemaVersion + } + } + /** * Record a successful API request for an account. * Tracks per model family with daily reset. diff --git a/packages/core/src/account-types.ts b/packages/core/src/account-types.ts index b85c510..47ae022 100644 --- a/packages/core/src/account-types.ts +++ b/packages/core/src/account-types.ts @@ -90,6 +90,17 @@ export interface AccountMetadataV3 { resetTime?: string }[] cachedQuotaUpdatedAt?: number + /** + * Captured plan tier ID from the most recent `loadCodeAssist` response. + * Raw upstream string (e.g. `"free-tier"`) — never normalised. + */ + capturedTierId?: string + /** Raw paid-tier ID from the most recent `loadCodeAssist` response. */ + capturedPaidTierId?: string + /** Epoch ms when `capturedTierId` was last recorded. */ + capturedTierAt?: number + /** Schema version of the most recent tier capture, even when paid tier is absent. */ + capturedTierSchemaVersion?: number /** Daily request counts per model family, resets when date changes */ dailyRequestCounts?: { date: string diff --git a/packages/core/src/auth-types.ts b/packages/core/src/auth-types.ts index 10b2b90..a60e40f 100644 --- a/packages/core/src/auth-types.ts +++ b/packages/core/src/auth-types.ts @@ -33,4 +33,10 @@ export interface RefreshParts { export interface ProjectContextResult { auth: OAuthAuthDetails effectiveProjectId: string + /** + * Plan tier captured from the `loadCodeAssist` payload at project-context + * resolution time. Present only when the upstream returned a non-empty + * `currentTier.id`; absent when the payload lacked tier info. + */ + capturedTier?: { id: string; paidId?: string; capturedAt: number } } diff --git a/packages/core/src/project.test.ts b/packages/core/src/project.test.ts index 9f1224c..bd6fdc1 100644 --- a/packages/core/src/project.test.ts +++ b/packages/core/src/project.test.ts @@ -116,6 +116,55 @@ describe('project bootstrap', () => { expect(fetchSpy).toHaveBeenCalledTimes(1) }) + it('Fix2-tier-capture: capturedTier is returned from ensureProjectContext when loadCodeAssist returns currentTier', async () => { + const capturedNow = 1_785_000_000_000 + spyOn(Date, 'now').mockImplementation(() => capturedNow) + const fetchSpy = mock().mockResolvedValue( + mockResponse({ + cloudaicompanionProject: { id: 'my-project' }, + currentTier: { id: 'pro-tier' }, + paidTier: { id: 'g1-pro-tier' }, + }), + ) + const { fetchWithAgyCliTransport } = await import('./agy-transport.ts') + ;(fetchWithAgyCliTransport as any).mockImplementation(fetchSpy) + + const auth = { + type: 'oauth' as const, + access: 'access-token', + refresh: 'bare-token', + expires: Date.now() + 60_000, + } + + const result = await ensureProjectContext(auth) + + expect(result.capturedTier).toEqual({ + id: 'pro-tier', + paidId: 'g1-pro-tier', + capturedAt: capturedNow, + }) + }) + + it('Fix2-tier-absent: capturedTier is absent when loadCodeAssist returns no currentTier', async () => { + const fetchSpy = mock().mockResolvedValue( + mockResponse({ cloudaicompanionProject: { id: 'my-project' } }), + ) + const { fetchWithAgyCliTransport } = await import('./agy-transport.ts') + ;(fetchWithAgyCliTransport as any).mockImplementation(fetchSpy) + + const auth = { + type: 'oauth' as const, + access: 'access-token', + refresh: 'bare-token-2', + expires: Date.now() + 60_000, + } + + const result = await ensureProjectContext(auth) + + // No tier in payload — must be absent, not defaulted. + expect(result.capturedTier).toBeUndefined() + }) + it('does not retry managed-project provisioning after a cached failure expires', async () => { let now = 1_000 spyOn(Date, 'now').mockImplementation(() => now) diff --git a/packages/core/src/project.ts b/packages/core/src/project.ts index f2645db..02a03be 100644 --- a/packages/core/src/project.ts +++ b/packages/core/src/project.ts @@ -40,6 +40,11 @@ interface LoadCodeAssistPayload { currentTier?: { id?: string } + paidTier?: + | string + | { + id?: string + } allowedTiers?: AntigravityUserTier[] } @@ -295,6 +300,7 @@ export async function ensureProjectContext( const persistManagedProject = async ( managedProjectId: string, + capturedTier?: ProjectContextResult['capturedTier'], ): Promise => { const updatedAuth: OAuthAuthDetails = { ...auth, @@ -305,7 +311,11 @@ export async function ensureProjectContext( }), } - return { auth: updatedAuth, effectiveProjectId: managedProjectId } + return { + auth: updatedAuth, + effectiveProjectId: managedProjectId, + capturedTier, + } } // Try to resolve a managed project from Antigravity if possible. @@ -313,10 +323,24 @@ export async function ensureProjectContext( accessToken, parts.projectId ?? fallbackProjectId, ) + // Capture tier from the loadCodeAssist payload. The raw id is stored + // as-is; absent payload or missing tier leaves capturedTier undefined. + const capturedTierId = loadPayload?.currentTier?.id + const paidTierId = + typeof loadPayload?.paidTier === 'string' + ? loadPayload.paidTier + : loadPayload?.paidTier?.id + const tierFromPayload: ProjectContextResult['capturedTier'] = capturedTierId + ? { + id: capturedTierId, + ...(paidTierId ? { paidId: paidTierId } : {}), + capturedAt: Date.now(), + } + : undefined const resolvedManagedProjectId = extractManagedProjectId(loadPayload) if (resolvedManagedProjectId) { - return persistManagedProject(resolvedManagedProjectId) + return persistManagedProject(resolvedManagedProjectId, tierFromPayload) } // No managed project found - try to auto-provision one via onboarding. @@ -337,7 +361,7 @@ export async function ensureProjectContext( log.debug('Successfully provisioned managed project', { provisionedProjectId, }) - return persistManagedProject(provisionedProjectId) + return persistManagedProject(provisionedProjectId, tierFromPayload) } log.warn( @@ -352,11 +376,19 @@ export async function ensureProjectContext( } if (parts.projectId) { - return { auth, effectiveProjectId: parts.projectId } + return { + auth, + effectiveProjectId: parts.projectId, + capturedTier: tierFromPayload, + } } // No project id present in auth; fall back to the hardcoded id for requests. - return { auth, effectiveProjectId: fallbackProjectId } + return { + auth, + effectiveProjectId: fallbackProjectId, + capturedTier: tierFromPayload, + } } if (!cacheKey) { diff --git a/packages/core/src/quota-manager.test.ts b/packages/core/src/quota-manager.test.ts index 13213c4..04efe4f 100644 --- a/packages/core/src/quota-manager.test.ts +++ b/packages/core/src/quota-manager.test.ts @@ -7,6 +7,7 @@ import { createQuotaManager, type FetchAccountQuota, fetchQuotaSummary, + type RetrieveUserQuotaSummaryBucket, type RetrieveUserQuotaSummaryResponse, } from './quota-manager.ts' import type { AccountQuotaResult } from './quota-types.ts' @@ -805,18 +806,18 @@ describe('aggregateQuotaSummary', () => { const gemini = summary.groups.gemini! expect(gemini.windows).toHaveLength(2) - // weekly first - expect(gemini.windows![0]!.window).toBe('weekly') - expect(gemini.windows![0]!.remainingFraction).toBeCloseTo(0.9214, 3) - expect(gemini.windows![1]!.window).toBe('5h') - expect(gemini.windows![1]!.remainingFraction).toBeCloseTo(0.9886, 3) + // shortest-first: 5h before weekly + expect(gemini.windows![0]!.window).toBe('5h') + expect(gemini.windows![0]!.remainingFraction).toBeCloseTo(0.9886, 3) + expect(gemini.windows![1]!.window).toBe('weekly') + expect(gemini.windows![1]!.remainingFraction).toBeCloseTo(0.9214, 3) const nonGemini = summary.groups['non-gemini']! expect(nonGemini.windows).toHaveLength(2) - expect(nonGemini.windows![0]!.window).toBe('weekly') - expect(nonGemini.windows![0]!.remainingFraction).toBeCloseTo(0.9852, 3) - expect(nonGemini.windows![1]!.window).toBe('5h') - expect(nonGemini.windows![1]!.remainingFraction).toBeCloseTo(0.9556, 3) + expect(nonGemini.windows![0]!.window).toBe('5h') + expect(nonGemini.windows![0]!.remainingFraction).toBeCloseTo(0.9556, 3) + expect(nonGemini.windows![1]!.window).toBe('weekly') + expect(nonGemini.windows![1]!.remainingFraction).toBeCloseTo(0.9852, 3) }) it('maps Free (weekly-only) groups to pools', () => { @@ -851,25 +852,56 @@ describe('aggregateQuotaSummary', () => { expect(summary.groups['non-gemini']!.resetTime).toBe('2026-07-24T18:41:52Z') }) - it('preserves window order: weekly first, then 5h', () => { + it('preserves window order: shortest first', () => { const response = loadFixture('pro-ruqs.json') const summary = aggregateQuotaSummary(response) for (const group of Object.values(summary.groups)) { if (!group?.windows) continue - for (let i = 1; i < group.windows.length; i++) { - const prev = group.windows[i - 1]! - const curr = group.windows[i]! - // All weeklies come before any 5h - if (prev.window === 'weekly' && curr.window === '5h') continue - if (prev.window === 'weekly' && curr.window === 'weekly') continue - if (prev.window === '5h' && curr.window === '5h') continue - // A 5h before a weekly = wrong order - expect(prev.window).not.toBe('5h') - } + expect(group.windows).toHaveLength(2) + expect(group.windows![0]!.window).toBe('5h') + expect(group.windows![1]!.window).toBe('weekly') } }) + it('sorts unknown window kinds last, deterministically', () => { + const buckets: RetrieveUserQuotaSummaryBucket[] = [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly', + window: 'weekly', + resetTime: '2026-07-28T00:00:00Z', + remainingFraction: 0.9, + }, + { + bucketId: 'gemini-daily', + displayName: 'Daily', + window: 'daily', + resetTime: '2026-07-25T00:00:00Z', + remainingFraction: 0.5, + }, + { + bucketId: 'gemini-5h', + displayName: '5h', + window: '5h', + resetTime: '2026-07-24T12:00:00Z', + remainingFraction: 0.8, + }, + ] + const summary = aggregateQuotaSummary({ + groups: [ + { displayName: 'Gemini', buckets, description: 'Models: A, B, C' }, + ], + }) + const windows = summary.groups.gemini!.windows! + expect(windows).toHaveLength(3) + // Known windows sorted shortest-first: 5h → weekly + expect(windows[0]!.window).toBe('5h') + expect(windows[1]!.window).toBe('weekly') + // Unknown window ('daily') sorts last + expect(windows[2]!.window).toBe('daily') + }) + it('counts models from the description minus the prefix label', () => { const response: RetrieveUserQuotaSummaryResponse = { groups: [ @@ -1176,4 +1208,67 @@ describe('fetchQuotaSummary', () => { ) expect(result.summary.groups[0]!.buckets[0]!.remainingFraction).toBe(0.42) }) + + it('N1: a transient 500 on the managed attempt does NOT fall through to the projectId fallback', async () => { + // N1 regression: before the fix, tryBody returned null for ANY failure + // (403, 500, network). The fallback was entered on all of them, which + // could return a different project\'s quota data on transient errors. + // Now the fallback is gated on 403 only. + const triedProjects: string[] = [] + const fetchVia = async ( + _url: string, + init: RequestInit, + ): Promise => { + const body = JSON.parse((init as any).body ?? '{}') + triedProjects.push(body.project as string) + // Return 500 for the managed project attempt (transient error). + return new Response('internal error', { status: 500 }) + } + // Both managed and regular project IDs are provided. The 500 on the + // managed attempt must NOT trigger the projectId fallback. + await expect( + fetchQuotaSummary({ + accessToken: 'tok', + managedProjectId: 'managed-proj', + projectId: 'regular-proj', + endpoints: ENDPOINTS, + fetchVia: fetchVia as any, + }), + ).rejects.toThrow() + // Only the managed project should have been tried — the regular-proj + // fallback must not run on a transient error. + expect(triedProjects).toHaveLength(1) + expect(triedProjects[0]).toBe('managed-proj') + }) + + it('N1: a 500 on the managed attempt does not clobber the 403-fallback path', async () => { + // Complement: a 403 on the managed project DOES enter the fallback. + // This ensures the gating change did not accidentally disable the + // existing 403-fallback behavior. + const triedProjects: string[] = [] + const summary: RetrieveUserQuotaSummaryResponse = { groups: [] } + const fetchVia = async ( + _url: string, + init: RequestInit, + ): Promise => { + const body = JSON.parse((init as any).body ?? '{}') + triedProjects.push(body.project as string) + if (body.project === 'managed-proj') { + return new Response('{}', { status: 403 }) + } + return new Response(JSON.stringify(summary), { status: 200 }) + } + const result = await fetchQuotaSummary({ + accessToken: 'tok', + managedProjectId: 'managed-proj', + projectId: 'regular-proj', + endpoints: ENDPOINTS, + fetchVia: fetchVia as any, + }) + // The 403 on managed-proj must trigger the fallback to regular-proj. + expect(triedProjects).toHaveLength(2) + expect(triedProjects[0]).toBe('managed-proj') + expect(triedProjects[1]).toBe('regular-proj') + expect(result.summary).toEqual(summary) + }) }) diff --git a/packages/core/src/quota-manager.ts b/packages/core/src/quota-manager.ts index 366b211..5d7f864 100644 --- a/packages/core/src/quota-manager.ts +++ b/packages/core/src/quota-manager.ts @@ -631,11 +631,27 @@ function parseDescriptionModelCount(description: string): number { return entries.length } +/** + * Map a quota window kind to its duration in milliseconds for + * deterministic sort order. Unknown windows get MAX_SAFE_INTEGER so + * they sort last without random reordering across runs. + */ +function windowDurationMs(window: 'weekly' | '5h' | string): number { + switch (window) { + case '5h': + return 5 * 60 * 60 * 1000 + case 'weekly': + return 7 * 24 * 60 * 60 * 1000 + default: + return Number.MAX_SAFE_INTEGER + } +} + /** * Aggregate a retrieveUserQuotaSummary response into a QuotaSummary. * * Each RUQS group maps to a pool via bucketId prefix. Within a pool, - * windows are stored ordered (weekly first, then 5h). The pool's + * windows are stored shortest-first (5h before weekly, etc.). The pool's * `remainingFraction`/`resetTime` derive from the most-constrained window. */ export function aggregateQuotaSummary( @@ -657,11 +673,10 @@ export function aggregateQuotaSummary( } if (windows.length === 0) continue - // Order: weekly first, then 5h. - windows.sort((a, b) => { - const order: Record = { weekly: 0, '5h': 1 } - return (order[a.window] ?? 2) - (order[b.window] ?? 2) - }) + // Order: shortest window first so the binding 5h window leads visually. + windows.sort( + (a, b) => windowDurationMs(a.window) - windowDurationMs(b.window), + ) const constrained = mostConstrainedWindow(windows) // Pick the first RECOGNIZED bucket for pool derivation. Older @@ -741,10 +756,18 @@ export async function fetchQuotaSummary( throw new Error('No endpoints configured for fetchQuotaSummary') } + // Discriminated result so the fallback is only entered on a 403, not on + // transient errors (429, 5xx, network). A transient failure on the managed + // project ID must not fall through to the regular projectId — that could + // return a DIFFERENT project's quota data instead of signalling a retry. + type TryBodyResult = + | { ok: true; summary: RetrieveUserQuotaSummaryResponse } + | { ok: false; reason: '403' | 'transient' } + const tryBody = async ( endpoint: string, projectId: string, - ): Promise => { + ): Promise => { const body = { project: projectId } try { const response = await transport( @@ -763,7 +786,10 @@ export async function fetchQuotaSummary( ) if (response.ok) { - return (await response.json()) as RetrieveUserQuotaSummaryResponse + return { + ok: true, + summary: (await response.json()) as RetrieveUserQuotaSummaryResponse, + } } const status = response.status @@ -771,7 +797,7 @@ export async function fetchQuotaSummary( errors.push( `retrieveUserQuotaSummary 403 at ${endpoint} (project=${projectId.slice(0, 12)}…)`, ) - return null + return { ok: false, reason: '403' } } if (status === 429 || status >= 500) { @@ -781,19 +807,19 @@ export async function fetchQuotaSummary( ) // Endpoint failover: the caller may have multiple endpoints; // continue to the next one with the same project ID. - return null + return { ok: false, reason: 'transient' } } const message = await response.text().catch(() => '') errors.push( `retrieveUserQuotaSummary ${status} at ${endpoint}${message ? `: ${message.trim().slice(0, 200)}` : ''}`, ) - return null + return { ok: false, reason: 'transient' } } catch (error) { errors.push( `retrieveUserQuotaSummary network error at ${endpoint}: ${error instanceof Error ? error.message : String(error)}`, ) - return null + return { ok: false, reason: 'transient' } } } @@ -801,16 +827,21 @@ export async function fetchQuotaSummary( // For each project, iterate the endpoint list so a 500/429 on the // primary endpoint falls through to the next entry. const primary = options.managedProjectId ?? options.projectId + let primaryGot403 = false if (primary) { for (const endpoint of options.endpoints) { const result = await tryBody(endpoint, primary) - if (result) return { summary: result } + if (result.ok) return { summary: result.summary } + if (result.reason === '403') primaryGot403 = true } } - // If primary failed with 403 and we used managedProjectId, - // retry with regular projectId as fallback (only when distinct). + // Retry with the regular projectId ONLY when the managed-project attempt + // got a 403 (i.e. the managed project does not own the user). A transient + // failure (429, 5xx, network) must NOT fall through here — the regular + // projectId could return a different project's quota data. const fallbackId = + primaryGot403 && options.managedProjectId && options.projectId && options.managedProjectId !== options.projectId @@ -819,7 +850,7 @@ export async function fetchQuotaSummary( if (fallbackId) { for (const endpoint of options.endpoints) { const result = await tryBody(endpoint, fallbackId) - if (result) return { summary: result } + if (result.ok) return { summary: result.summary } } } @@ -979,6 +1010,15 @@ export interface FetchGeminiCliQuotaOptions { endpoints: readonly string[] timeoutMs?: number userAgent?: string + /** + * Optional transport override. Production callers omit this; the e2e + * harness and tests inject a stub. Mirrors `fetchQuotaSummary`'s seam. + */ + fetchVia?: ( + url: string, + options: RequestInit, + extra: { timeoutMs: number }, + ) => Promise } export async function fetchGeminiCliQuota( @@ -986,11 +1026,13 @@ export async function fetchGeminiCliQuota( ): Promise { const timeoutMs = options.timeoutMs ?? QUOTA_MANAGER_DEFAULT_TIMEOUT_MS const userAgent = options.userAgent ?? buildAntigravityHarnessUserAgent() + const transport = options.fetchVia ?? defaultTransport + const errors: string[] = [] for (const endpoint of options.endpoints) { const body = options.projectId ? { project: options.projectId } : {} try { - const response = await defaultTransport( + const response = await transport( `${endpoint}/v1internal:retrieveUserQuota`, { method: 'POST', @@ -1010,10 +1052,25 @@ export async function fetchGeminiCliQuota( const status = response.status if (status === 429 || status >= 500) { + errors.push(`fetchGeminiCliQuota ${status} at ${endpoint}`) continue } + // Non-retryable server response (e.g. 403) — treat as no CLI quota. return { buckets: [] } - } catch {} + } catch (error) { + // Transport-level failure (network abort, DNS, timeout). Collect the + // error and try the next endpoint; if all endpoints fail, throw so + // the caller can distinguish a transient network failure from an + // authoritative 'no CLI quota configured' response. + errors.push(error instanceof Error ? error.message : String(error)) + } + } + + // All endpoints produced transport-level errors — propagate so the + // outer caller (quota.ts .catch) can surface the real failure reason + // rather than the generic 'No Gemini CLI quota available' message. + if (errors.length > 0) { + throw new Error(errors.join('; ') || 'fetchGeminiCliQuota failed') } return { buckets: [] } diff --git a/packages/core/src/quota-types.ts b/packages/core/src/quota-types.ts index c1b3adc..0e7df08 100644 --- a/packages/core/src/quota-types.ts +++ b/packages/core/src/quota-types.ts @@ -27,12 +27,74 @@ export interface QuotaGroupSummary { modelCount: number /** * Per-window breakdown from the retrieveUserQuotaSummary response. - * `weekly` first, then `5h`. Omitted in legacy cached shapes — + * shortest-first (5h before weekly, etc.). Omitted in legacy cached shapes — * consumers treat a single remainingFraction as one unlabeled window. */ windows?: QuotaWindowEntry[] } +/** + * Migrate pre-pool quota keys while loading persisted account data. Current + * keys pass through unchanged, making this safe for every storage read. + */ +export function normalizeLegacyCachedQuota( + raw: AccountMetadataV3['cachedQuota'], +): Partial> | undefined { + if (!raw) return raw + + const hasLegacy = + 'gemini-pro' in raw || + 'gemini-flash' in raw || + 'claude' in raw || + 'gpt-oss' in raw + if (!hasLegacy) return raw + + const earlierResetTime = ( + a: string | undefined, + b: string | undefined, + ): string | undefined => { + if (!a) return b + if (!b) return a + return a < b ? a : b + } + + const minFraction = ( + a: QuotaGroupSummary | undefined, + b: QuotaGroupSummary | undefined, + ): QuotaGroupSummary | undefined => { + if (!a) return b + if (!b) return a + const fa = a.remainingFraction ?? 1 + const fb = b.remainingFraction ?? 1 + const winner = fa <= fb ? a : b + const loser = fa <= fb ? b : a + return { + ...winner, + resetTime: earlierResetTime(winner.resetTime, loser.resetTime), + } + } + + const gemini = minFraction( + raw.gemini as QuotaGroupSummary | undefined, + minFraction( + raw['gemini-pro'] as QuotaGroupSummary | undefined, + raw['gemini-flash'] as QuotaGroupSummary | undefined, + ), + ) + const nonGemini = minFraction( + raw['non-gemini'] as QuotaGroupSummary | undefined, + minFraction( + raw.claude as QuotaGroupSummary | undefined, + raw['gpt-oss'] as QuotaGroupSummary | undefined, + ), + ) + + return { + ...(gemini !== undefined ? { gemini } : {}), + ...(nonGemini !== undefined ? { 'non-gemini': nonGemini } : {}), + } +} + export interface PerModelQuotaEntry { modelId: string displayName?: string diff --git a/packages/e2e-tests/src/harness.ts b/packages/e2e-tests/src/harness.ts index a220266..4bba3d1 100644 --- a/packages/e2e-tests/src/harness.ts +++ b/packages/e2e-tests/src/harness.ts @@ -37,6 +37,8 @@ export interface E2eHarness { createPlugin(options?: { clientOverrides?: Record extraDependencies?: Partial + /** Override specific fields in the project-level config written to disk. */ + configOverrides?: Record }): Promise /** Tear down all built plugins + the mock server + temp dirs. */ dispose(): Promise @@ -109,7 +111,10 @@ export async function createE2eHarness( soft_quota_threshold_percent: 100, quota_refresh_interval_minutes: 0, proactive_rotation_threshold_percent: 0, + background_quota_refresh: false, auto_update: false, + // Allow per-test overrides (e.g. enabling background_quota_refresh). + ...(pluginOptions.configOverrides ?? {}), }), ) const client = createFakeClient(pluginOptions.clientOverrides) diff --git a/packages/e2e-tests/src/plugin-flow.e2e.test.ts b/packages/e2e-tests/src/plugin-flow.e2e.test.ts index 1ce4698..ca54dbe 100644 --- a/packages/e2e-tests/src/plugin-flow.e2e.test.ts +++ b/packages/e2e-tests/src/plugin-flow.e2e.test.ts @@ -322,6 +322,94 @@ describe('plugin flow (e2e)', () => { }) }) + it('V3 poller: background_quota_refresh=true exercises the default-on path and refreshes idle accounts', async () => { + // The harness globally sets background_quota_refresh:false to prevent + // network calls outside the fetch interceptor. This test overrides that + // and verifies the poller path through the real factory: + // - plugin boots with background_quota_refresh:true + // - accounts are loaded (seeded by beforeEach) + // - the poller fires (stale freshness gate) and calls the mock server + // - the sidebar state file is updated with non-empty quota + await withHarness(async (h) => { + const sidebarPath = process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + if (!sidebarPath) throw new Error('sidebar path missing') + + // Use a longer startup jitter (350ms) so we have a reliable window + // to write a stale sidebar AFTER the auth loader (which writes a fresh + // checkedAt) but BEFORE the first poller tick. This makes the only + // write that can advance past `pollBefore` be the poller's own tick. + const plugin = await h.createPlugin({ + configOverrides: { + background_quota_refresh: true, + background_quota_refresh_interval_minutes: 1, + }, + extraDependencies: { + // random=0.012 → jitter = floor(0.012 * 30000) = 360ms. + // Auth loader runs in ~20ms; we write the stale sidebar at ~250ms + // — leaving ~110ms before the first tick fires. + clock: { random: () => 0.012 }, + }, + }) + + // Warm up the auth loader so the AccountManager is populated. + // This will write a fresh checkedAt to the sidebar (~T+20ms). + await plugin.auth.loader( + async () => ({ + type: 'oauth' as const, + refresh: 'refresh-a|project-a|managed-a', + access: 'access-a', + expires: Date.now() + 3_600_000, + }), + {} as Parameters[1], + ) + + // Give the auth loader\'s async sidebar write time to land, then + // forcibly overwrite with a stale checkedAt so the poller\'s freshness + // gate passes on its first tick (~360ms from start). Must bypass the + // merge logic (which rejects stale writes) — write directly to disk. + await new Promise((r) => setTimeout(r, 200)) + require('node:fs').writeFileSync( + sidebarPath, + JSON.stringify({ + version: 1, + checkedAt: Date.now() - 60_000, + accounts: [], + activeRouting: {}, + routingAuthoritative: false, + }), + ) + + // Capture pollBefore AFTER the stale write. Now the ONLY write that + // can advance checkedAt past this point is the poller\'s first tick. + const pollBefore = Date.now() + + const deadline = pollBefore + 5_000 + let sidebarCheckedAt = 0 + while (Date.now() < deadline) { + await new Promise((r) => setTimeout(r, 100)) + try { + const raw = require('node:fs').readFileSync( + sidebarPath, + 'utf8', + ) as string + const parsed = JSON.parse(raw) as { checkedAt?: number } + sidebarCheckedAt = parsed.checkedAt ?? 0 + if (sidebarCheckedAt > pollBefore) break + } catch { + // File not written yet — retry. + } + } + + await plugin.dispose() + + // The poller must have fired, passed the freshness gate on the stale + // sidebar, and written a fresh checkedAt. This assertion can ONLY be + // satisfied by the poller — not by the auth loader (which ran before + // pollBefore and wrote the now-overwritten stale checkedAt). + expect(sidebarCheckedAt).toBeGreaterThan(pollBefore) + }) + }) + it('dispose() releases every subsystem and clears the auth loader runtime', async () => { await withHarness(async (h) => { const plugin = await h.createPlugin() diff --git a/packages/opencode/assets/antigravity.schema.json b/packages/opencode/assets/antigravity.schema.json index e3ef5d1..77ba2d2 100644 --- a/packages/opencode/assets/antigravity.schema.json +++ b/packages/opencode/assets/antigravity.schema.json @@ -360,6 +360,16 @@ ], "additionalProperties": false }, + "background_quota_refresh": { + "default": true, + "type": "boolean" + }, + "background_quota_refresh_interval_minutes": { + "default": 5, + "type": "number", + "minimum": 1, + "maximum": 60 + }, "auto_update": { "default": true, "type": "boolean", diff --git a/packages/opencode/src/plugin/account-command-oauth.ts b/packages/opencode/src/plugin/account-command-oauth.ts index 846870e..79495e5 100644 --- a/packages/opencode/src/plugin/account-command-oauth.ts +++ b/packages/opencode/src/plugin/account-command-oauth.ts @@ -83,6 +83,17 @@ export function createAccountCommandOAuthService( return entry } + // Wrap listAccounts so a lock-contention or I/O failure in the error + // branches below does not replace the real OAuth failure text with a + // generic apply-error message. Mirrors the stage-4 guarded pattern. + const safeListAccounts = async (): Promise => { + try { + return await options.listAccounts() + } catch { + return [] + } + } + return { async start(sessionId) { const authorization = await options.authorize() @@ -107,7 +118,7 @@ export function createAccountCommandOAuthService( redirectUri, createdAt: now(), }) - return { url: authorization.url, accounts: await options.listAccounts() } + return { url: authorization.url, accounts: await safeListAccounts() } }, async finish(sessionId, callbackInput, label) { @@ -115,7 +126,7 @@ export function createAccountCommandOAuthService( if (!pending) { return { text: 'OAuth session expired. Please start again.', - accounts: await options.listAccounts(), + accounts: await safeListAccounts(), } } @@ -127,13 +138,13 @@ export function createAccountCommandOAuthService( } catch { return { text: 'OAuth authentication failed: could not parse the callback. Please start a new OAuth flow.', - accounts: await options.listAccounts(), + accounts: await safeListAccounts(), } } if ('error' in callback) { return { text: `OAuth authentication failed: ${callback.error}. Please start a new OAuth flow.`, - accounts: await options.listAccounts(), + accounts: await safeListAccounts(), } } @@ -146,13 +157,13 @@ export function createAccountCommandOAuthService( } catch { return { text: 'OAuth exchange failed due to a network error. Please start a new OAuth flow.', - accounts: await options.listAccounts(), + accounts: await safeListAccounts(), } } if (result.type === 'failed') { return { text: 'OAuth authentication failed. Please start a new OAuth flow and try again.', - accounts: await options.listAccounts(), + accounts: await safeListAccounts(), } } @@ -168,7 +179,7 @@ export function createAccountCommandOAuthService( } catch { return { text: 'OAuth account could not be saved to disk. Please start a new OAuth flow.', - accounts: await options.listAccounts(), + accounts: await safeListAccounts(), } } diff --git a/packages/opencode/src/plugin/auth-loader.test.ts b/packages/opencode/src/plugin/auth-loader.test.ts index fdd3bc1..a357b13 100644 --- a/packages/opencode/src/plugin/auth-loader.test.ts +++ b/packages/opencode/src/plugin/auth-loader.test.ts @@ -59,6 +59,7 @@ describe('createAuthLoader', () => { }, ], requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), dispose: mock(async () => {}), } const { lifecycle, replacements } = createLifecycle() @@ -159,6 +160,7 @@ describe('createAuthLoader', () => { }, ], requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), dispose: mock(async () => {}), } const secondManager = { @@ -174,6 +176,7 @@ describe('createAuthLoader', () => { }, ], requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), dispose: mock(async () => {}), } const managers = [firstManager, secondManager] @@ -262,6 +265,7 @@ describe('createAuthLoader', () => { }, ], requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), dispose: mock(async () => {}), } const managerB = { @@ -277,6 +281,7 @@ describe('createAuthLoader', () => { }, ], requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), dispose: mock(async () => {}), } const managers = [managerA, managerB] @@ -381,6 +386,7 @@ describe('createAuthLoader', () => { }, ], requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), dispose: mock(async () => {}), } const secondManager = { @@ -396,6 +402,7 @@ describe('createAuthLoader', () => { }, ], requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), dispose: mock(async () => {}), } const managers = [firstManager, secondManager] diff --git a/packages/opencode/src/plugin/auth-loader.ts b/packages/opencode/src/plugin/auth-loader.ts index 12d0f08..2297ed6 100644 --- a/packages/opencode/src/plugin/auth-loader.ts +++ b/packages/opencode/src/plugin/auth-loader.ts @@ -1,7 +1,10 @@ import { createHash } from 'node:crypto' +import { getHealthTracker } from '@cortexkit/antigravity-auth-core' import { buildSidebarMachineStateFromAccounts, + isAccountCurrent, setSidebarMachineState, + toCapturedTier, } from '../sidebar-state' import { AccountManager } from './accounts' import { isOAuthAuth } from './auth' @@ -180,21 +183,28 @@ export function createAuthLoader({ // without waiting for the first fetch to complete. await setSidebarMachineState( buildSidebarMachineStateFromAccounts( - accountManager.getAccounts().map((entry) => ({ - index: entry.index, - label: entry.label, - enabled: entry.enabled, - current: false, - coolingDownUntil: entry.coolingDownUntil, - cachedQuota: entry.cachedQuota, - // Stamp the sidebar snapshot so the projection can detect a - // stale cache that landed on the wrong account (the manager's - // `cachedQuotaAccountId` is keyed to whatever account actually - // produced the snapshot — the live refresh-token hash is the - // expected identity at this slot). - cachedQuotaAccountId: entry.cachedQuotaAccountId, - currentQuotaAccountId: refreshTokenIdentity(entry.parts.refreshToken), - })), + accountManager.getAccounts().map((entry) => { + const activeByFamily = accountManager.getActiveIndexByFamily() + return { + index: entry.index, + label: entry.label, + enabled: entry.enabled, + current: isAccountCurrent(entry.index, activeByFamily), + coolingDownUntil: entry.coolingDownUntil, + healthScore: getHealthTracker().getScore(entry.index), + cachedQuota: entry.cachedQuota, + // Stamp the sidebar snapshot so the projection can detect a + // stale cache that landed on the wrong account (the manager's + // `cachedQuotaAccountId` is keyed to whatever account actually + // produced the snapshot — the live refresh-token hash is the + // expected identity at this slot). + cachedQuotaAccountId: entry.cachedQuotaAccountId, + currentQuotaAccountId: refreshTokenIdentity( + entry.parts.refreshToken, + ), + tier: toCapturedTier(entry), + } + }), ), ) } diff --git a/packages/opencode/src/plugin/background-quota-refresh.test.ts b/packages/opencode/src/plugin/background-quota-refresh.test.ts new file mode 100644 index 0000000..0d3034c --- /dev/null +++ b/packages/opencode/src/plugin/background-quota-refresh.test.ts @@ -0,0 +1,1981 @@ +import { afterEach, beforeEach, describe, expect, it, mock } from 'bun:test' +import { createHash } from 'node:crypto' +import { mkdirSync, mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' + +/** Mirror of the private quotaAccountIdentity used in production code. */ +function quotaAccountIdentity(refreshToken: string): string { + return createHash('sha256').update(refreshToken).digest('hex').slice(0, 16) +} + +import type { AccountMetadataV3 } from '@cortexkit/antigravity-auth-core' + +import { + readSidebarState, + SIDEBAR_STATE_ENV, + setSidebarMachineState, +} from '../sidebar-state' +import { + BackgroundQuotaRefresh, + type BackgroundQuotaRefreshOptions, + type PollerAccountView, +} from './background-quota-refresh' +import { createAntigravityPlugin } from './index' +import type { PluginClient, PluginInput } from './types' + +// ─── Minimal harness helpers ───────────────────────────────────────────────── + +function makeMinimalClient(): PluginClient { + return { + app: { log: mock(async () => {}) }, + auth: { set: mock(async () => {}) }, + session: { + abort: mock(async () => {}), + messages: mock(async () => ({ data: [] })), + prompt: mock(async () => {}), + }, + tui: { showToast: mock(async () => {}) }, + } as unknown as PluginClient +} + +function makePluginInput(client: PluginClient, directory: string): PluginInput { + return { + client, + project: {} as PluginInput['project'], + directory, + worktree: directory, + experimental_workspace: { register: mock(() => {}) }, + serverUrl: new URL('http://localhost:4096'), + $: (() => {}) as unknown as PluginInput['$'], + } +} + +// Build a fake AccountManager that satisfies BackgroundQuotaRefresh's usage. +type StubAccount = ReturnType[number] + +function makeStubAccount(refreshToken = 'tok-abc'): StubAccount { + return { + index: 0, + label: 'Account', + enabled: true, + parts: { refreshToken }, + cachedQuota: undefined, + cachedQuotaAccountId: undefined, + coolingDownUntil: undefined, + } +} + +function makeAccountManager( + accounts: StubAccount[] = [makeStubAccount()], + activeByFamily: { claude: number; gemini: number } = { claude: 0, gemini: 0 }, +): PollerAccountView { + return { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: (): AccountMetadataV3[] => + accounts.map((a) => ({ + refreshToken: a.parts.refreshToken, + enabled: a.enabled, + addedAt: 0, + lastUsed: 0, + })), + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => activeByFamily, + } +} + +// Build a QuotaManager stub that returns configurable results. +function makeQuotaManager( + opts: { + result?: 'ok' | 'error' | 'pending' + onRefreshAccounts?: () => void + delay?: number + } = {}, +) { + const { result = 'ok', delay = 0 } = opts + const refreshAccounts = mock( + async (accounts: Array<{ refreshToken: string }>, _options: unknown) => { + if (opts.onRefreshAccounts) opts.onRefreshAccounts() + if (delay > 0) await new Promise((r) => setTimeout(r, delay)) + if (result === 'error') throw new Error('quota fetch failed') + if (result === 'ok') { + return accounts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { + groups: { + 'non-gemini': { remainingFraction: 0.5, modelCount: 1 }, + }, + }, + })) + } + return [] + }, + ) + return { + refreshAccounts, + dispose: mock(async () => {}), + } +} + +type TierStubAccount = StubAccount & { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number +} + +function makeTierManager(account: TierStubAccount): PollerAccountView { + return { + getAccounts: () => [account], + getAccountsForQuotaCheck: () => [ + { + refreshToken: account.parts.refreshToken, + enabled: account.enabled, + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock((_index, patch) => Object.assign(account, patch)), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } +} + +type TickablePoller = { + runTick: () => Promise +} + +function makeTwoPollers(options: { + stateFile: string + now: number + firstManager: PollerAccountView + secondManager: PollerAccountView + firstQuotaManager?: ReturnType + secondQuotaManager?: ReturnType + firstLoadAccountTier?: BackgroundQuotaRefreshOptions['loadAccountTier'] + secondLoadAccountTier?: BackgroundQuotaRefreshOptions['loadAccountTier'] +}) { + const first = new BackgroundQuotaRefresh({ + intervalMs: 5 * 60_000, + sidebarStateFile: options.stateFile, + getAccountManager: () => options.firstManager, + quotaManager: (options.firstQuotaManager ?? + makeQuotaManager()) as unknown as import('./quota').QuotaManager, + now: () => options.now, + random: () => 0, + loadAccountTier: options.firstLoadAccountTier, + }) + const second = new BackgroundQuotaRefresh({ + intervalMs: 5 * 60_000, + sidebarStateFile: options.stateFile, + getAccountManager: () => options.secondManager, + quotaManager: (options.secondQuotaManager ?? + makeQuotaManager()) as unknown as import('./quota').QuotaManager, + now: () => options.now, + random: () => 0, + loadAccountTier: options.secondLoadAccountTier, + }) + + return { + first, + second, + tick: (poller: BackgroundQuotaRefresh) => + (poller as unknown as TickablePoller).runTick(), + dispose: () => Promise.all([first.dispose(), second.dispose()]), + } +} + +// ─── Suite ─────────────────────────────────────────────────────────────────── + +describe('BackgroundQuotaRefresh', () => { + let dir: string + let stateFile: string + let savedSidebarEnv: string | undefined + + beforeEach(() => { + savedSidebarEnv = process.env[SIDEBAR_STATE_ENV] + dir = mkdtempSync(join(tmpdir(), 'bg-quota-test-')) + stateFile = join(dir, 'sidebar-state.json') + process.env[SIDEBAR_STATE_ENV] = stateFile + }) + + afterEach(async () => { + // Restore rather than delete — a delete drops resolution to the real state dir. + if (savedSidebarEnv !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebarEnv + else delete process.env[SIDEBAR_STATE_ENV] + rmSync(dir, { recursive: true, force: true }) + }) + + // ── Timer: idempotent start/stop ────────────────────────────────────────── + + it('start is idempotent — a second call before the first tick does not double-schedule', async () => { + const calls: number[] = [] + const manager = makeAccountManager() + const quotaManager = makeQuotaManager({ + onRefreshAccounts: () => calls.push(Date.now()), + }) + const poller = new BackgroundQuotaRefresh({ + intervalMs: 5 * 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + // Delay the first tick past dispose so we can check nothing ran. + random: () => 1, // max startup jitter → no immediate tick + }) + poller.start() + poller.start() + await poller.dispose() + // No tick ran within the jitter window — both starts share one timer. + expect(calls).toHaveLength(0) + }) + + it('dispose after dispose is idempotent', async () => { + const manager = makeAccountManager() + const quotaManager = makeQuotaManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 1, + }) + poller.start() + await poller.dispose() + await poller.dispose() // must not throw + }) + + // ── V4: unexpected throw before internal reschedule still reschedules ────── + + it('V4 liveness: a throw escaping runTick before any scheduleNext still reschedules', async () => { + // runTick has three internal reschedule paths (freshness skip, lock held, + // lock throws). A throw that ESCAPES the function before any of those + // paths previously meant no next tick was queued — the poller died. + // The fix adds a reschedule in the outer .catch() so the poller + // degrades to a delayed retry rather than silently stopping. + // + // We verify this by observing that scheduleNext IS called after the throw, + // rather than waiting for the full ERROR_JITTER_MS delay (15 s). + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const manager = makeAccountManager() + const quotaManager = makeQuotaManager() + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + + // Intercept scheduleNext to observe calls after the catch. + type PollerInternals = { + runTick: () => Promise + scheduleNext: (ms: number) => void + } + const inner = poller as unknown as PollerInternals + const scheduleNextCalls: number[] = [] + const origScheduleNext = inner.scheduleNext.bind(poller) + inner.scheduleNext = (ms: number) => { + scheduleNextCalls.push(ms) + origScheduleNext(ms) + } + + // Replace runTick to throw before any internal scheduleNext can run. + const origRunTick = inner.runTick.bind(poller) + let firstCall = true + inner.runTick = async () => { + if (firstCall) { + firstCall = false + throw new Error('simulated unexpected escape from runTick') + } + return origRunTick() + } + + poller.start() + // The startup scheduleNext fires synchronously with jitteredStartDelay + // (random=0 → 0 ms) — tick runs almost immediately, throws, outer catch + // queues the recovery scheduleNext. Allow a brief event-loop drain. + await new Promise((r) => setTimeout(r, 50)) + await poller.dispose() + + // The outer catch must have called scheduleNext (the recovery reschedule). + // scheduleNextCalls[0] is the startup call; [1] is from the outer catch. + expect(scheduleNextCalls.length).toBeGreaterThanOrEqual(2) + // The recovery delay includes ERROR_JITTER_MS (15_000). + const recoveryDelay = scheduleNextCalls[1] + expect(recoveryDelay).toBeGreaterThanOrEqual(15_000) + }) + + // ── Error in tick does not kill the timer ───────────────────────────────── + + it('an error inside a tick does not prevent subsequent ticks', async () => { + let callCount = 0 + let secondCallResolve!: () => void + const secondCallSeen = new Promise((r) => { + secondCallResolve = r + }) + + const manager = makeAccountManager() + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const quotaManager = { + refreshAccounts: mock(async () => { + callCount++ + if (callCount === 1) throw new Error('first tick error') + secondCallResolve() + return [] + }), + dispose: mock(async () => {}), + } + + const poller = new BackgroundQuotaRefresh({ + // Very short interval so the second tick fires within the timeout. + intervalMs: 100, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + await Promise.race([ + secondCallSeen, + new Promise((_, reject) => + setTimeout( + () => reject(new Error('timeout waiting for second tick')), + 5_000, + ), + ), + ]) + await poller.dispose() + expect(callCount).toBeGreaterThanOrEqual(2) + }) + + // ── In-flight tick not double-entered ──────────────────────────────────── + + it('a slow tick prevents re-entry while it is in flight', async () => { + let active = 0 + let maxActive = 0 + let tickDone!: () => void + const tickBlocked = new Promise((r) => { + tickDone = r + }) + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const quotaManager = { + refreshAccounts: mock(async () => { + active++ + maxActive = Math.max(maxActive, active) + await tickBlocked + active-- + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + // Give the first tick a moment to start. + await new Promise((r) => setTimeout(r, 20)) + tickDone() + await poller.dispose() + expect(maxActive).toBe(1) + }) + + // ── Async dispose awaits in-flight tick ─────────────────────────────────── + + it('dispose awaits an in-flight tick before resolving', async () => { + const events: string[] = [] + + let tickStartedResolve!: () => void + const tickStartedP = new Promise((r) => { + tickStartedResolve = r + }) + let tickRelease!: () => void + const tickGate = new Promise((r) => { + tickRelease = r + }) + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const quotaManager = { + refreshAccounts: mock(async () => { + events.push('tick:start') + tickStartedResolve() + await tickGate + events.push('tick:finish') + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + await tickStartedP + // Tick is now in flight. Dispose must wait for it. + const disposeP = poller.dispose().then(() => { + events.push('dispose:done') + }) + tickRelease() + await disposeP + // dispose:done must follow tick:finish + expect(events).toEqual(['tick:start', 'tick:finish', 'dispose:done']) + }) + + it('S3 boundary: re-entry guard prevents inFlight overwrite when scheduleNext fires mid-tick', async () => { + // The natural test scenario (intervalMs < tick duration) never reaches the + // guard because runTick only calls scheduleNext AFTER it completes, so the + // second timer always fires after inFlight is already null. + // This test triggers the guard directly via reflection: call scheduleNext(0) + // while tick 1 is blocked, then assert the inFlight reference is unchanged. + // Without the guard the 0ms timer would overwrite inFlight; with it, the + // timer callback exits early and the original promise is preserved. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + let tickStartedResolve!: () => void + const tickStartedP = new Promise((r) => { + tickStartedResolve = r + }) + let releaseFirstTick!: () => void + const firstTickGate = new Promise((r) => { + releaseFirstTick = r + }) + + const quotaManager = { + refreshAccounts: mock(async () => { + tickStartedResolve() + await firstTickGate + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + + // Reach private members via cast to verify the guard's effect. + type PollerInternals = { + inFlight: Promise | null + scheduleNext: (ms: number) => void + } + const inner = poller as unknown as PollerInternals + + poller.start() + await tickStartedP // tick 1 is now blocked inside refreshAccounts + + const originalInFlight = inner.inFlight + expect(originalInFlight).not.toBeNull() + + // Force a 0ms timer while the tick is in flight. Without the guard, this + // overwrites inFlight; with it, the timer callback returns immediately. + inner.scheduleNext(0) + await new Promise((r) => setTimeout(r, 20)) // let the 0ms timer fire + + // Guard must have prevented the overwrite. + expect(inner.inFlight).toBe(originalInFlight) + + releaseFirstTick() + await poller.dispose() + }) + + // ── Freshness gate ──────────────────────────────────────────────────────── + + it('freshness gate: skips when checkedAt is fresher than threshold (5-minute interval)', async () => { + const intervalMs = 5 * 60_000 // 300 000 ms + // Threshold = max(300000-60000, floor(300000/2)) = max(240000, 150000) = 240 000 ms + // A checkedAt 239 s ago is still within the threshold → skip. + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 239_000, accounts: [] }, + { stateFile }, + ) + + const quotaManager = makeQuotaManager() + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => now, + }) + poller.start() + // Allow one event loop tick for the timer. + await new Promise((r) => setTimeout(r, 50)) + await poller.dispose() + expect(quotaManager.refreshAccounts).not.toHaveBeenCalled() + }) + + it('freshness gate: refreshes when checkedAt is staler than threshold (5-minute interval)', async () => { + const intervalMs = 5 * 60_000 + // checkedAt 241 s ago → stale → should refresh. + let refreshCalled = false + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 241_000, accounts: [] }, + { stateFile }, + ) + + const quotaManager = { + refreshAccounts: mock(async () => { + refreshCalled = true + refreshResolve() + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => now, + }) + poller.start() + await Promise.race([ + refreshedP, + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 3_000), + ), + ]) + await poller.dispose() + expect(refreshCalled).toBe(true) + }) + + it('freshness gate at 1-minute minimum (floor(intervalMs/2) = 30 s)', async () => { + const intervalMs = 60_000 // 1 minute + // Threshold = max(60000-60000, floor(60000/2)) = max(0, 30000) = 30 000 ms + // checkedAt 29 s ago → skip. + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 29_000, accounts: [] }, + { stateFile }, + ) + + const quotaManager = makeQuotaManager() + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => now, + }) + poller.start() + await new Promise((r) => setTimeout(r, 50)) + await poller.dispose() + expect(quotaManager.refreshAccounts).not.toHaveBeenCalled() + }) + + it('freshness gate at 1-minute minimum refreshes when stale (31 s age)', async () => { + const intervalMs = 60_000 + // checkedAt 31 s ago → stale → refresh. + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 31_000, accounts: [] }, + { stateFile }, + ) + + const quotaManager = { + refreshAccounts: mock(async () => { + refreshResolve() + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => now, + }) + poller.start() + await Promise.race([ + refreshedP, + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 3_000), + ), + ]) + await poller.dispose() + expect(quotaManager.refreshAccounts).toHaveBeenCalled() + }) + + it('freshness gate at 60-minute interval uses intervalMs - 60 000 threshold', async () => { + const intervalMs = 60 * 60_000 // 3 600 000 ms + // Threshold = max(3600000-60000, floor(3600000/2)) = max(3540000, 1800000) = 3 540 000 ms + // checkedAt 3 539 s ago → fresh → skip. + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 3_539_000, accounts: [] }, + { stateFile }, + ) + + const quotaManager = makeQuotaManager() + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => now, + }) + poller.start() + await new Promise((r) => setTimeout(r, 50)) + await poller.dispose() + expect(quotaManager.refreshAccounts).not.toHaveBeenCalled() + }) + + it('two pollers: a fresh snapshot from one writer does not starve the other tier slot', async () => { + const now = Date.now() + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const firstQuotaManager = makeQuotaManager() + const secondQuotaManager = makeQuotaManager() + const secondLoadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: now, + })) + const pair = makeTwoPollers({ + stateFile, + now, + firstManager: makeTierManager({ + ...makeStubAccount('tok-first-fresh-tier'), + capturedTierId: 'free-tier', + capturedPaidTierId: 'g1-pro-tier', + capturedTierAt: now, + capturedTierSchemaVersion: 1, + }), + secondManager: makeTierManager(makeStubAccount('tok-second-stale-tier')), + firstQuotaManager, + secondQuotaManager, + secondLoadAccountTier, + }) + + await pair.tick(pair.first) + await pair.tick(pair.second) + await pair.dispose() + + expect(firstQuotaManager.refreshAccounts).toHaveBeenCalledTimes(1) + expect(secondQuotaManager.refreshAccounts).not.toHaveBeenCalled() + expect(secondLoadAccountTier).toHaveBeenCalledTimes(1) + }) + + it('two pollers: a shared stale account resolves tier once in one lock window', async () => { + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 60_000, accounts: [] }, + { stateFile }, + ) + + const account = makeStubAccount('tok-shared-stale-tier') + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: now, + })) + const pair = makeTwoPollers({ + stateFile, + now, + firstManager: makeTierManager(account), + secondManager: makeTierManager(account), + firstLoadAccountTier: loadAccountTier, + secondLoadAccountTier: loadAccountTier, + }) + + await Promise.all([pair.tick(pair.first), pair.tick(pair.second)]) + await pair.dispose() + + expect(loadAccountTier).toHaveBeenCalledTimes(1) + }) + + it('two pollers: both tier slots eventually spend while their shared snapshot stays fresh', async () => { + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 60_000, accounts: [] }, + { stateFile }, + ) + + const firstLoadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: now, + })) + const secondLoadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: now, + })) + const pair = makeTwoPollers({ + stateFile, + now, + firstManager: makeTierManager(makeStubAccount('tok-first-stale-tier')), + secondManager: makeTierManager(makeStubAccount('tok-second-stale-tier')), + firstLoadAccountTier, + secondLoadAccountTier, + }) + + await pair.tick(pair.first) + await pair.tick(pair.second) + await pair.dispose() + + expect(firstLoadAccountTier).toHaveBeenCalledTimes(1) + expect(secondLoadAccountTier).toHaveBeenCalledTimes(1) + }) + + // ── Lock: held → skip ──────────────────────────────────────────────────── + + it('skips when the fenced lock is held by another process', async () => { + // Acquire the lock ourselves so the poller sees "held". + const { acquireFencedFileLock } = await import( + '@cortexkit/antigravity-auth-core/file-lock' + ) + const lock = await acquireFencedFileLock({ + path: stateFile, + name: 'bg-quota-poll', + ttlMs: 60_000, + renew: false, + }) + expect(lock).not.toBeNull() + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const quotaManager = makeQuotaManager() + const manager = makeAccountManager() + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: Date.now(), + })) + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + loadAccountTier, + }) + poller.start() + await new Promise((r) => setTimeout(r, 100)) + await poller.dispose() + await lock?.release() + + // Lock was held → refresh must not have been called. + expect(quotaManager.refreshAccounts).not.toHaveBeenCalled() + expect(loadAccountTier).not.toHaveBeenCalled() + }) + + // ── Lock throws → FAIL CLOSED ──────────────────────────────────────────── + + it('fail-closed: acquireFencedFileLock throw skips the tick without a fallback mutex', async () => { + // Use a NUL-byte in the path: the kernel rejects mkdir("...\0...") with + // ENOENT/EINVAL immediately, which causes acquireFencedFileLock to throw + // before touching any lock file. Two pollers share the same valid + // sidebarStateFile (for the freshness read), but tick 1 uses the bad lock + // path to trigger the throw, tick 2 uses the real path to confirm the + // timer rescheduled and a real refresh is possible. + // + // Crucially: no "claim marker" or secondary mutex files appear anywhere. + // This is the load-bearing invariant from the module-level comment. + const nullPath = join(dir, 'bad\x00path', 'state.json') + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const quotaManager = { + refreshAccounts: mock(async () => { + refreshResolve() + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: Date.now(), + })) + + // Poller 1: bad lock path → tick throws → fail closed → refresh NOT called. + const pollerBad = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + // The lock is attempted against sidebarStateFile (the valid path) by + // default. We need the lock to use the nullPath instead. Since the lock + // path is derived from sidebarStateFile (by appending `.bg-quota-poll.lock`) + // we point sidebarStateFile at nullPath so the lock write uses it. + // The freshness read will fail (file doesn't exist) and return checkedAt=0, + // which passes the gate — and then the lock attempt throws on NUL-byte. + sidebarStateFile: nullPath, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + loadAccountTier, + }) + pollerBad.start() + // Give it a moment to attempt and fail the lock. + await new Promise((r) => setTimeout(r, 100)) + await pollerBad.dispose() + + // Lock threw → refresh must NOT have been called on the bad poller. + expect(quotaManager.refreshAccounts).not.toHaveBeenCalled() + expect(loadAccountTier).not.toHaveBeenCalled() + + // Confirm NO claim/marker files were created — the fail-closed path + // must not fall back to a secondary mutex. + const { readdirSync } = await import('node:fs') + const anyLockFiles = readdirSync(dir).filter( + (f) => + f.includes('.lock') || f.includes('.claim') || f.includes('.marker'), + ) + expect(anyLockFiles).toHaveLength(0) + + // Poller 2: real path → refresh succeeds → confirms the timer is still + // functional after a fail-closed skip (not a property of pollerBad's timer + // but of the code path: the catch reschedules, so another poller can run). + const pollerGood = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + pollerGood.start() + await Promise.race([ + refreshedP, + new Promise((_, reject) => + setTimeout( + () => reject(new Error('timeout waiting for good refresh')), + 5_000, + ), + ), + ]) + await pollerGood.dispose() + expect(quotaManager.refreshAccounts).toHaveBeenCalledTimes(1) + }) + + // ── N-contender bound ───────────────────────────────────────────────────── + + it('N-contender bound: concurrent pollers against a slow refresh never all refresh', async () => { + // A slow refresh is simulated by inserting a delay. We run N pollers with + // zero startup jitter, all stale, and count how many refreshes execute + // BEFORE the first write lands. The freshness gate + lock means at most + // one poller should execute the refresh. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + let refreshes = 0 + let firstRefreshDone!: () => void + const firstRefreshP = new Promise((r) => { + firstRefreshDone = r + }) + + const quotaManager = { + refreshAccounts: mock(async () => { + refreshes++ + // Simulate a slow fetch — enough for all other pollers to attempt. + await new Promise((r) => setTimeout(r, 150)) + firstRefreshDone() + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + + const N = 4 + const pollers = Array.from({ length: N }, () => { + const p = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + p.start() + return p + }) + + // Wait for the first refresh to start and finish. + await Promise.race([ + firstRefreshP, + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + // At this point the first refresh just completed. Count before cleanup. + const refreshesSoFar = refreshes + await Promise.all(pollers.map((p) => p.dispose())) + + // Only 1 refresh should have run before the first write landed — + // the freshness gate blocks all others once checkedAt is updated. + expect(refreshesSoFar).toBe(1) + }) + + // ── Windows survive a polled refresh ───────────────────────────────────── + + it('windows-survive: a polled refresh writes per-window data into the sidebar state', async () => { + // This test exercises the real path: + // BackgroundQuotaRefresh.refresh() + // → quotaManager.refreshAccounts (returns per-window data) + // → manager.updateQuotaCache + // → pushSidebarQuotaSnapshot (reads the live view) + // → sidebar state file (verified via readSidebarState) + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const windowData = { + 'non-gemini': { + remainingFraction: 0.72, + modelCount: 1, + windows: [ + { + window: 'weekly' as const, + remainingFraction: 0.72, + resetTime: new Date( + Date.now() + 7 * 24 * 60 * 60_000, + ).toISOString(), + }, + { + window: '5h' as const, + remainingFraction: 0.91, + resetTime: new Date(Date.now() + 5 * 60 * 60_000).toISOString(), + }, + ], + }, + } + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const accounts = [ + { + index: 0, + label: 'Account', + enabled: true, + parts: { refreshToken: 'tok-win' }, + cachedQuota: undefined as typeof windowData | undefined, + cachedQuotaAccountId: undefined as string | undefined, + coolingDownUntil: undefined as number | undefined, + }, + ] + + const quotaManager = { + refreshAccounts: mock(async (accts: Array<{ refreshToken: string }>) => { + return accts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { groups: windowData }, + })) + }), + dispose: mock(async () => {}), + } + + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: (): AccountMetadataV3[] => [ + { + refreshToken: 'tok-win', + enabled: true, + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache: ( + _index: number, + groups: Partial< + Record + >, + _token: string, + ) => { + accounts[0]!.cachedQuota = groups as typeof windowData + // Use the real identity derived from the refresh token so the + // currentQuotaAccountId computed by the poller (same hash) matches. + accounts[0]!.cachedQuotaAccountId = quotaAccountIdentity('tok-win') + refreshResolve() + }, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await Promise.race([ + // Wait a bit extra to let pushSidebarQuotaSnapshot complete after updateQuotaCache. + refreshedP.then(() => new Promise((r) => setTimeout(r, 200))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + await poller.dispose() + + const state = readSidebarState(stateFile) + expect(state.accounts).toHaveLength(1) + const quota = state.accounts[0]?.quota['non-gemini'] + expect(quota?.remainingPercent).toBe(72) + // Per-window entries must survive through the sidebar read path. + expect(quota?.windows).toBeDefined() + expect(quota?.windows?.length).toBeGreaterThanOrEqual(1) + const weekly = quota?.windows?.find((w) => w.window === 'weekly') + expect(weekly).toBeDefined() + expect(weekly?.remainingPercent).toBe(72) + }) + + // ── V1: stale cachedQuotaAccountId is dropped after a polled refresh ──── + + it('V1 identity: poller writes currentQuotaAccountId from live token, not from cached stamp', async () => { + // An account whose cachedQuotaAccountId does NOT match its live + // refresh-token identity — simulating an account that was reordered + // or replaced. After a poll-written snapshot the sidebar must DROP the + // stale cached quota rather than displaying the wrong account's bars. + // + // Without the fix both fields were the same value so the + // staleness check never fired. + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const refreshToken = 'tok-fresh-account' + // A stale stamp that will never match the hash of refreshToken. + const staleStamp = 'stale-id-000000000000' + + const accounts = [ + { + index: 0, + label: 'Account', + enabled: true, + parts: { refreshToken }, + // Stale quota snapshot from a previous account at this index. + cachedQuota: { + 'non-gemini': { remainingFraction: 0.5, modelCount: 1 }, + } as ReturnType[0]['cachedQuota'], + cachedQuotaAccountId: staleStamp, + coolingDownUntil: undefined as number | undefined, + }, + ] + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const quotaManager = { + // The quota fetch itself succeeds and returns new data. + refreshAccounts: mock(async (accts: Array<{ refreshToken: string }>) => { + return accts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { + groups: { + 'non-gemini': { remainingFraction: 0.8, modelCount: 1 }, + }, + }, + })) + }), + dispose: mock(async () => {}), + } + + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: (_index: number, groups: unknown, _token: string) => { + // Update the in-memory cache but keep the STALE stamp — simulating + // an account where the identity hasn't been updated yet. + accounts[0]!.cachedQuota = groups as (typeof accounts)[0]['cachedQuota'] + // Intentionally DO NOT update cachedQuotaAccountId here: the stamp + // stays stale so the sidebar projection must detect the mismatch. + refreshResolve() + }, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await Promise.race([ + refreshedP.then(() => new Promise((r) => setTimeout(r, 200))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout waiting for poll')), 5_000), + ), + ]) + await poller.dispose() + + const state = readSidebarState(stateFile) + expect(state.accounts).toHaveLength(1) + // The stale stamp must cause the quota to be dropped from the sidebar. + const sidebarQuota = state.accounts[0]?.quota + expect(sidebarQuota?.['non-gemini']).toBeUndefined() + }) + + // ── Fix 1: poller keeps `current` flag ───────────────────────────────────────── + + it('Fix1-current: poller tick keeps current:true on the active account (index 1)', async () => { + // Two accounts; claude-active is index 1. After a poll tick the sidebar + // must carry current:true on account at index 1 and false on index 0. + // Before this fix the poller omitted getActiveIndexByFamily, defaulting + // isAccountCurrent to false for every account. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const accounts: StubAccount[] = [ + { ...makeStubAccount('tok-0'), index: 0 }, + { ...makeStubAccount('tok-1'), index: 1 }, + ] + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const quotaManager = { + refreshAccounts: mock(async (accts: Array<{ refreshToken: string }>) => { + return accts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { + groups: { 'non-gemini': { remainingFraction: 0.5, modelCount: 1 } }, + }, + })) + }), + dispose: mock(async () => {}), + } + + // claude-active = 1, so account at index 1 should be current. + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => + accounts.map((a) => ({ + refreshToken: a.parts.refreshToken, + enabled: true, + addedAt: 0, + lastUsed: 0, + })), + updateQuotaCache: (_i, _g, _t) => { + if (_i === accounts.length - 1) refreshResolve() + }, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 1, gemini: 1 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await Promise.race([ + refreshedP.then(() => new Promise((r) => setTimeout(r, 300))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + await poller.dispose() + + const state = readSidebarState(stateFile) + expect(state.accounts).toHaveLength(2) + const acct0 = state.accounts.find((a) => a.id === 'acct-0') + const acct1 = state.accounts.find((a) => a.id === 'acct-1') + // Index 1 is the active account — must be current. + expect(acct1?.current).toBe(true) + // Index 0 is not active. + expect(acct0?.current).toBe(false) + }) + + // ── Fix 2: tier round-trips through sidebar state file ──────────────────── + + it('Fix2-tier-poller: captured tier appears in sidebar after a poll tick', async () => { + // When the in-memory account carries capturedTierId/capturedTierAt the + // poller must include tier: { id, capturedAt } in the sidebar snapshot. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const capturedAt = Date.now() + const accounts: StubAccount[] = [ + { + ...makeStubAccount('tok-tier'), + index: 0, + capturedTierId: 'enterprise-tier', + capturedPaidTierId: 'g1-pro-tier', + capturedTierAt: capturedAt, + } as StubAccount & { capturedPaidTierId: string }, + ] + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const quotaManager = { + refreshAccounts: mock(async (accts: Array<{ refreshToken: string }>) => { + return accts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { + groups: { 'non-gemini': { remainingFraction: 0.9, modelCount: 1 } }, + }, + })) + }), + dispose: mock(async () => {}), + } + + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { + refreshToken: 'tok-tier', + enabled: true, + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache: (_i, _g, _t) => { + accounts[0]!.cachedQuotaAccountId = quotaAccountIdentity('tok-tier') + refreshResolve() + }, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await Promise.race([ + refreshedP.then(() => new Promise((r) => setTimeout(r, 300))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + await poller.dispose() + + const state = readSidebarState(stateFile) + expect(state.accounts).toHaveLength(1) + const acct = state.accounts[0] + expect(acct?.tier).toBeDefined() + expect(acct?.tier?.id).toBe('enterprise-tier') + expect(acct?.tier?.paidId).toBe('g1-pro-tier') + expect(acct?.tier?.capturedAt).toBe(capturedAt) + }) + + it('keeps a successful quota refresh when the updated account token differs', async () => { + const accounts: StubAccount[] = [makeStubAccount('bare-token')] + const quotaGroups = { + gemini: { remainingFraction: 0.5, modelCount: 1 }, + } + const updateQuotaCache = mock(() => {}) + const quotaManager = { + refreshAccounts: mock(async () => { + accounts[0]!.parts.refreshToken = 'updated-token' + return [ + { + index: 0, + status: 'ok' as const, + quota: { groups: quotaGroups }, + updatedAccount: { + refreshToken: 'updated-token', + addedAt: 0, + lastUsed: 0, + }, + }, + ] + }), + dispose: mock(async () => {}), + } + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { + refreshToken: 'bare-token', + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await new Promise((resolve) => setTimeout(resolve, 300)) + await poller.dispose() + + expect(updateQuotaCache).toHaveBeenCalledWith( + 0, + quotaGroups, + 'updated-token', + ) + }) + + it('Fix2-tier-absent: account without tier yields no tier key in sidebar', async () => { + // An account with no capturedTierId must not emit tier in the snapshot. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const accounts: StubAccount[] = [ + { ...makeStubAccount('tok-notier'), index: 0 }, + ] + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const quotaManager = { + refreshAccounts: mock(async (accts: Array<{ refreshToken: string }>) => + accts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { groups: {} }, + })), + ), + dispose: mock(async () => {}), + } + + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { + refreshToken: 'tok-notier', + enabled: true, + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache: () => { + refreshResolve() + }, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await Promise.race([ + refreshedP.then(() => new Promise((r) => setTimeout(r, 300))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + await poller.dispose() + + const state = readSidebarState(stateFile) + expect(state.accounts).toHaveLength(1) + const acct = state.accounts[0] + // Must be absent (undefined), not {} or { id: undefined }. + expect(acct?.tier).toBeUndefined() + }) + + // ── P2-A: tier resolved via loadAccountTier for existing accounts ────────────── + + it('P2-A: account with managedProjectId (existing account) gets tier after a tick via loadAccountTier', async () => { + // P2-A: existing accounts have managedProjectId packed in their refresh token. + // ensureProjectContext fast-paths on that field and never calls loadCodeAssist, + // so capturedTierId is never populated by the quota fetch path alone. + // The tier-staleness check must call loadAccountTier for the stalest account + // and persist the result via applyUpdatedAccount. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const refreshToken = 'tok-existing-account' + // Simulate an account with managedProjectId but NO captured tier yet. + const accounts: (StubAccount & { + capturedTierId?: string + capturedTierAt?: number + })[] = [ + { + ...makeStubAccount(refreshToken), + index: 0, + // capturedTierId absent -- the bug: existing account never got tier. + }, + ] + + let tierLoadCalled = false + const applyUpdatedAccountCalls: Array<{ + index: number + patch: { capturedTierId?: string; capturedTierAt?: number } + }> = [] + + const quotaManager = makeQuotaManager({ result: 'ok' }) + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { + refreshToken, + managedProjectId: 'proj-abc123', // existing account has managedProjectId + enabled: true, + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache: mock((_i, groups, _t) => { + accounts[0]!.cachedQuotaAccountId = quotaAccountIdentity(refreshToken) + }), + applyUpdatedAccount: mock((index, patch) => { + applyUpdatedAccountCalls.push({ index, patch }) + if (patch.capturedTierId !== undefined) { + accounts[0]!.capturedTierId = patch.capturedTierId + accounts[0]!.capturedTierAt = patch.capturedTierAt + } + }), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + let tierResolve!: () => void + const tierResolved = new Promise((r) => { + tierResolve = r + }) + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + loadAccountTier: async (account) => { + tierLoadCalled = true + // Must be called with the AccountMetadataV3 that has managedProjectId. + expect(account.managedProjectId).toBe('proj-abc123') + const result = { id: 'free-tier', capturedAt: Date.now() } + // Signal after returning so the assertions run AFTER applyUpdatedAccount. + setTimeout(tierResolve, 0) + return result + }, + }) + poller.start() + + await Promise.race([ + tierResolved.then(() => new Promise((r) => setTimeout(r, 300))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout waiting for tier')), 5_000), + ), + ]) + await poller.dispose() + + // The tier loader must have been called. + expect(tierLoadCalled).toBe(true) + // applyUpdatedAccount must have been called with the resolved tier. + const tierCall = applyUpdatedAccountCalls.find( + (c) => c.patch.capturedTierId !== undefined, + ) + expect(tierCall).toBeDefined() + expect(tierCall?.patch.capturedTierId).toBe('free-tier') + expect(typeof tierCall?.patch.capturedTierAt).toBe('number') + }) + + it('P2-A-stale: account with outdated capturedTierAt triggers tier refresh', async () => { + // An account that already has a tier but its capturedTierAt is older than + // TIER_STALENESS_TTL_MS must be re-resolved. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const refreshToken = 'tok-stale-tier' + const oldTierAt = Date.now() - 25 * 60 * 60 * 1_000 // 25h ago (> 24h TTL) + const accounts: (StubAccount & { + capturedTierId?: string + capturedTierAt?: number + })[] = [ + { + ...makeStubAccount(refreshToken), + index: 0, + capturedTierId: 'old-tier', + capturedTierAt: oldTierAt, + }, + ] + + let tierLoadCalled = false + const quotaManager = makeQuotaManager({ result: 'ok' }) + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + let tierResolve!: () => void + const tierResolved = new Promise((r) => { + tierResolve = r + }) + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), // real time -- oldTierAt IS stale + loadAccountTier: async (_account) => { + tierLoadCalled = true + setTimeout(tierResolve, 0) + return { id: 'new-tier', capturedAt: Date.now() } + }, + }) + poller.start() + + await Promise.race([ + tierResolved.then(() => new Promise((r) => setTimeout(r, 100))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + await poller.dispose() + + expect(tierLoadCalled).toBe(true) + }) + + it('backfills a missing paid tier field despite a fresh tier timestamp', async () => { + const now = 1_700_000_000_000 + const refreshToken = 'tok-paid-tier-backfill' + const accounts: Array< + StubAccount & { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number + } + > = [ + { + ...makeStubAccount(refreshToken), + capturedTierId: 'free-tier', + capturedTierAt: now - 60 * 60 * 1_000, + }, + ] + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + paidId: 'g1-pro-tier', + capturedAt: now, + })) + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock((_index, patch) => { + Object.assign(accounts[0]!, patch) + }), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: makeQuotaManager({ + result: 'ok', + }) as unknown as import('./quota').QuotaManager, + loadAccountTier, + now: () => now, + }) + + await (poller as unknown as { refresh: () => Promise }).refresh() + + expect(loadAccountTier).toHaveBeenCalledTimes(1) + expect(accounts[0]?.capturedPaidTierId).toBe('g1-pro-tier') + }) + + it('records a paid-tier-schema attempt when the upstream paid tier is absent', async () => { + const now = 1_700_000_000_000 + const refreshToken = 'tok-empty-paid-tier' + const accounts: Array< + StubAccount & { + capturedTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number + } + > = [ + { + ...makeStubAccount(refreshToken), + capturedTierId: 'free-tier', + capturedTierAt: now - 60 * 60 * 1_000, + }, + ] + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: now, + })) + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock((_index, patch) => { + Object.assign(accounts[0]!, patch) + }), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: makeQuotaManager({ + result: 'ok', + }) as unknown as import('./quota').QuotaManager, + loadAccountTier, + now: () => now, + }) + const refresh = ( + poller as unknown as { refresh: () => Promise } + ).refresh.bind(poller) + + await refresh() + await refresh() + + expect(loadAccountTier).toHaveBeenCalledTimes(1) + expect(accounts[0]?.capturedTierSchemaVersion).toBe(1) + }) + + it('does not re-resolve a fresh tier captured under the paid-tier schema', async () => { + const now = 1_700_000_000_000 + const refreshToken = 'tok-current-tier-schema' + const accounts: Array< + StubAccount & { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number + } + > = [ + { + ...makeStubAccount(refreshToken), + capturedTierId: 'free-tier', + capturedPaidTierId: 'g1-pro-tier', + capturedTierAt: now - 60 * 60 * 1_000, + }, + ] + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + paidId: 'g1-pro-tier', + capturedAt: now, + })) + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: makeQuotaManager({ + result: 'ok', + }) as unknown as import('./quota').QuotaManager, + loadAccountTier, + now: () => now, + }) + + await (poller as unknown as { refresh: () => Promise }).refresh() + + expect(loadAccountTier).not.toHaveBeenCalled() + }) + + it('P2-A-fresh: account with fresh capturedTierAt does NOT trigger tier refresh', async () => { + // An account refreshed 1h ago (< 24h TTL) must not trigger loadAccountTier. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const refreshToken = 'tok-fresh-tier' + const freshTierAt = Date.now() - 1 * 60 * 60 * 1_000 // 1h ago (well within 24h) + const accounts: (StubAccount & { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number + })[] = [ + { + ...makeStubAccount(refreshToken), + index: 0, + capturedTierId: 'free-tier', + capturedPaidTierId: 'g1-pro-tier', + capturedTierAt: freshTierAt, + capturedTierSchemaVersion: 1, + }, + ] + + let tierLoadCalled = false + let refreshDone = false + const quotaManager = { + refreshAccounts: mock(async () => { + refreshDone = true + return [] + }), + dispose: mock(async () => {}), + } + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + loadAccountTier: async () => { + tierLoadCalled = true + return null + }, + }) + poller.start() + + // Wait for the quota refresh to complete (a tick ran). + await Promise.race([ + new Promise((r) => { + const interval = setInterval(() => { + if (refreshDone) { + clearInterval(interval) + r() + } + }, 10) + }), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + // Small extra drain so any async tier check would have fired. + await new Promise((r) => setTimeout(r, 200)) + await poller.dispose() + + // Fresh tier -- tier loader must NOT have been invoked. + expect(tierLoadCalled).toBe(false) + }) +}) + +// ─── Wiring tests (through the real factory) ────────────────────────────────── + +describe('BackgroundQuotaRefresh wiring through createAntigravityPlugin', () => { + let tempDir: string + let stateFile: string + let fetchSpy: ReturnType + let savedSidebarEnvWiring: string | undefined + + beforeEach(() => { + savedSidebarEnvWiring = process.env[SIDEBAR_STATE_ENV] + tempDir = mkdtempSync(join(tmpdir(), 'bg-quota-wiring-')) + stateFile = join(tempDir, 'sidebar-state.json') + process.env[SIDEBAR_STATE_ENV] = stateFile + // Prevent the version check from hitting the network. + const { spyOn } = require('bun:test') as typeof import('bun:test') + fetchSpy = spyOn(globalThis, 'fetch').mockImplementation( + (async () => + new Response('not-found', { status: 404 })) as unknown as typeof fetch, + ) + }) + + afterEach(async () => { + fetchSpy.mockRestore() + mock.restore() + // Restore rather than delete — a delete drops resolution to the real state dir. + if (savedSidebarEnvWiring !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebarEnvWiring + else delete process.env[SIDEBAR_STATE_ENV] + rmSync(tempDir, { recursive: true, force: true }) + }) + + async function makePlugin( + directory: string, + config: Record, + pollerCreatedCb?: (p: BackgroundQuotaRefresh) => void, + ) { + mkdirSync(join(directory, '.opencode'), { recursive: true }) + await Bun.write( + join(directory, '.opencode', 'antigravity.json'), + JSON.stringify(config), + ) + const client = makeMinimalClient() + const input = makePluginInput(client, directory) + return createAntigravityPlugin('google', { + _onPollerCreated: pollerCreatedCb, + })(input) + } + + it('background_quota_refresh: false → NO poller instance is constructed', async () => { + const pluginDir = join(tempDir, 'plugin-disabled') + let pollerCreated = false + const plugin = await makePlugin( + pluginDir, + { background_quota_refresh: false }, + () => { + pollerCreated = true + }, + ) + await plugin.dispose() + // _onPollerCreated must not have been called — the block is guarded by + // `if (config.background_quota_refresh)` and the seam is wired INSIDE it. + expect(pollerCreated).toBe(false) + }) + + it('background_quota_refresh: true → exactly one poller is constructed and its dispose is registered', async () => { + const pluginDir = join(tempDir, 'plugin-enabled') + const capturedPollers: BackgroundQuotaRefresh[] = [] + const plugin = await makePlugin( + pluginDir, + { + background_quota_refresh: true, + background_quota_refresh_interval_minutes: 60, + }, + (p) => capturedPollers.push(p), + ) + // Exactly one instance was constructed. + expect(capturedPollers).toHaveLength(1) + // Replace dispose with a spy BEFORE plugin.dispose() is called so we can + // confirm the plugin's producer plumbing reaches through to the instance. + let pollerDisposeCalled = false + const original = capturedPollers[0]!.dispose.bind(capturedPollers[0]!) + capturedPollers[0]!.dispose = async () => { + pollerDisposeCalled = true + return original() + } + await plugin.dispose() + expect(pollerDisposeCalled).toBe(true) + }) + + it('plugin dispose awaits the poller in the PRODUCER phase (poller stops before the sidebar drain)', async () => { + // Evidence that the poller is a producer (not consumer): the lifecycle + // disposes producers BEFORE calling drainSidebarWrites. We instrument + // the poller's dispose to record its completion timestamp, and confirm + // plugin.dispose() resolves — which requires the producer chain to have + // completed (poller stopped, any in-flight tick awaited). + // + // A more surgical ordering probe would require intercepting lifecycle + // internals; this test focuses on the externally observable guarantee: + // plugin.dispose() fully awaits the poller, and the poller's dispose + // resolves within a tight deadline. + const pluginDir = join(tempDir, 'plugin-producer') + const capturedPollers: BackgroundQuotaRefresh[] = [] + const plugin = await makePlugin( + pluginDir, + { + background_quota_refresh: true, + background_quota_refresh_interval_minutes: 60, + }, + (p) => capturedPollers.push(p), + ) + expect(capturedPollers).toHaveLength(1) + + // Wrap dispose to measure how long the poller took to stop. + let pollerDisposeResolved = false + const orig = capturedPollers[0]!.dispose.bind(capturedPollers[0]!) + capturedPollers[0]!.dispose = async () => { + await orig() + pollerDisposeResolved = true + } + const disposed = await Promise.race([ + plugin.dispose().then(() => 'done'), + new Promise((r) => setTimeout(() => r('timeout'), 5_000)), + ]) + expect(disposed).toBe('done') + // The poller's dispose must have resolved before plugin.dispose() resolved. + expect(pollerDisposeResolved).toBe(true) + }) +}) diff --git a/packages/opencode/src/plugin/background-quota-refresh.ts b/packages/opencode/src/plugin/background-quota-refresh.ts new file mode 100644 index 0000000..4b13624 --- /dev/null +++ b/packages/opencode/src/plugin/background-quota-refresh.ts @@ -0,0 +1,475 @@ +/** + * Background quota poller. + * + * Runs a jittered, `unref()`'d timer that periodically refreshes quota for + * all accounts and pushes ONE sidebar snapshot per tick. Three-layer + * cross-process dedup: + * + * 1. Freshness gate — the only hard correctness bound for quota snapshots. + * Reads `checkedAt` from the on-disk sidebar state; skips quota work when + * it is fresher than + * `max(intervalMs − 60_000, floor(intervalMs / 2))`. The /2 floor + * ensures the gate still fires at the 1-minute minimum interval. + * + * 2. Advisory fenced lock — an optimization that lets a single process do + * the network work when several wake up simultaneously. Held → skip. + * + * 3. Lock throws → FAIL CLOSED (skip this tick, add jitter to next). A + * fallback mutex is NOT used here. The correct reasoning is: the + * freshness gate (layer 1) already bounds how many refreshes can land; + * a secondary "claim marker" file is just a worse re-implementation of + * the fenced lock we just failed to acquire — it replicates the same + * ownership problem without the protocol guarantees, and any transient + * I/O error that made the lock throw will likely affect a claim file + * too. The previous poller attempt fell into this trap across three + * review rounds; this comment exists so a future reviewer stops + * re-litigating it. + */ + +import { createHash } from 'node:crypto' +import type { + AccountMetadataV3, + QuotaGroup, + QuotaGroupSummary, +} from '@cortexkit/antigravity-auth-core' +import { + acquireFencedFileLock, + type FencedFileLock, +} from '@cortexkit/antigravity-auth-core/file-lock' + +import { readSidebarState, toCapturedTier } from '../sidebar-state' +import type { QuotaManager } from './quota' +import { pushSidebarQuotaSnapshot } from './quota' + +/** + * The subset of AccountManager that the poller needs. Using a narrow + * interface keeps the module decoupled from the large concrete class and + * lets tests pass minimal stubs without satisfying 90+ methods. + */ +export interface PollerAccountView { + getAccounts(): Array<{ + index: number + label?: string + enabled?: boolean + parts: { refreshToken: string } + cachedQuota?: Partial> + cachedQuotaAccountId?: string + coolingDownUntil?: number + /** Captured plan tier ID (raw upstream string). */ + capturedTierId?: string + /** Captured paid-tier ID (raw upstream string). */ + capturedPaidTierId?: string + /** Epoch ms when capturedTierId was last recorded. */ + capturedTierAt?: number + /** Schema version of the most recent tier capture. */ + capturedTierSchemaVersion?: number + }> + /** Must return the full AccountMetadataV3 shape so the quota manager + * can use token-refresh and project-context fields inside each fetch. */ + getAccountsForQuotaCheck(): AccountMetadataV3[] + updateQuotaCache( + accountIndex: number, + quotaGroups: Partial>, + expectedRefreshToken?: string, + ): void + /** + * Apply a subset of quota-result fields (e.g. captured tier) onto the + * in-memory account record without touching the quota cache itself. + */ + applyUpdatedAccount( + accountIndex: number, + patch: { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number + }, + expectedRefreshToken?: string, + ): void + requestSaveToDisk(): void + /** + * Returns the active account index per model family so the sidebar + * snapshot can stamp the `current` flag on the right account. + */ + getActiveIndexByFamily(): { claude: number; gemini: number } +} + +/** + * Opaque identity derived from a refresh token. Three callers keep + * independent copies: this module, `command-data.ts` (ships into the TUI + * compiled tree and cannot reach the core barrel), and `index.ts`. All + * three must produce the same 16-char hex prefix. + */ +function quotaAccountIdentity(refreshToken: string): string { + return createHash('sha256').update(refreshToken).digest('hex').slice(0, 16) +} + +/** Minimum poll interval (1 minute, matching the config schema min). */ +const MIN_INTERVAL_MS = 60_000 + +/** Spread startup jitter so N processes don't all fire at once on launch. */ +const STARTUP_JITTER_MS = 30_000 + +/** Additional jitter added to the next tick when the lock mechanism throws. */ +const ERROR_JITTER_MS = 15_000 + +/** TTL for the advisory poll lock. Long enough to cover a slow network fetch. */ +const POLL_LOCK_TTL_MS = 60_000 + +/** Maximum age of a captured tier before the poller refreshes it. 24 h. */ +const TIER_STALENESS_TTL_MS = 24 * 60 * 60 * 1_000 + +/** A capture at this version distinguishes an absent paid tier from old data. */ +const CAPTURED_TIER_SCHEMA_VERSION = 1 + +export interface BackgroundQuotaRefreshOptions { + intervalMs: number + sidebarStateFile: string + /** + * Supplier for the live account pool. Called inside each locked phase, so it + * observes any concurrent add/remove that completed while waiting for the + * lock. + */ + getAccountManager: () => PollerAccountView | null + quotaManager: QuotaManager + /** + * Optional tier-refresh callback. When provided, the poller calls this + * once per tick for the account whose capturedTierAt is missing or older + * than TIER_STALENESS_TTL_MS. The callback handles its own token-refresh + * and must resolve to a tier object (or null on failure / nothing to + * report) -- failure must never reject, since tier is best-effort. + * + * Wired in index.ts through the same agyTransport seam the quota manager + * uses, so e2e tests can inject a mock alongside the mock quota fetch. + */ + loadAccountTier?: ( + account: AccountMetadataV3, + ) => Promise<{ id: string; paidId?: string; capturedAt: number } | null> + /** + * Clock / randomness seam. Defaults to the system clock. + */ + now?: () => number + random?: () => number +} + +/** + * Per-loader-instance background quota poller. + * + * Must be created after the plugin has initialised (so `getAccountManager()` + * is populated) and disposed via the producer lifecycle phase. + */ +export class BackgroundQuotaRefresh { + private readonly intervalMs: number + private readonly sidebarStateFile: string + private readonly getAccountManager: () => PollerAccountView | null + private readonly quotaManager: QuotaManager + private readonly loadAccountTier?: ( + account: AccountMetadataV3, + ) => Promise<{ id: string; paidId?: string; capturedAt: number } | null> + private readonly now: () => number + private readonly random: () => number + + private timer: ReturnType | null = null + /** Resolves when the currently-running tick completes (or immediately if none). */ + private inFlight: Promise | null = null + private disposed = false + + constructor(options: BackgroundQuotaRefreshOptions) { + this.intervalMs = options.intervalMs + this.sidebarStateFile = options.sidebarStateFile + this.getAccountManager = options.getAccountManager + this.quotaManager = options.quotaManager + this.loadAccountTier = options.loadAccountTier + this.now = options.now ?? (() => Date.now()) + this.random = options.random ?? Math.random + } + + /** Start the background timer. Idempotent: a second call is a no-op. */ + start(): void { + if (this.disposed || this.timer !== null) return + this.scheduleNext(this.jitteredStartDelay()) + } + + /** Stop the timer and await any in-flight tick. */ + async dispose(): Promise { + if (this.disposed) return + this.disposed = true + if (this.timer !== null) { + clearTimeout(this.timer) + this.timer = null + } + await this.inFlight + } + + // ── internals ───────────────────────────────────────────────────────────── + + private jitteredStartDelay(): number { + return Math.floor(this.random() * STARTUP_JITTER_MS) + } + + private jitteredInterval(extraJitterMs = 0): number { + // ±10% on top of the base interval prevents phase-locking across restarts. + const jitter = Math.floor(this.random() * this.intervalMs * 0.1) + return this.intervalMs + jitter + extraJitterMs + } + + private scheduleNext(delayMs: number): void { + if (this.disposed) return + this.timer = setTimeout(() => { + this.timer = null + // Re-entry guard: if the previous tick is still in flight (refresh + // duration > interval), skip this timer callback. The in-flight tick + // calls scheduleNext itself when it completes, so the poller continues. + // Without this guard, assigning this.inFlight here would overwrite the + // first tick's promise and dispose() would await the wrong one. + if (this.inFlight !== null) return + this.inFlight = this.runTick() + .catch(() => { + // Errors inside runTick are already caught per-phase. An unexpected + // throw here (before any internal reschedule) means no next tick + // was queued — reschedule with error jitter so the poller degrades + // to a delayed tick rather than dying permanently. + this.scheduleNext(this.jitteredInterval(ERROR_JITTER_MS)) + }) + .finally(() => { + this.inFlight = null + }) + }, delayMs) + // Unreffing lets the process exit while the timer is pending — the poller + // must never be the reason an idle session stays alive. + if (this.timer.unref) this.timer.unref() + } + + private async runTick(): Promise { + // ── 1. Quota freshness gate ────────────────────────────────────────── + // + // Bound: `max(intervalMs − 60_000, floor(intervalMs / 2))`. + // + // - intervalMs − 60_000: at 5-minute intervals this is 240 s — tight + // enough to prevent double refreshes while allowing drift. + // - floor(intervalMs / 2): the fallback that keeps the gate non-trivial + // at the 1-minute minimum (floor(60000/2) = 30 000 ms, i.e. 30 s). + const freshnessThresholdMs = Math.max( + this.intervalMs - MIN_INTERVAL_MS, + Math.floor(this.intervalMs / 2), + ) + const state = readSidebarState(this.sidebarStateFile) + const ageMs = this.now() - state.checkedAt + + // ── 2. Advisory fenced lock ────────────────────────────────────────── + let lock: FencedFileLock | null = null + try { + lock = await acquireFencedFileLock({ + path: this.sidebarStateFile, + name: 'bg-quota-poll', + ttlMs: POLL_LOCK_TTL_MS, + // No renewal: the tick completes well within the TTL, and a renewal + // timer would need its own dispose chain. + renew: false, + }) + } catch { + // Lock mechanism threw (I/O error, filesystem issue). FAIL CLOSED: + // skip this tick. See module-level invariant comment for why we do not + // invent a fallback mutex here. + this.scheduleNext(this.jitteredInterval(ERROR_JITTER_MS)) + return + } + + if (lock === null) { + // Another process holds the lock — its tier phase owns this tick's + // metadata work as well as any quota snapshot write. + this.scheduleNext(this.jitteredInterval()) + return + } + + try { + if (ageMs >= freshnessThresholdMs) await this.refresh() + else await this.refreshTierSlot() + } finally { + await lock.release().catch(() => { + // Release is best-effort; the TTL expires it if we can't reach it. + }) + } + + if (!this.disposed) { + this.scheduleNext(this.jitteredInterval()) + } + } + + private async refresh(): Promise { + const refreshedQuota = await this.refreshQuota() + await this.refreshTierSlot() + if (refreshedQuota) await this.pushSnapshot() + } + + private async refreshQuota(): Promise { + if (this.disposed) return false + + const manager = this.getAccountManager() + if (!manager) return false + + const accounts = manager.getAccountsForQuotaCheck() + if (accounts.length === 0) return false + + // Refresh all accounts, then stamp the live manager once per result. + // Using the shared quotaManager deduplicates in-flight refreshes with + // any concurrent modal or fetch-interceptor refresh. + let results: Awaited> + try { + results = await this.quotaManager.refreshAccounts(accounts, { + indexFor: (account) => accounts.indexOf(account), + // Do not force: the quota manager's per-account dedup and backoff + // apply — an account that already refreshed recently is skipped. + force: false, + }) + } catch { + return false + } + + if (this.disposed) return false + + // Re-resolve live indices after the await: a concurrent add/remove may + // have shifted positions. + const liveAccounts = manager.getAccounts() + const liveIndexByToken = new Map() + for (const entry of liveAccounts) { + liveIndexByToken.set(entry.parts.refreshToken, entry.index) + } + + let anyUpdated = false + for (let i = 0; i < results.length; i++) { + const result = results[i] + if (result?.status !== 'ok' || !result.quota?.groups) continue + const candidateTokens = [ + result.updatedAccount?.refreshToken, + accounts[i]?.refreshToken, + ].filter((token): token is string => Boolean(token)) + const liveIndex = candidateTokens + .map((token) => liveIndexByToken.get(token)) + .find((index): index is number => index !== undefined) + if (liveIndex === undefined) continue + const liveToken = liveAccounts[liveIndex]?.parts.refreshToken + if (!liveToken) continue + manager.updateQuotaCache(liveIndex, result.quota.groups, liveToken) + // Persist tier captured during ensureProjectContext (same tick, no extra + // network call) so the sidebar snapshot carries it after every poll. + const uaTier = toCapturedTier(result.updatedAccount ?? {}) + if (uaTier !== undefined) { + manager.applyUpdatedAccount( + liveIndex, + { + capturedTierId: uaTier.id, + ...(uaTier.paidId !== undefined + ? { capturedPaidTierId: uaTier.paidId } + : {}), + capturedTierAt: uaTier.capturedAt, + capturedTierSchemaVersion: CAPTURED_TIER_SCHEMA_VERSION, + }, + liveToken, + ) + } + anyUpdated = true + } + if (anyUpdated) manager.requestSaveToDisk() + + return true + } + + private async refreshTierSlot(): Promise { + if (this.disposed || !this.loadAccountTier) return + + const manager = this.getAccountManager() + if (!manager) return + + const liveAccounts = manager.getAccounts() + const accounts = manager.getAccountsForQuotaCheck() + const accountByToken = new Map( + accounts.map((account) => [account.refreshToken, account]), + ) + const now = this.now() + let stalestEntry: (typeof liveAccounts)[number] | undefined + for (const entry of liveAccounts) { + const needsPaidTierBackfill = + entry.capturedPaidTierId === undefined && + entry.capturedTierSchemaVersion !== CAPTURED_TIER_SCHEMA_VERSION + const isTierStale = + entry.capturedTierAt === undefined || + now - entry.capturedTierAt >= TIER_STALENESS_TTL_MS || + needsPaidTierBackfill + if ( + isTierStale && + (stalestEntry === undefined || + (entry.capturedTierAt ?? 0) < (stalestEntry.capturedTierAt ?? 0)) + ) { + stalestEntry = entry + } + } + + if (stalestEntry === undefined) return + + const refreshToken = stalestEntry.parts.refreshToken + const accountForTier = accountByToken.get(refreshToken) + if (!accountForTier) return + + try { + const tier = await this.loadAccountTier(accountForTier) + if (tier === null || this.disposed) return + const liveEntry = manager + .getAccounts() + .find((entry) => entry.parts.refreshToken === refreshToken) + if (!liveEntry) return + manager.applyUpdatedAccount( + liveEntry.index, + { + capturedTierId: tier.id, + ...(tier.paidId !== undefined + ? { capturedPaidTierId: tier.paidId } + : {}), + capturedTierAt: tier.capturedAt, + capturedTierSchemaVersion: CAPTURED_TIER_SCHEMA_VERSION, + }, + refreshToken, + ) + manager.requestSaveToDisk() + } catch { + // Tier is best-effort; a failure must never break the tick. + } + } + + private async pushSnapshot(): Promise { + // ONE sidebar write per poll. `pushSidebarQuotaSnapshot` reads the live + // account view (with the just-updated cached values) and writes atomically. + const getAccounts = () => { + const m = this.getAccountManager() + return m ? m.getAccounts() : null + } + + await pushSidebarQuotaSnapshot( + () => { + const accts = getAccounts() + if (!accts) return null + return accts.map((entry) => ({ + index: entry.index, + label: entry.label, + enabled: entry.enabled, + coolingDownUntil: entry.coolingDownUntil, + cachedQuota: entry.cachedQuota, + cachedQuotaAccountId: entry.cachedQuotaAccountId, + // Derive current identity from the LIVE refresh token, independent + // of whatever stamp the cached snapshot carries. Mirrors index.ts's + // two correct call sites so the sidebar can detect a stale snapshot + // after an account reorder. + currentQuotaAccountId: quotaAccountIdentity(entry.parts.refreshToken), + tier: toCapturedTier(entry), + })) + }, + 0, + () => { + const m = this.getAccountManager() + return m ? m.getActiveIndexByFamily() : null + }, + ).catch(() => { + // Sidebar write failure is not fatal; the next tick will retry. + }) + } +} diff --git a/packages/opencode/src/plugin/command-data.test.ts b/packages/opencode/src/plugin/command-data.test.ts index 453607c..6d955fa 100644 --- a/packages/opencode/src/plugin/command-data.test.ts +++ b/packages/opencode/src/plugin/command-data.test.ts @@ -74,6 +74,9 @@ interface AccountFixture { cachedQuotaAccountId?: string accountIneligible?: boolean coolingDownUntil?: number + healthScore?: number + capturedTierId?: string + capturedTierAt?: number } function makeAccountFixture( @@ -226,6 +229,9 @@ function makeHarness(options: { cachedQuotaAccountId: entry.cachedQuotaAccountId, accountIneligible: entry.accountIneligible, coolingDownUntil: entry.coolingDownUntil, + healthScore: entry.healthScore, + capturedTierId: entry.capturedTierId, + capturedTierAt: entry.capturedTierAt, })) }, getAccountsForQuotaCheck() { @@ -366,14 +372,19 @@ async function readSidebar(stateFile: string): Promise { } let dir: string +let savedSidebarEnv: string | undefined describe('createCommandDataService', () => { beforeEach(() => { + savedSidebarEnv = process.env[SIDEBAR_STATE_ENV] dir = mkdtempSync(join(tmpdir(), 'agy-command-data-')) }) afterEach(() => { - delete process.env[SIDEBAR_STATE_ENV] + // Restore rather than delete — a delete drops resolution to the real state dir. + if (savedSidebarEnv !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebarEnv + else delete process.env[SIDEBAR_STATE_ENV] rmSync(dir, { recursive: true, force: true }) }) @@ -517,21 +528,22 @@ describe('createCommandDataService', () => { }) }) - it('silently skips unknown quota keys in a legacy snapshot (tolerant read)', async () => { - // Older Antigravity revisions persisted quota under keys we no - // longer render (`claude`, `gpt-4`, ad-hoc pool names). The - // projection must IGNORE unknown keys rather than crash the - // dialog or surface them as keys without a label. + it('normalizes legacy quota keys and silently skips truly unknown keys', async () => { + // S1 + N3: the legacy keys (claude → non-gemini, gpt-oss → non-gemini, + // gemini-pro/gemini-flash → gemini) are now NORMALIZED at read time, + // not silently dropped. Truly unknown keys (e.g. 'gpt-4') are still + // dropped. The canonical 'gemini' key carries through unchanged. const harness = makeHarness({ accounts: [ makeAccountFixture({ refreshToken: 'refresh-a', label: 'Account A', cachedQuota: { - // Legacy / unknown keys that the renderer doesn't support. + // Legacy key: maps to non-gemini with MIN fraction. claude: { remainingFraction: 0.9 }, + // Truly unknown key: must be dropped (no mapping). 'gpt-4': { remainingFraction: 0.8 }, - // Supported key — must still render. + // Canonical key: must render as-is. gemini: { remainingFraction: 0.5 }, } as unknown as QuotaGroupFixture, }), @@ -541,15 +553,14 @@ describe('createCommandDataService', () => { const rows = await harness.service.listAccounts() expect(rows).toHaveLength(1) - expect(rows[0]?.quota).toEqual([ - { - key: 'gemini', - label: 'Gemini', - remainingPercent: 50, - resetAt: undefined, - windows: undefined, - }, - ]) + // Both canonical and normalized-from-legacy keys must render. + const quota = rows[0]?.quota + expect(quota?.find((q) => q.key === 'gemini')?.remainingPercent).toBe(50) + expect(quota?.find((q) => q.key === 'non-gemini')?.remainingPercent).toBe( + 90, + ) + // The truly unknown 'gpt-4' key must not appear in the output. + expect(quota?.find((q) => q.key === ('gpt-4' as never))).toBeUndefined() }) it('refreshQuota() fetches every account, including an uncached new account, through the shared quota manager', async () => { @@ -849,6 +860,24 @@ describe('createCommandDataService', () => { expect(harness.storage.accounts[1]?.enabled).toBe(true) }) + it('toggleAccountEnabled() carries healthScore through the command row seam', async () => { + // An account whose tracker score is 42 must keep that score after + // a command mutate — the row seam must forward healthScore so the + // sidebar does not silently reset to the 100 default. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-a', + label: 'Alpha', + healthScore: 42, + }), + ], + }) + const rows = await harness.service.toggleAccountEnabled(0) + expect(rows).not.toBeNull() + expect(rows?.[0]?.healthScore).toBe(42) + }) + it('toggleAccountEnabled() rejects out-of-range indices without touching storage', async () => { const harness = makeHarness({ accounts: [ @@ -1433,6 +1462,199 @@ describe('createCommandDataService', () => { expect(harness.storage.activeIndexByFamily?.claude).toBe(1) expect(harness.storage.activeIndexByFamily?.gemini).toBe(1) }) + + it('refreshQuotaRespectingBackoff() folds successful results into the live cache', async () => { + // A stale account whose cached quota is 10% non-gemini. After a + // successful non-forced refresh returning 80%, the live cache must + // carry the fresh 80% — NOT the stale 10%. + const refreshResults = new Map([ + [ + 'refresh-a', + { + index: 0, + status: 'ok', + quota: { + groups: { + 'non-gemini': { + remainingFraction: 0.8, + resetTime: new Date(0).toISOString(), + modelCount: 1, + }, + }, + modelCount: 1, + }, + updatedAccount: { + refreshToken: 'refresh-a', + addedAt: 0, + lastUsed: 0, + }, + }, + ], + ]) + + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-a', + label: 'Alpha', + cachedQuota: { 'non-gemini': { remainingFraction: 0.1 } }, + cachedQuotaUpdatedAt: 1, + }), + ], + refreshResults, + }) + + // Live cache carries stale 10% before the call. + const before = harness.liveView[0]?.cachedQuota as + | QuotaGroupFixture + | undefined + expect(before?.['non-gemini']?.remainingFraction).toBe(0.1) + + await harness.service.refreshQuotaRespectingBackoff() + + // Live cache updated to fresh 80%. + const after = harness.liveView[0]?.cachedQuota as + | QuotaGroupFixture + | undefined + expect(after?.['non-gemini']?.remainingFraction).toBe(0.8) + // Identity stamp must be set. + expect(harness.liveView[0]?.cachedQuotaAccountId).toBe( + quotaAccountIdentity('refresh-a'), + ) + // Must have been a non-forced refresh. + expect(harness.quotaCallLog.every((call) => call.force === false)).toBe( + true, + ) + }) + + it('refreshQuotaRespectingBackoff() preserves existing backoff behaviour when all accounts are fresh', async () => { + // When refreshResults is empty (no stubs match — simulating the quota + // manager skipping every account due to backoff/freshness), the + // method must NOT crash, NOT update the cache, and NOT write a + // different sidebar. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'fresh-account', + label: 'Fresh', + cachedQuota: { 'non-gemini': { remainingFraction: 0.5 } }, + }), + ], + refreshResults: new Map(), // Empty → every result is 'error' + }) + + const beforeFraction = ( + harness.liveView[0]?.cachedQuota as QuotaGroupFixture | undefined + )?.['non-gemini']?.remainingFraction + + await harness.service.refreshQuotaRespectingBackoff() + + // Stale cache unchanged — no successful refresh occurred. + const afterFraction = ( + harness.liveView[0]?.cachedQuota as QuotaGroupFixture | undefined + )?.['non-gemini']?.remainingFraction + expect(afterFraction).toBe(beforeFraction) + // Identity stamp must NOT have been set (no update happened). + expect(harness.liveView[0]?.cachedQuotaAccountId).toBeUndefined() + }) +}) + +describe('P2-B: writeSidebar lands tier in the state FILE', () => { + let dir: string + let savedSidebarEnvTier: string | undefined + + beforeEach(() => { + savedSidebarEnvTier = process.env[SIDEBAR_STATE_ENV] + dir = mkdtempSync(join(tmpdir(), 'agy-tier-sidebar-')) + }) + + afterEach(() => { + if (savedSidebarEnvTier !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebarEnvTier + else delete process.env[SIDEBAR_STATE_ENV] + rmSync(dir, { recursive: true, force: true }) + }) + + it('P2-B-file: refreshQuota writes tier to sidebar state file (not just the returned row)', async () => { + // Previously writeSidebar built SidebarAccountRedactionInput WITHOUT tier, + // so the field landed in the CommandAccountRow but was thrown away before + // reaching the disk. This test asserts the state FILE carries it. + const capturedAt = 1_700_000_001_000 + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-tier-test', + capturedTierId: 'free-tier', + capturedTierAt: capturedAt, + cachedQuota: { 'non-gemini': { remainingFraction: 0.5 } }, + cachedQuotaAccountId: quotaAccountIdentity('refresh-tier-test'), + }), + ], + refreshResults: new Map([ + [ + 'refresh-tier-test', + { + index: 0, + status: 'ok' as const, + quota: { + groups: { + 'non-gemini': { remainingFraction: 0.6, modelCount: 1 }, + }, + modelCount: 1, + }, + updatedAccount: { + refreshToken: 'refresh-tier-test', + addedAt: 0, + lastUsed: 0, + }, + }, + ], + ]), + }) + + await harness.service.refreshQuota() + const state = await readSidebar(harness.stateFile) + + expect(state.accounts).toHaveLength(1) + const acct = state.accounts[0] + // Tier must reach the FILE, not just the returned row. + expect(acct?.tier).toBeDefined() + expect(acct?.tier?.id).toBe('free-tier') + expect(acct?.tier?.capturedAt).toBe(capturedAt) + }) + + it('P2-B-absent: account without tier produces no tier key in the sidebar file', async () => { + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-no-tier', + cachedQuota: { 'non-gemini': { remainingFraction: 0.5 } }, + cachedQuotaAccountId: quotaAccountIdentity('refresh-no-tier'), + }), + ], + refreshResults: new Map([ + [ + 'refresh-no-tier', + { + index: 0, + status: 'ok' as const, + quota: { groups: {}, modelCount: 0 }, + updatedAccount: { + refreshToken: 'refresh-no-tier', + addedAt: 0, + lastUsed: 0, + }, + }, + ], + ]), + }) + + await harness.service.refreshQuota() + const state = await readSidebar(harness.stateFile) + + expect(state.accounts).toHaveLength(1) + expect(state.accounts[0]?.tier).toBeUndefined() + }) }) // Keep `mock` import alive for symmetry with sibling suites. diff --git a/packages/opencode/src/plugin/command-data.ts b/packages/opencode/src/plugin/command-data.ts index e517ca4..eb72733 100644 --- a/packages/opencode/src/plugin/command-data.ts +++ b/packages/opencode/src/plugin/command-data.ts @@ -44,8 +44,10 @@ import { createHash } from 'node:crypto' import { buildSidebarMachineStateFromAccounts, + normalizeLegacyCachedQuota, type SidebarAccountRedactionInput, setSidebarMachineState, + toCapturedTier, } from '../sidebar-state' /** @@ -143,6 +145,8 @@ export interface CommandAccountRow { * without relying on unstable numeric indices. */ coolingDownUntil?: number + /** Health score in [0, 100], forwarded from the live tracker. */ + healthScore?: number quota: Array<{ key: 'gemini' | 'non-gemini' label: string @@ -154,6 +158,8 @@ export interface CommandAccountRow { resetAt?: number }> }> + /** Captured plan tier. Absent when unknown; never defaulted. */ + tier?: { id: string; capturedAt: number } } /** @@ -186,6 +192,13 @@ interface LiveAccountSnapshot { cachedQuotaAccountId?: string accountIneligible?: boolean coolingDownUntil?: number + healthScore?: number + /** Captured plan tier from the most recent loadCodeAssist response. */ + capturedTierId?: string + /** Captured paid-tier ID from the most recent loadCodeAssist response. */ + capturedPaidTierId?: string + /** Epoch ms when capturedTierId was last recorded. */ + capturedTierAt?: number } function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { @@ -193,11 +206,19 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { // (the refresh token changed, or an index shift placed another account's // snapshot at this position). Drop the stale cache rather than rendering // the wrong account's quota percentages. - const cached = + const rawCached = entry.cachedQuotaAccountId && entry.cachedQuotaAccountId !== quotaAccountIdentity(entry.refreshToken) ? undefined : entry.cachedQuota + // Normalize legacy pool keys (gemini-pro, gemini-flash, claude, gpt-oss) + // to canonical keys before the SUPPORTED_QUOTA_KEYS loop. Without this + // a legacy snapshot produces empty quota rows, and writeSidebar would + // permanently re-emit legacy keys on every dialog open until the next + // real quota refresh rewrites the on-disk snapshot. + const cached = normalizeLegacyCachedQuota( + rawCached as Parameters[0], + ) const quota: CommandAccountRow['quota'] = [] for (const key of SUPPORTED_QUOTA_KEYS) { const cachedEntry = cached?.[key] @@ -246,7 +267,9 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { enabled: entry.enabled, current: entry.active, coolingDownUntil: entry.coolingDownUntil, + healthScore: entry.healthScore, quota, + tier: toCapturedTier(entry), } } @@ -399,8 +422,20 @@ export interface CommandDataService { * refresh token, bump `cachedQuotaUpdatedAt`, and push a label-only * sidebar snapshot. Returns the freshly persisted rows so the * dialog can re-render in place. + * + * Bypass flag is intentional: explicit user-triggered refreshes + * ("Refresh" button, direct slash-command with `refresh` argument) + * must always fetch, even if the quota manager has backed off. */ refreshQuota(): Promise + /** + * Non-forced quota check that RESPECTS the quota manager's per-account + * backoff. Used by the dialog OPEN path — which is a VIEW, not a user + * request to fetch. Accounts already fresh (or in backoff) are skipped; + * the result is discarded (the sidebar poller carries the freshness + * guarantee for idle accounts). + */ + refreshQuotaRespectingBackoff(): Promise /** * Pin `index` as the active account for every family the dialog * tracks (claude + gemini). Mutates the live AccountManager AND the @@ -452,54 +487,63 @@ export function createCommandDataService( const projectRows = (): CommandAccountRow[] => accountManagerView.getAccounts().map(toCommandAccountRow) - const writeSidebar = (rows: CommandAccountRow[]): void => { - const accounts: SidebarAccountRedactionInput[] = rows.map((row) => { - const gemini = row.quota.find((q) => q.key === 'gemini') - const nonGemini = row.quota.find((q) => q.key === 'non-gemini') - // Map a CommandAccountRow quota entry → cachedQuota pool shape - // so projectQuotaPoolForSidebar (the canonical projection seam) - // carries the per-window breakdown into the sidebar state. - const toPool = ( - q: - | { - remainingPercent: number | null - resetAt?: number - windows?: CommandAccountRow['quota'][number]['windows'] - } - | undefined, - ) => { - if (!q || q.remainingPercent == null) return undefined - return { - remainingFraction: q.remainingPercent / 100, - resetTime: - typeof q.resetAt === 'number' && Number.isFinite(q.resetAt) - ? new Date(q.resetAt).toISOString() - : undefined, - windows: q.windows?.map((w) => ({ - window: w.window, - remainingFraction: w.remainingPercent / 100, - resetTime: - typeof w.resetAt === 'number' && Number.isFinite(w.resetAt) - ? new Date(w.resetAt).toISOString() - : '', - })), - } - } + // Map a CommandAccountRow back to the SidebarAccountRedactionInput shape. + // Extracted to prevent the "sixth dropped field" pattern: every field in + // CommandAccountRow that must reach the sidebar lives here once, not as an + // inline literal that a future author has to remember to update in two places. + const toRedactionInput = ( + row: CommandAccountRow, + ): SidebarAccountRedactionInput => { + const gemini = row.quota.find((q) => q.key === 'gemini') + const nonGemini = row.quota.find((q) => q.key === 'non-gemini') + // Map a CommandAccountRow quota entry to cachedQuota pool shape + // so projectQuotaPoolForSidebar (the canonical projection seam) + // carries the per-window breakdown into the sidebar state. + const toPool = ( + q: + | { + remainingPercent: number | null + resetAt?: number + windows?: CommandAccountRow['quota'][number]['windows'] + } + | undefined, + ) => { + if (!q || q.remainingPercent == null) return undefined return { - index: row.index, - label: row.label, - enabled: row.enabled, - current: row.current, - cachedQuota: { - gemini: toPool(gemini), - 'non-gemini': toPool(nonGemini), - }, - // The stamp check has already been done by `toCommandAccountRow` - // before the rows reach this writer; nothing further for the - // sidebar projection to validate here. + remainingFraction: q.remainingPercent / 100, + resetTime: + typeof q.resetAt === 'number' && Number.isFinite(q.resetAt) + ? new Date(q.resetAt).toISOString() + : undefined, + windows: q.windows?.map((w) => ({ + window: w.window, + remainingFraction: w.remainingPercent / 100, + resetTime: + typeof w.resetAt === 'number' && Number.isFinite(w.resetAt) + ? new Date(w.resetAt).toISOString() + : '', + })), } - }) - // Fire-and-forget — the sidebar writer is fenced by its own queue, + } + return { + index: row.index, + label: row.label, + enabled: row.enabled, + current: row.current, + healthScore: row.healthScore, + // Tier passes the redaction boundary unchanged (plan metadata, not PII). + // The stamp check was already done by toCommandAccountRow. + tier: row.tier, + cachedQuota: { + gemini: toPool(gemini), + 'non-gemini': toPool(nonGemini), + }, + } + } + + const writeSidebar = (rows: CommandAccountRow[]): void => { + const accounts = rows.map(toRedactionInput) + // Fire-and-forget -- the sidebar writer is fenced by its own queue, // so a transient lock contention cannot block the dialog response. void setSidebarMachineState( buildSidebarMachineStateFromAccounts(accounts, { checkedAt: now() }), @@ -611,6 +655,77 @@ export function createCommandDataService( return rows }, + async refreshQuotaRespectingBackoff() { + const accountsForQuota = accountManagerView.getAccountsForQuotaCheck() + if (accountsForQuota.length === 0) return + + // Non-forced: the quota manager's per-account backoff and dedup + // apply. Accounts that are fresh or in backoff are skipped. Unlike + // the earlier fire-and-forget path, successful results ARE folded + // into the live cache so the sidebar snapshot reflects fresh data + // immediately — the dialog renders ONLY from the sidebar file, and + // the quota manager's push path does not update AccountManager + // entries for this caller. + let results: Awaited> + try { + results = await quotaManager.refreshAccounts(accountsForQuota, { + indexFor: (account) => accountsForQuota.indexOf(account), + force: false, + }) + } catch { + // Non-critical: dialog open must never fail because a background + // refresh check encountered backoff or a transient error. + return + } + + // Build updates akin to the forced refresh path — resolve live + // indexes after the network call (a concurrent add/remove may have + // shifted positions), fold successful results into the live cache, + // then push a sidebar snapshot. + const updates: Array<{ + refreshToken: string + groups?: Partial< + Record + > + }> = [] + for (const result of results) { + const refreshToken = + result.updatedAccount?.refreshToken ?? + accountsForQuota[result.index]?.refreshToken + if (!refreshToken) continue + const groups = + result.status === 'ok' && result.quota?.groups + ? result.quota.groups + : undefined + updates.push({ refreshToken, groups }) + } + + if (updates.length === 0) return + + const liveIndexByRefreshToken = new Map() + for (const entry of accountManagerView.getAccounts()) { + liveIndexByRefreshToken.set(entry.refreshToken, entry.index) + } + let liveQuotaChanged = false + for (const update of updates) { + const liveIndex = liveIndexByRefreshToken.get(update.refreshToken) + if (liveIndex === undefined || !update.groups) continue + accountManagerView.updateQuotaCache( + liveIndex, + update.groups, + update.refreshToken, + ) + liveQuotaChanged = true + } + if (liveQuotaChanged) accountManagerView.requestSaveToDisk() + + // Push a fresh sidebar snapshot — the live view now carries the + // just-refreshed quota plus any skipped (backoff/fresh) accounts + // whose cached percentages are unchanged. + const rows = projectRows() + writeSidebar(rows) + }, + async setCurrentAccount(index) { return mutateLiveAndStorage({ action: 'setCurrent', index }) }, @@ -738,7 +853,12 @@ export function createCommandDataService( // (the core's buildStorageSnapshot clamps negative sentinels // to 0, and auth-doctor treats a negative activeIndex as // corruption — cf. auth-doctor.ts:149-156). - return index < nextAccounts.length ? index : 0 + // + // Use tokenIdx (the storage-snapshot position) rather than the + // outer-scope `index` (live-read position): a concurrent mutation + // between the live read and the lock can diverge the two, and the + // storage-level cursor must match the storage-level array. + return tokenIdx < nextAccounts.length ? tokenIdx : 0 } const legacyClaude = current.activeIndexByFamily?.claude ?? current.activeIndex diff --git a/packages/opencode/src/plugin/commands.test.ts b/packages/opencode/src/plugin/commands.test.ts index 81ca4b3..af0fa60 100644 --- a/packages/opencode/src/plugin/commands.test.ts +++ b/packages/opencode/src/plugin/commands.test.ts @@ -245,6 +245,7 @@ describe('createCommandExecuteBefore', () => { }, ], refreshQuota: async () => [], + refreshQuotaRespectingBackoff: async () => {}, } as never, { isTuiConnected: () => true }, ) @@ -751,9 +752,137 @@ describe('applyCommand', () => { expect(state.accounts[0]?.quota['non-gemini']?.remainingPercent).toBe(50) expect(state.accounts[1]?.quota).toEqual({}) } finally { - if (previousSidebarFile === undefined) { - delete process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE - } else { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + + it('V2 windows: per-window quota arrays survive the post-apply sidebar projection', async () => { + // After an account or quota apply, createSidebarRefresher rebuilds + // cachedQuota from CommandAccountRow.quota. Without the fix the + // `windows` array was not carried and the sidebar collapsed to a + // single aggregate bar after every apply. + const sidebarFile = join(dir, 'sidebar-windows-v2.json') + const previousSidebarFile = process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = sidebarFile + try { + const resetAt = Date.now() + 7 * 24 * 60 * 60_000 + const dialogRows: CommandAccountRow[] = [ + { + id: 'acct-0', + index: 0, + label: 'Account 1', + enabled: true, + current: true, + quota: [ + { + key: 'non-gemini' as const, + label: 'Non-Gemini', + remainingPercent: 60, + resetAt, + windows: [ + { + window: '5h' as const, + remainingPercent: 85, + resetAt: resetAt - 1000, + }, + { window: 'weekly' as const, remainingPercent: 60, resetAt }, + ], + }, + ], + }, + ] + + const refresher = createSidebarRefresher(() => []) + await refresher(dialogRows) + + const state = readSidebarState(sidebarFile) + expect(state.accounts).toHaveLength(1) + const quota = state.accounts[0]?.quota['non-gemini'] + expect(quota?.remainingPercent).toBe(60) + // Per-window entries must survive the post-apply sidebar write. + expect(quota?.windows).toBeDefined() + expect(quota?.windows?.length).toBeGreaterThanOrEqual(2) + const weekly = quota?.windows?.find((w) => w.window === 'weekly') + expect(weekly).toBeDefined() + expect(weekly?.remainingPercent).toBe(60) + } finally { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + + it('carries healthScore from dialog rows into the sidebar state file', async () => { + // Covers the M1 seam AND the M3 gap: the final projection inside + // createSidebarRefresher must not re-map fields (dropping healthScore), + // and the sidebar state file (not just the returned row) must carry + // the tracker score. A score of 42 should appear as health: 42 in the + // persisted state - not reset to the default 100. + const sidebarFile = join(dir, 'sidebar-health-seam.json') + const previousSidebarFile = process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = sidebarFile + try { + const refresher = createSidebarRefresher(() => []) + const dialogRows: CommandAccountRow[] = [ + { + id: 'acct-0', + index: 0, + label: 'Alpha', + enabled: true, + current: true, + healthScore: 42, + quota: [ + { + key: 'gemini' as const, + label: 'Gemini', + remainingPercent: 80, + windows: [ + { + window: '5h' as const, + remainingPercent: 80, + resetAt: Date.now() + 5 * 60 * 1000, + }, + { + window: 'weekly' as const, + remainingPercent: 75, + resetAt: Date.now() + 7 * 24 * 60 * 60_000, + }, + ], + }, + ], + }, + { + id: 'acct-1', + index: 1, + label: 'Beta', + enabled: false, + current: false, + healthScore: 17, + quota: [], + }, + ] + + await refresher(dialogRows) + + const state = readSidebarState(sidebarFile) + expect(state.accounts).toHaveLength(2) + // Core seam assertion: tracker scores must reach the state file. + expect(state.accounts[0]?.health).toBe(42) + expect(state.accounts[1]?.health).toBe(17) + // Per-window data must also survive the same projection. + const gemini = state.accounts[0]?.quota.gemini + expect(gemini?.windows).toBeDefined() + expect(gemini?.windows?.length).toBeGreaterThanOrEqual(2) + } finally { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile } } @@ -810,9 +939,9 @@ describe('applyCommand', () => { expect(state.accounts[0]?.cooldownUntil).toBe(1_900_000_000_000) expect(state.accounts[1]?.cooldownUntil).toBeUndefined() } finally { - if (previousSidebarFile === undefined) { - delete process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE - } else { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile } } @@ -878,9 +1007,9 @@ describe('applyCommand', () => { expect(state.accounts[0]?.cooldownUntil).toBe(1_900_000_000_000) expect(state.accounts[1]?.cooldownUntil).toBeUndefined() } finally { - if (previousSidebarFile === undefined) { - delete process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE - } else { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile } } @@ -936,9 +1065,9 @@ describe('applyCommand', () => { // account at live index 0. expect(state.accounts[0]?.cooldownUntil).toBeUndefined() } finally { - if (previousSidebarFile === undefined) { - delete process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE - } else { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile } } @@ -996,15 +1125,19 @@ describe('applyCommand', () => { expect(state.accounts[0]?.cooldownUntil).toBe(1_800_000_000_000) expect(state.accounts[1]?.cooldownUntil).toBeUndefined() } finally { - if (previousSidebarFile === undefined) { - delete process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE - } else { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile } } }) - it('starts a quota refresh when the quota dialog opens without delaying its cached payload', async () => { + it('starts a backoff-respecting quota check when the quota dialog opens without delaying its cached payload', async () => { + // N6: the dialog OPEN path now calls refreshQuotaRespectingBackoff + // (non-forced) instead of refreshQuota (forced). The fire-and-forget + // semantics are preserved: if buildDialogPayload accidentally awaits + // the call this test hangs and surfaces as a timeout. const listAccounts = mock(async () => [ { id: 'acct-0', @@ -1015,13 +1148,8 @@ describe('applyCommand', () => { quota: [], }, ]) - // refreshQuota returns a never-resolving promise. If buildDialogPayload - // ever accidentally `await`s the refresh path (instead of the fire- - // and-forget `void` it does today) this test will hang and the - // runner will report it as a timeout — exactly the regression we - // want to surface. let refreshStarted = false - const refreshQuota = mock(() => { + const refreshQuotaRespectingBackoff = mock(() => { refreshStarted = true return new Promise(() => {}) }) @@ -1030,12 +1158,12 @@ describe('applyCommand', () => { client: {} as never, sessionID: 'session-1', settings: ctx.settings, - commandData: { listAccounts, refreshQuota } as never, + commandData: { listAccounts, refreshQuotaRespectingBackoff } as never, }) expect(payload.knobs.accounts).toHaveLength(1) expect(refreshStarted).toBe(true) - expect(refreshQuota).toHaveBeenCalledTimes(1) + expect(refreshQuotaRespectingBackoff).toHaveBeenCalledTimes(1) }) it('returns the expired-pending result from add-oauth-finish', async () => { diff --git a/packages/opencode/src/plugin/commands.ts b/packages/opencode/src/plugin/commands.ts index 5f1cd3c..6dfc584 100644 --- a/packages/opencode/src/plugin/commands.ts +++ b/packages/opencode/src/plugin/commands.ts @@ -25,6 +25,7 @@ import { isTuiConnected as defaultIsTuiConnected } from '../rpc/notifications' import type { CommandModalName } from '../rpc/protocol' import { buildSidebarMachineStateFromAccounts, + type SidebarAccountRedactionInput, setSidebarMachineState, } from '../sidebar-state' import type { AccountCommandOAuthService } from './account-command-oauth' @@ -154,10 +155,12 @@ export async function buildDialogPayload( const accounts = context.commandData ? await context.commandData.listAccounts() : [] - // Render the cached snapshot immediately while the shared manager - // refreshes all accounts; the mounted panel polls the fenced state file. + // Render the cached snapshot immediately. Opening the dialog is a + // VIEW — it must respect the quota manager's per-account backoff so + // a rate-limited account isn't hammered every time the dialog opens. + // Explicit user-triggered refreshes (inside the dialog) keep force:true. if (context.commandData) { - void context.commandData.refreshQuota().catch(() => {}) + void context.commandData.refreshQuotaRespectingBackoff().catch(() => {}) } return { command, @@ -751,6 +754,8 @@ export function createSidebarRefresher( gemini?: { remainingFraction?: number; resetTime?: string } 'non-gemini'?: { remainingFraction?: number; resetTime?: string } } + /** Captured plan tier. Absent when unknown. */ + tier?: { id: string; capturedAt: number } }> | null, ): (accounts?: CommandAccountRow[]) => Promise { return async (dialogAccounts) => { @@ -785,7 +790,14 @@ export function createSidebarRefresher( index: entry.index, label: entry.label, enabled: entry.enabled, + current: entry.current, coolingDownUntil: liveCooldown, + healthScore: entry.healthScore, + tier: entry.tier, + // Rebuild cachedQuota carrying the windows array so the sidebar + // doesn't collapse from per-window rows to a single aggregate bar + // after every apply. Mirrors the QuotaGroupSummary shape expected + // by projectQuotaPoolForSidebar inside buildSidebarMachineStateFromAccounts. cachedQuota: Object.fromEntries( entry.quota.flatMap((group) => { if (group.remainingPercent == null) return [] @@ -797,6 +809,19 @@ export function createSidebarRefresher( ...(group.resetAt === undefined ? {} : { resetTime: new Date(group.resetAt).toISOString() }), + ...(group.windows && group.windows.length > 0 + ? { + windows: group.windows.map((w) => ({ + window: w.window, + remainingFraction: w.remainingPercent / 100, + resetTime: + typeof w.resetAt === 'number' && + Number.isFinite(w.resetAt) + ? new Date(w.resetAt).toISOString() + : '', + })), + } + : {}), }, ], ] @@ -807,15 +832,14 @@ export function createSidebarRefresher( : getAccounts() if (!accounts || accounts.length === 0) return try { + // Pass the already-projected accounts directly so every field the + // intermediate map built (healthScore, cachedQuota with windows, ...) + // reaches the redactor unchanged. A second .map here would silently + // drop any field not explicitly forwarded -- that is the class of + // defect this cast prevents. await setSidebarMachineState( buildSidebarMachineStateFromAccounts( - accounts.map((entry) => ({ - index: entry.index, - label: entry.label, - enabled: entry.enabled, - coolingDownUntil: entry.coolingDownUntil, - cachedQuota: entry.cachedQuota, - })), + accounts as SidebarAccountRedactionInput[], ), ) } catch { diff --git a/packages/opencode/src/plugin/config/schema.ts b/packages/opencode/src/plugin/config/schema.ts index 7f4ecad..d5aa8ca 100644 --- a/packages/opencode/src/plugin/config/schema.ts +++ b/packages/opencode/src/plugin/config/schema.ts @@ -513,6 +513,31 @@ export const AntigravityConfigSchema = z.object({ // Auto-Update // ========================================================================= + // ========================================================================= + // Background Quota Refresh + // ========================================================================= + + /** + * Enable a background timer that periodically refreshes quota for all + * accounts. Without this, idle sessions show stale sidebar bars + * because Antigravity is a poll-only API (no quota headers in responses). + * + * @default true + */ + background_quota_refresh: z.boolean().default(true), + + /** + * How often the background poller refreshes quota (in minutes). + * Valid range: 1–60. + * + * @default 5 + */ + background_quota_refresh_interval_minutes: z + .number() + .min(1) + .max(60) + .default(5), + /** * Enable automatic plugin updates. * @default true @@ -603,6 +628,8 @@ export const DEFAULT_CONFIG: AntigravityConfig = { quota_refresh_interval_minutes: 30, soft_quota_cache_ttl_minutes: 'auto', proactive_rotation_threshold_percent: 20, + background_quota_refresh: true, + background_quota_refresh_interval_minutes: 5, auto_update: true, signature_cache: { enabled: true, diff --git a/packages/opencode/src/plugin/debug.test.ts b/packages/opencode/src/plugin/debug.test.ts index 8249385..0288384 100644 --- a/packages/opencode/src/plugin/debug.test.ts +++ b/packages/opencode/src/plugin/debug.test.ts @@ -126,9 +126,8 @@ describe('debug sink redaction', () => { }) it('masks full project IDs in debug logs', async () => { - const { initializeDebug, startAntigravityDebugRequest } = await import( - './debug' - ) + const { initializeDebug, startAntigravityDebugRequest, closeDebugLog } = + await import('./debug') initializeDebug({ ...DEFAULT_CONFIG, @@ -154,8 +153,9 @@ describe('debug sink redaction', () => { projectId: knownProjectId, }) - // Allow the writeStream a tick to flush the line. - await new Promise((r) => setTimeout(r, 50)) + // Close the write stream so all buffered data is flushed to disk + // before reading the log file. + await closeDebugLog() // The debug log file is the most recent `antigravity-debug-*.log` in // the log dir. @@ -178,9 +178,8 @@ describe('debug sink redaction', () => { }) it('masks fingerprint User-Agent headers in the recorded headers dump', async () => { - const { initializeDebug, startAntigravityDebugRequest } = await import( - './debug' - ) + const { initializeDebug, startAntigravityDebugRequest, closeDebugLog } = + await import('./debug') initializeDebug({ ...DEFAULT_CONFIG, @@ -204,7 +203,7 @@ describe('debug sink redaction', () => { streaming: false, }) - await new Promise((r) => setTimeout(r, 50)) + await closeDebugLog() const files = readdirSync(logDir) const logFile = files diff --git a/packages/opencode/src/plugin/debug.ts b/packages/opencode/src/plugin/debug.ts index a1aaf1e..7fe9f7d 100644 --- a/packages/opencode/src/plugin/debug.ts +++ b/packages/opencode/src/plugin/debug.ts @@ -60,6 +60,30 @@ function closeDebugStream(): void { } } +/** + * Close the debug log stream and resolve once all buffered writes + * have been flushed to disk. Tests use this to wait deterministically + * for log lines before reading the log file; the internal reinitialize + * path calls {@link closeDebugStream} directly (fire-and-forget). + */ +export function closeDebugLog(): Promise { + const current = debugState?.logStream + if (!current) return Promise.resolve() + // Detach the reference so reinitialization won't double-close. + if (debugState) { + debugState.logStream = null + } + return new Promise((resolve) => { + current.once('close', () => resolve()) + current.once('error', () => resolve()) + try { + current.end() + } catch { + resolve() + } + }) +} + /** * Get the OS-specific config directory. */ diff --git a/packages/opencode/src/plugin/index.ts b/packages/opencode/src/plugin/index.ts index d1f7aeb..179316a 100644 --- a/packages/opencode/src/plugin/index.ts +++ b/packages/opencode/src/plugin/index.ts @@ -5,7 +5,12 @@ import { createAutoUpdateCheckerHook } from '../hooks/auto-update-checker' import { drainNotifications, pushNotification } from '../rpc/notifications' import { getRpcDir } from '../rpc/rpc-dir' import { startRpcServer } from '../rpc/rpc-server' -import { drainSidebarWrites, getSidebarStateFile } from '../sidebar-state' +import { + drainSidebarWrites, + getSidebarStateFile, + isAccountCurrent, + toCapturedTier, +} from '../sidebar-state' import { createAccountAccessService, promptAccountIndexForVerification, @@ -13,6 +18,7 @@ import { } from './account-access' import { createAccountCommandOAuthService } from './account-command-oauth' import { createAuthLoader } from './auth-loader' +import { BackgroundQuotaRefresh } from './background-quota-refresh' import { initDiskSignatureCache, shutdownDiskSignatureCache } from './cache' import { applyAntigravityProviderCatalog, @@ -29,8 +35,8 @@ import { createSidebarRefresher, } from './commands' import { initRuntimeConfig, loadConfig } from './config' -import { getUserConfigPath as getUserConfigDir } from './config/loader' -import { initializeDebug } from './debug' +import { getUserConfigPath } from './config/loader' +import { closeDebugLog, initializeDebug } from './debug' import { type PluginDependencies, type PluginDependencyOverrides, @@ -47,9 +53,17 @@ import { type OperatorSettingsController, } from './operator-settings' import { persistAccountPool } from './persist-account-pool' -import { createOpenCodeQuotaManager, type QuotaManager } from './quota' +import { + createOpenCodeQuotaManager, + makeTierLoader, + type QuotaManager, +} from './quota' import { createSessionRecoveryHook } from './recovery' -import { initHealthTracker, initTokenTracker } from './rotation' +import { + getHealthTracker, + initHealthTracker, + initTokenTracker, +} from './rotation' import { AgySessionRegistry } from './session-context' import { clearAccounts, @@ -90,6 +104,12 @@ export interface CreateAntigravityPluginOptions { * randomness. Defaults to production implementations. */ dependencies?: PluginDependencyOverrides + /** + * Test-only seam: called synchronously after the background poller is + * constructed, before `start()`. Lets wiring tests capture the instance + * without reaching into lifecycle internals. + */ + _onPollerCreated?: (poller: BackgroundQuotaRefresh) => void } export function registerQuotaManagerProducer( @@ -151,6 +171,12 @@ export const createAntigravityPlugin = drainSidebarWrites, }) const quotaManager = createOpenCodeQuotaManager(client, providerId, { + // Route quota fetches through the injected transport so e2e and + // custom-host deployments stay on the loopback/mock server for both + // the request path and the background poller. Without this the + // poller always hits the production Antigravity endpoint even when + // the caller injected a mock transport. + fetchVia: dependencies.agyTransport, // Bind to the live AccountManager so every refresh (manual or // background) pushes the freshly-updated quota percentages into the // sidebar without an extra RPC. The wrapper reads lazily so the @@ -158,10 +184,12 @@ export const createAntigravityPlugin = getAccountsForSidebar: () => { const manager = lifecycle.getAccountManager() if (!manager) return null + const activeByFamily = manager.getActiveIndexByFamily() return manager.getAccounts().map((entry) => ({ index: entry.index, label: entry.label, enabled: entry.enabled, + current: isAccountCurrent(entry.index, activeByFamily), coolingDownUntil: entry.coolingDownUntil, cachedQuota: entry.cachedQuota, // Carry the identity stamp so the sidebar projection can @@ -169,8 +197,14 @@ export const createAntigravityPlugin = // after an index shift (see `redactAccountForSidebar`). cachedQuotaAccountId: entry.cachedQuotaAccountId, currentQuotaAccountId: quotaAccountIdentity(entry.parts.refreshToken), + tier: toCapturedTier(entry), })) }, + getActiveIndexByFamily: () => { + const manager = lifecycle.getAccountManager() + if (!manager) return null + return manager.getActiveIndexByFamily() + }, }) // Producer phase: the quota manager emits fire-and-forget sidebar // writes after every refresh. Its dispose() awaits any in-flight @@ -180,6 +214,30 @@ export const createAntigravityPlugin = // after drainSidebarWrites() already asserted the queue was empty. registerQuotaManagerProducer(lifecycle, quotaManager) + // Background quota poller: per-loader-instance timer that keeps idle + // account sidebar bars fresh between real requests. Registered as a + // producer so it is stopped and awaited BEFORE the sidebar drain, + // guaranteeing any final poll write lands before the drain flushes. + if (config.background_quota_refresh) { + const poller = new BackgroundQuotaRefresh({ + intervalMs: config.background_quota_refresh_interval_minutes * 60_000, + sidebarStateFile: getSidebarStateFile(), + getAccountManager: () => lifecycle.getAccountManager(), + quotaManager, + // Resolve plan tier for accounts with stale capturedTierAt. Uses + // the same token-refresh path as the quota manager without a separate + // network seam -- tier is best-effort metadata, not quota. + loadAccountTier: makeTierLoader(client, providerId), + // Thread the injected clock so e2e tests can control startup + // jitter and timing without needing a real 30-second wait. + now: dependencies.clock.now, + random: dependencies.clock.random, + }) + options._onPollerCreated?.(poller) + poller.start() + lifecycle.register({ dispose: () => poller.dispose() }, 'producer') + } + // Operator settings controller backs the /antigravity-* slash commands. // The controller loads existing persisted settings at first read, mutates // runtime config immediately, and serializes through the fenced-lock @@ -187,7 +245,10 @@ export const createAntigravityPlugin = const operatorSettings: OperatorSettingsController = createOperatorSettingsController({ projectConfigPath: join(directory, '.opencode', 'antigravity.json'), - userConfigPath: join(getUserConfigDir(), 'antigravity.json'), + // getUserConfigPath() already returns the full file path including + // 'antigravity.json' — do NOT join again or the path double-nests + // into /antigravity.json/antigravity.json. + userConfigPath: getUserConfigPath(), }) setRuntimeLogLevel(operatorSettings.get().log_level) lifecycle.register({ dispose: () => operatorSettings.dispose() }) @@ -217,19 +278,22 @@ export const createAntigravityPlugin = getAccounts: () => { const manager = lifecycle.getAccountManager() if (!manager) return [] - const current = - manager.getCurrentAccountForFamily('claude')?.index ?? 0 + const activeByFamily = manager.getActiveIndexByFamily() return manager.getAccounts().map((entry, index) => ({ index, refreshToken: entry.parts.refreshToken, label: entry.label, enabled: entry.enabled !== false, - active: index === current, + active: isAccountCurrent(index, activeByFamily), cachedQuota: entry.cachedQuota, cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, cachedQuotaAccountId: entry.cachedQuotaAccountId, accountIneligible: entry.accountIneligible, coolingDownUntil: entry.coolingDownUntil, + healthScore: getHealthTracker().getScore(entry.index), + capturedTierId: entry.capturedTierId, + capturedPaidTierId: entry.capturedPaidTierId, + capturedTierAt: entry.capturedTierAt, })) }, getAccountsForQuotaCheck: () => { @@ -376,10 +440,12 @@ export const createAntigravityPlugin = const refreshSidebar = createSidebarRefresher(() => { const manager = lifecycle.getAccountManager() if (!manager) return null + const activeByFamily = manager.getActiveIndexByFamily() return manager.getAccounts().map((entry) => ({ index: entry.index, label: entry.label, enabled: entry.enabled, + current: isAccountCurrent(entry.index, activeByFamily), coolingDownUntil: entry.coolingDownUntil, cachedQuota: entry.cachedQuota, // Carry the identity stamp so the sidebar projection can @@ -389,6 +455,7 @@ export const createAntigravityPlugin = currentQuotaAccountId: quotaAccountIdentity( entry.parts.refreshToken, ), + tier: toCapturedTier(entry), })) }) const result = await applyCommand(request, { @@ -407,6 +474,9 @@ export const createAntigravityPlugin = drain: drainNotifications, }) lifecycle.register({ dispose: () => rpcServer.stop() }) + // Flush the debug log stream after the sidebar drain so the last + // buffered lines land on disk before the process exits. Best-effort. + lifecycle.register({ dispose: () => closeDebugLog().catch(() => {}) }) return { dispose: async () => { diff --git a/packages/opencode/src/plugin/lifecycle.test.ts b/packages/opencode/src/plugin/lifecycle.test.ts index 5d73d25..93bd41a 100644 --- a/packages/opencode/src/plugin/lifecycle.test.ts +++ b/packages/opencode/src/plugin/lifecycle.test.ts @@ -1,6 +1,11 @@ -import { describe, expect, it, mock } from 'bun:test' +import { afterEach, beforeEach, describe, expect, it, mock } from 'bun:test' +import { mkdirSync, readdirSync, readFileSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' import type { AccountManager } from './accounts' +import { DEFAULT_CONFIG } from './config' +import { initializeDebug } from './debug' import { createPluginLifecycle } from './lifecycle' import type { ProactiveRefreshQueue } from './refresh-queue' @@ -344,3 +349,107 @@ describe('PluginLifecycle RPC ownership', () => { expect(stop).toHaveBeenCalledTimes(1) }) }) + +describe('PluginLifecycle debug log disposal', () => { + let logDir: string + + beforeEach(() => { + logDir = join( + tmpdir(), + `agy-lifecycle-debug-${Math.random().toString(36).slice(2)}`, + ) + mkdirSync(logDir, { recursive: true }) + }) + + afterEach(() => { + // Reset module-global debug state before removing the temp dir so any + // stream targeting logDir is closed and the global points nowhere live. + // Without this, a stale WriteStream in debugState outlives the directory; + // under --isolate this is benign today (each file gets a fresh registry) + // but correct teardown guards against future tests added to this suite. + initializeDebug({ ...DEFAULT_CONFIG, debug: false }) + rmSync(logDir, { recursive: true, force: true }) + }) + + it('calls closeDebugLog as a consumer-phase disposable during lifecycle disposal', async () => { + const events: string[] = [] + const lifecycle = createPluginLifecycle({ + sessionRegistry: { clear: () => {} }, + shutdownDiskSignatureCache: async () => {}, + clearFetchState: () => {}, + drainSidebarWrites: async () => { + events.push('sidebar:drain') + }, + }) + // Simulate the plugin registration: closeDebugLog is registered as a + // consumer so it runs AFTER the sidebar drain. + lifecycle.register( + { + dispose: async () => { + events.push('debug-log:close') + }, + }, + 'consumer', + ) + + await lifecycle.dispose() + + // The debug log close must happen after the sidebar drain — consumers + // are disposed in phase 3, after the drain in phase 2. + expect(events.indexOf('sidebar:drain')).toBeLessThan( + events.indexOf('debug-log:close'), + ) + }) + + it('flushes buffered debug lines to disk when closeDebugLog runs during dispose', async () => { + // Initialize debug with a known log dir so closeDebugLog targets a + // predictable file path. + const { + initializeDebug, + getLogFilePath, + closeDebugLog, + startAntigravityDebugRequest, + } = await import('./debug') + + initializeDebug({ + ...DEFAULT_CONFIG, + debug: true, + debug_tui: false, + log_dir: logDir, + }) + + // Write a log line that would be buffered in the WriteStream. + startAntigravityDebugRequest({ + originalUrl: 'https://example.com/v1', + resolvedUrl: 'https://example.com/v1', + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: '"dispose-test-body"', + streaming: false, + }) + + const lifecycle = createPluginLifecycle({ + sessionRegistry: { clear: () => {} }, + shutdownDiskSignatureCache: async () => {}, + clearFetchState: () => {}, + }) + lifecycle.register({ + dispose: () => closeDebugLog().catch(() => {}), + }) + + await lifecycle.dispose() + + // After dispose, the debug log file must exist and contain the + // test marker written before shutdown. + const logPath = getLogFilePath() + expect(logPath).toBeTruthy() + const files = readdirSync(logDir) + const logFile = files + .filter((f) => f.startsWith('antigravity-debug-') && f.endsWith('.log')) + .sort() + .pop() + expect(logFile).toBeTruthy() + const contents = readFileSync(join(logDir, logFile!), 'utf8') + expect(contents).toContain('dispose-test-body') + }) +}) diff --git a/packages/opencode/src/plugin/quota.test.ts b/packages/opencode/src/plugin/quota.test.ts index 5225605..5aea92e 100644 --- a/packages/opencode/src/plugin/quota.test.ts +++ b/packages/opencode/src/plugin/quota.test.ts @@ -78,15 +78,21 @@ describe('classifyQuotaGroup', () => { describe('pushSidebarQuotaSnapshot', () => { let dir: string let stateFile: string + let savedSidebarEnv: string | undefined beforeEach(() => { + // Save preload-pinned value so afterEach can restore it instead of + // deleting — a delete drops resolution to the operator's real state dir. + savedSidebarEnv = process.env[SIDEBAR_STATE_ENV] dir = mkdtempSync(join(tmpdir(), 'agy-quota-sidebar-')) stateFile = join(dir, 'sidebar-state.json') process.env[SIDEBAR_STATE_ENV] = stateFile }) afterEach(() => { - delete process.env[SIDEBAR_STATE_ENV] + if (savedSidebarEnv !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebarEnv + else delete process.env[SIDEBAR_STATE_ENV] rmSync(dir, { recursive: true, force: true }) }) @@ -287,6 +293,75 @@ describe('pushSidebarQuotaSnapshot', () => { } }) + it('marks the active claude account as current when getActiveIndexByFamily is passed', async () => { + const getAccounts = (): QuotaSnapshotAccount[] => [ + { + index: 0, + label: 'Active', + enabled: true, + cachedQuota: { + 'non-gemini': { remainingFraction: 0.8, modelCount: 1 }, + }, + }, + { + index: 1, + label: 'Idle', + enabled: true, + }, + ] + + await pushSidebarQuotaSnapshot(getAccounts, 0, () => ({ + claude: 0, + gemini: 0, + })) + + const state = read() + expect(state.accounts).toHaveLength(2) + expect(state.accounts[0]?.current).toBe(true) + expect(state.accounts[1]?.current).toBe(false) + }) + + it('marks both accounts current when each family points to a different index', async () => { + const getAccounts = (): QuotaSnapshotAccount[] => [ + { index: 0, label: 'Claude', enabled: true }, + { index: 1, label: 'Middle', enabled: true }, + { index: 2, label: 'Gemini', enabled: true }, + ] + + await pushSidebarQuotaSnapshot(getAccounts, 0, () => ({ + claude: 0, + gemini: 2, + })) + + const state = read() + expect(state.accounts).toHaveLength(3) + expect(state.accounts[0]?.current).toBe(true) + expect(state.accounts[1]?.current).toBe(false) + expect(state.accounts[2]?.current).toBe(true) + }) + + it('defaults current to false when getActiveIndexByFamily is omitted (backward compat)', async () => { + const getAccounts = (): QuotaSnapshotAccount[] => [ + { index: 0, label: 'Acc', enabled: true }, + ] + + await pushSidebarQuotaSnapshot(getAccounts, 0) + + const state = read() + expect(state.accounts[0]?.current).toBe(false) + }) + + it('returns null from getActiveIndexByFamily → all accounts false', async () => { + const getAccounts = (): QuotaSnapshotAccount[] => [ + { index: 0, label: 'A', enabled: true }, + ] + + await pushSidebarQuotaSnapshot(getAccounts, 0, () => null) + + const state = read() + expect(state.accounts[0]?.current).toBe(false) + }) + it('fences the real quota wrapper sidebar enqueue before the lifecycle drain', async () => { const events: string[] = [] let releaseFetch!: () => void @@ -371,6 +446,10 @@ describe('pushSidebarQuotaSnapshot', () => { 'fetch:start', 'fetch:start', 'fetch:start', + // fetchGeminiCliQuota now also uses fetchVia (2 endpoints via the + // FetchGeminiCliQuotaOptions.fetchVia seam added for N2 testability). + 'fetch:start', + 'fetch:start', 'sidebar:write-start', 'lifecycle:drain', 'drain:sees-sidebar-write', @@ -382,4 +461,169 @@ describe('pushSidebarQuotaSnapshot', () => { fetchSpy.mockRestore() } }) + + it('N2: CLI rejection carries the real error message instead of the generic no-CLI-configured string', async () => { + // fetchGeminiCliQuota now propagates transport-level errors (network abort, + // DNS, socket hang, timeout) by collecting them per-endpoint and throwing + // when all endpoints fail. That throw reaches the outer .catch() in + // quota.ts, which sets geminiCliFetchError. Before N2 the catch was + // present but never fired — fetchGeminiCliQuota silently swallowed errors + // and returned { buckets: [] }, landing on the generic message. + // + // The new FetchGeminiCliQuotaOptions.fetchVia seam lets the test inject + // a transport that throws, driving the outer .catch() directly. + const CLI_THROW_MSG = 'socket hang up' + + const fetchSpy = spyOn(globalThis, 'fetch').mockImplementation((async ( + input: unknown, + ) => { + const url = + typeof input === 'string' + ? input + : input instanceof URL + ? input.toString() + : (input as Request).url + + if (url.includes('retrieveUserQuotaSummary')) { + return new Response( + JSON.stringify({ + groups: [ + { + displayName: 'Gemini Models', + buckets: [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly', + window: 'weekly', + resetTime: '2026-01-08T00:00:00Z', + remainingFraction: 0.6, + }, + ], + }, + ], + }), + { status: 200 }, + ) + } + if (url.includes('retrieveUserQuota')) { + // Simulate a transport-level rejection (socket hang up). + // fetchGeminiCliQuota collects this as an error and re-throws + // at end-of-loop, triggering the outer .catch() in quota.ts. + throw new Error(CLI_THROW_MSG) + } + return new Response( + JSON.stringify({ access_token: 'access-token', expires_in: 3600 }), + { status: 200 }, + ) + }) as unknown as typeof fetch) + + const client = { + auth: { set: mock(async () => {}) }, + } as unknown as PluginClient + const account: AccountMetadataV3 = { + refreshToken: `n2-rejection-${Date.now()}-${Math.random()}`, + managedProjectId: 'managed-n2', + projectId: 'project-n2', + addedAt: 0, + lastUsed: 0, + } + const manager = createOpenCodeQuotaManager(client, 'google') + let result: import('./quota.ts').AccountQuotaResult | undefined + try { + const results = await manager.refreshAccounts([account], { + indexFor: () => 0, + force: true, + }) + result = results[0] + } finally { + fetchSpy.mockRestore() + } + + // Summary succeeded — overall status must stay 'ok'. + expect(result?.status).toBe('ok') + if (result?.status !== 'ok') return + + // Summary groups must survive the CLI rejection. + expect(result.quota?.groups).toBeDefined() + + // The annotation must carry the REAL thrown message, not the generic + // 'No Gemini CLI quota available' that indicates a permanent absence. + // Multiple endpoints may each contribute the message (joined by '; '). + expect(result.geminiCliQuota?.error).toContain(CLI_THROW_MSG) + expect(result.geminiCliQuota?.error).not.toBe( + 'No Gemini CLI quota available', + ) + }) + + it('N2: HTTP-500 on CLI endpoint does NOT kill the summary (parallel-fetch isolation)', async () => { + // Complement: an HTTP 500 is treated as a transport error by the updated + // fetchGeminiCliQuota (errors[] + re-throw), so the same .catch() path + // fires and the summary still flows through. + const fetchSpy = spyOn(globalThis, 'fetch').mockImplementation((async ( + input: unknown, + ) => { + const url = + typeof input === 'string' + ? input + : input instanceof URL + ? input.toString() + : (input as Request).url + + if (url.includes('retrieveUserQuotaSummary')) { + return new Response( + JSON.stringify({ + groups: [ + { + displayName: 'Gemini Models', + buckets: [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly', + window: 'weekly', + resetTime: '2026-01-08T00:00:00Z', + remainingFraction: 0.6, + }, + ], + }, + ], + }), + { status: 200 }, + ) + } + if (url.includes('retrieveUserQuota')) { + return new Response('internal error', { status: 500 }) + } + return new Response( + JSON.stringify({ access_token: 'access-token', expires_in: 3600 }), + { status: 200 }, + ) + }) as unknown as typeof fetch) + + const client = { + auth: { set: mock(async () => {}) }, + } as unknown as PluginClient + const account: AccountMetadataV3 = { + refreshToken: `n2-500-${Date.now()}-${Math.random()}`, + managedProjectId: 'managed-n2b', + projectId: 'project-n2b', + addedAt: 0, + lastUsed: 0, + } + const manager = createOpenCodeQuotaManager(client, 'google') + let result: import('./quota.ts').AccountQuotaResult | undefined + try { + const results = await manager.refreshAccounts([account], { + indexFor: () => 0, + force: true, + }) + result = results[0] + } finally { + fetchSpy.mockRestore() + } + + expect(result?.status).toBe('ok') + if (result?.status !== 'ok') return + expect(result.quota?.groups).toBeDefined() + expect(result.geminiCliQuota?.error).toBeTruthy() + }) }) diff --git a/packages/opencode/src/plugin/quota.ts b/packages/opencode/src/plugin/quota.ts index bc913fb..2f9a1d3 100644 --- a/packages/opencode/src/plugin/quota.ts +++ b/packages/opencode/src/plugin/quota.ts @@ -29,6 +29,7 @@ import { fetchGeminiCliQuota, fetchQuotaSummary, type GeminiCliQuotaSummary, + getHealthTracker, type QuotaManager, type QuotaSummary, } from '@cortexkit/antigravity-auth-core' @@ -40,6 +41,7 @@ import { } from '../constants' import { buildSidebarMachineStateFromAccounts, + isAccountCurrent, setSidebarMachineState, } from '../sidebar-state' import { @@ -50,7 +52,7 @@ import { import { logQuotaFetch, logQuotaStatus } from './debug' import { buildAntigravityHarnessUserAgent } from './fingerprint' import { createLogger } from './logger' -import { ensureProjectContext } from './project' +import { ensureProjectContext, loadManagedProject } from './project' import { refreshAccessToken } from './token' import type { OAuthAuthDetails, PluginClient } from './types' @@ -118,6 +120,15 @@ export function createOpenCodeQuotaManager( cachedQuotaAccountId?: string currentQuotaAccountId?: string }> | null + /** + * Optional provider for the active-account indexes per model family. + * Wired by the plugin entry so every quota-refresh sidebar snapshot + * carries the real `current` flag — not a hardcoded `false`. + */ + getActiveIndexByFamily?: () => { + claude: number + gemini: number + } | null /** * Optional transport adapter used for both `fetchAvailableModels` * and the project-context lookup. When omitted, the production @@ -143,6 +154,7 @@ export function createOpenCodeQuotaManager( const originalRefreshAccount = manager.refreshAccount const originalRefreshAccounts = manager.refreshAccounts const getAccountsForSidebar = options.getAccountsForSidebar + const getActiveIndexByFamily = options.getActiveIndexByFamily let disposed = false const inFlight = new Set>() @@ -153,6 +165,7 @@ export function createOpenCodeQuotaManager( await pushSidebarQuotaSnapshot( getAccountsForSidebar, manager.getBackoffUntil(account), + getActiveIndexByFamily, ).catch(() => { // Sidebar persistence remains best-effort when lock contention // outlives its retry budget. @@ -290,11 +303,18 @@ export async function pushSidebarQuotaSnapshot( cachedQuota?: AccountMetadataV3['cachedQuota'] cachedQuotaAccountId?: string currentQuotaAccountId?: string + /** Captured plan tier to surface in the sidebar state file. */ + tier?: { id: string; paidId?: string; capturedAt: number } }> | null, backoffUntil: number = 0, + getActiveIndexByFamily?: () => { + claude: number + gemini: number + } | null, ): Promise { const accounts = getAccounts() if (!accounts || accounts.length === 0) return + const activeByFamily = getActiveIndexByFamily?.() ?? null try { await setSidebarMachineState( buildSidebarMachineStateFromAccounts( @@ -302,11 +322,15 @@ export async function pushSidebarQuotaSnapshot( index: entry.index, label: entry.label, enabled: entry.enabled, - current: false, + current: activeByFamily + ? isAccountCurrent(entry.index, activeByFamily) + : false, coolingDownUntil: entry.coolingDownUntil, cachedQuota: entry.cachedQuota, cachedQuotaAccountId: entry.cachedQuotaAccountId, currentQuotaAccountId: entry.currentQuotaAccountId, + healthScore: getHealthTracker().getScore(entry.index), + tier: entry.tier, })), { checkedAt: Date.now(), @@ -404,7 +428,11 @@ function makeFetchAccountQuota( const projectContext = await ensureProjectContext(auth) auth = projectContext.auth - const updatedAccount = applyAccountUpdates(account, auth) + const updatedAccount = applyAccountUpdates( + account, + auth, + projectContext.capturedTier, + ) if (rotatedRefresh && client) { await persistRotatedRefresh(client, providerId, auth).catch(() => {}) @@ -456,6 +484,12 @@ function makeFetchAccountQuota( } })() + // CLI fetch is independent of the summary fetch. A CLI failure must + // NOT kill the summary result, but it also must not be laundered into + // "No Gemini CLI quota available" (a permanent-looking status) when + // the real cause is a transient network error. Capture the error + // message separately so the annotated result can carry it. + let geminiCliFetchError: string | undefined const fetchGeminiCliPayload = fetchGeminiCliQuota({ accessToken: auth.access ?? '', projectId: projectContext.effectiveProjectId, @@ -464,13 +498,9 @@ function makeFetchAccountQuota( timeoutMs: 10_000, ...(fetchVia ? { fetchVia } : {}), }).catch((error: unknown) => { - // The previous sequence (`await fetchGeminiCliQuota`) let - // any throw bubble to the outer try/catch and produced an - // undefined buckets downstream. Preserve that semantics so - // the parallel path is a drop-in replacement. - log.debug('fetchGeminiCliQuota failed', { - error: error instanceof Error ? error.message : String(error), - }) + geminiCliFetchError = + error instanceof Error ? error.message : String(error) + log.debug('fetchGeminiCliQuota failed', { error: geminiCliFetchError }) return { buckets: undefined } as Awaited< ReturnType > @@ -490,9 +520,12 @@ function makeFetchAccountQuota( ? { ...geminiCliQuotaResult, error: - geminiCliQuotaResult.models.length === 0 + // A real fetch exception is a transient failure, not a + // "no CLI configured" scenario — propagate the actual message. + geminiCliFetchError ?? + (geminiCliQuotaResult.models.length === 0 ? 'No Gemini CLI quota available' - : undefined, + : undefined), } : geminiCliQuotaResult @@ -546,6 +579,7 @@ function buildAuthFromAccount(account: AccountMetadataV3): OAuthAuthDetails { function applyAccountUpdates( account: AccountMetadataV3, auth: OAuthAuthDetails, + capturedTier?: { id: string; paidId?: string; capturedAt: number }, ): AccountMetadataV3 | undefined { const parts = parseRefreshParts(auth.refresh) if (!parts.refreshToken) { @@ -557,12 +591,30 @@ function applyAccountUpdates( refreshToken: parts.refreshToken, projectId: parts.projectId ?? account.projectId, managedProjectId: parts.managedProjectId ?? account.managedProjectId, + // Persist the captured tier alongside the project-context write. Only + // present when the loadCodeAssist payload returned a non-empty currentTier.id. + ...(capturedTier + ? { + capturedTierId: capturedTier.id, + ...(capturedTier.paidId !== undefined + ? { capturedPaidTierId: capturedTier.paidId } + : {}), + capturedTierAt: capturedTier.capturedAt, + } + : {}), } const changed = updated.refreshToken !== account.refreshToken || updated.projectId !== account.projectId || - updated.managedProjectId !== account.managedProjectId + updated.managedProjectId !== account.managedProjectId || + updated.capturedTierId !== account.capturedTierId || + updated.capturedPaidTierId !== account.capturedPaidTierId || + // capturedAt represents when the tier was LAST CONFIRMED, not when it + // changed -- always update it on a successful observation so consumers + // can gate staleness on that timestamp even when the id stays the same. + (capturedTier !== undefined && + updated.capturedTierAt !== account.capturedTierAt) return changed ? updated : undefined } @@ -582,3 +634,57 @@ async function persistRotatedRefresh( }, }) } + +/** + * Build a per-account tier-loader callback for the background poller. + * + * Calls `loadManagedProject` (loadCodeAssist) directly, bypassing the + * `ensureProjectContext` cache that fast-paths on `managedProjectId` and + * never returns a tier for existing accounts. One call per account per 24 h. + * + * Uses the same token-refresh infrastructure as `makeFetchAccountQuota` so + * an expired access token does not silently fail the tier lookup. + * + * `loadManagedProject` uses the production TLS transport (`fetchWithAgyCliTransport`) + * and is not interceptable via `fetchVia` -- the same design constraint applies to + * `ensureProjectContext`. Tier lookup is best-effort; any failure resolves `null`. + */ +export function makeTierLoader( + client: PluginClient | undefined, + providerId: string, +): ( + account: AccountMetadataV3, +) => Promise<{ id: string; paidId?: string; capturedAt: number } | null> { + return async (account) => { + try { + let auth = buildAuthFromAccount(account) + if (accessTokenExpired(auth)) { + const refreshed = await refreshAccessToken( + auth, + client as PluginClient, + providerId, + ) + if (!refreshed) return null + auth = refreshed + } + + const accessToken = auth.access + if (!accessToken) return null + + const payload = await loadManagedProject(accessToken) + if (!payload?.currentTier?.id) return null + const paidTierId = + typeof payload.paidTier === 'string' + ? payload.paidTier + : payload.paidTier?.id + + return { + id: payload.currentTier.id, + ...(paidTierId ? { paidId: paidTierId } : {}), + capturedAt: Date.now(), + } + } catch { + return null + } + } +} diff --git a/packages/opencode/src/sidebar-state.test.ts b/packages/opencode/src/sidebar-state.test.ts index cbf24d1..41141b2 100644 --- a/packages/opencode/src/sidebar-state.test.ts +++ b/packages/opencode/src/sidebar-state.test.ts @@ -26,6 +26,8 @@ import { acquireFencedFileLock } from '@cortexkit/antigravity-auth-core' import { buildSidebarMachineStateFromAccounts, DEFAULT_SIDEBAR_STATE, + isAccountCurrent, + normalizeLegacyCachedQuota, pruneActiveRouting, readSidebarState, redactAccountForSidebar, @@ -106,6 +108,27 @@ async function waitFor( } } +describe('isAccountCurrent', () => { + it('returns true when index matches the claude active index', () => { + expect(isAccountCurrent(0, { claude: 0, gemini: 0 })).toBe(true) + }) + + it('returns true when index matches the gemini active index but not claude', () => { + expect(isAccountCurrent(2, { claude: 0, gemini: 2 })).toBe(true) + }) + + it('returns false when index matches neither family', () => { + expect(isAccountCurrent(1, { claude: 0, gemini: 0 })).toBe(false) + }) + + it('returns true for both indices when families point to different accounts', () => { + const active = { claude: 0, gemini: 2 } + expect(isAccountCurrent(0, active)).toBe(true) + expect(isAccountCurrent(2, active)).toBe(true) + expect(isAccountCurrent(1, active)).toBe(false) + }) +}) + describe('setSidebarMachineState — freshness merge', () => { let fixture: Fixture @@ -493,6 +516,19 @@ describe('redaction', () => { } }) + it('defaults healthScore to 100 when missing', () => { + // The doc comment says "Defaults to 100 when missing" — the + // implementation must NOT render 0 just because no producer + // passed the field. + const redacted = redactAccountForSidebar({ index: 0 }) + expect(redacted.health).toBe(100) + }) + + it('passes through an explicit healthScore', () => { + const redacted = redactAccountForSidebar({ index: 0, healthScore: 42 }) + expect(redacted.health).toBe(42) + }) + it('replaces profile labels with privacy-safe ordinal labels', () => { const redacted = redactAccountForSidebar({ index: 0, @@ -578,32 +614,32 @@ describe('windows rework — producer seam tests', () => { remainingFraction: 0.92, resetTime: '2026-07-28T18:24:21Z', windows: [ - { - window: 'weekly', - remainingFraction: 0.92, - resetTime: '2026-07-28T18:24:21Z', - }, { window: '5h', remainingFraction: 0.99, resetTime: '2026-07-24T20:43:21Z', }, + { + window: 'weekly', + remainingFraction: 0.92, + resetTime: '2026-07-28T18:24:21Z', + }, ], }, 'non-gemini': { remainingFraction: 0.96, resetTime: '2026-07-24T18:41:52Z', windows: [ - { - window: 'weekly', - remainingFraction: 0.99, - resetTime: '2026-07-31T13:41:52Z', - }, { window: '5h', remainingFraction: 0.96, resetTime: '2026-07-24T18:41:52Z', }, + { + window: 'weekly', + remainingFraction: 0.99, + resetTime: '2026-07-31T13:41:52Z', + }, ], }, }, @@ -613,26 +649,30 @@ describe('windows rework — producer seam tests', () => { expect(gemini).toBeDefined() expect(gemini!.remainingPercent).toBe(92) expect(gemini!.windows).toHaveLength(2) - expect(gemini!.windows![0]!.window).toBe('weekly') - expect(gemini!.windows![0]!.remainingPercent).toBe(92) - expect(gemini!.windows![1]!.window).toBe('5h') - expect(gemini!.windows![1]!.remainingPercent).toBe(99) + expect(gemini!.windows![0]!.window).toBe('5h') + expect(gemini!.windows![0]!.remainingPercent).toBe(99) + expect(gemini!.windows![1]!.window).toBe('weekly') + expect(gemini!.windows![1]!.remainingPercent).toBe(92) const nonGemini = redacted.quota['non-gemini'] expect(nonGemini).toBeDefined() expect(nonGemini!.windows).toHaveLength(2) - expect(nonGemini!.windows![0]!.remainingPercent).toBe(99) - expect(nonGemini!.windows![1]!.remainingPercent).toBe(96) + expect(nonGemini!.windows![0]!.remainingPercent).toBe(96) + expect(nonGemini!.windows![1]!.remainingPercent).toBe(99) }) - it('redactAccountForSidebar handles legacy cachedQuota without windows', () => { + it('redactAccountForSidebar normalizes legacy cachedQuota keys and drops windows', () => { + // Uses real legacy keys (gemini-pro) to exercise normalizeLegacyCachedQuota, + // not the canonical 'gemini' key that bypasses the normalization path. const redacted = redactAccountForSidebar({ index: 0, cachedQuota: { - gemini: { remainingFraction: 0.5 }, + 'gemini-pro': { remainingFraction: 0.5 }, }, }) + // Legacy key must be normalized to 'gemini'. expect(redacted.quota.gemini?.remainingPercent).toBe(50) + // No windows were in the legacy snapshot — none must be invented. expect(redacted.quota.gemini?.windows).toBeUndefined() }) @@ -675,16 +715,16 @@ describe('windows rework — producer seam tests', () => { remainingFraction: 0.7, resetTime: '2026-08-01T00:00:00Z', windows: [ - { - window: 'weekly', - remainingFraction: 0.7, - resetTime: '2026-08-01T00:00:00Z', - }, { window: '5h', remainingFraction: 0.85, resetTime: '2026-07-25T00:00:00Z', }, + { + window: 'weekly', + remainingFraction: 0.7, + resetTime: '2026-08-01T00:00:00Z', + }, ], }, }, @@ -701,15 +741,103 @@ describe('windows rework — producer seam tests', () => { const gemini = state.accounts[0]!.quota.gemini expect(gemini).toBeDefined() expect(gemini!.windows).toHaveLength(2) - expect(gemini!.windows![0]!.window).toBe('weekly') - expect(gemini!.windows![0]!.remainingPercent).toBe(70) - expect(gemini!.windows![1]!.window).toBe('5h') - expect(gemini!.windows![1]!.remainingPercent).toBe(85) + expect(gemini!.windows![0]!.window).toBe('5h') + expect(gemini!.windows![0]!.remainingPercent).toBe(85) + expect(gemini!.windows![1]!.window).toBe('weekly') + expect(gemini!.windows![1]!.remainingPercent).toBe(70) } finally { rmSync(root, { recursive: true, force: true }) } }) + it('normalizeLegacyCachedQuota: gemini-pro + gemini-flash collapse to gemini with MIN fraction', () => { + // M5: verify the two legacy gemini keys collapse with the most-constrained + // (lowest remainingFraction) winning. Also verifies the M4 resetTime fix: + // the winner (gemini-flash, 0.3) carries resetTime=Aug-01; gemini-pro (0.5) + // has an earlier resetTime=Jul-26 — the merged result must use Jul-26. + const result = normalizeLegacyCachedQuota({ + 'gemini-pro': { + remainingFraction: 0.5, + resetTime: '2026-07-26T00:00:00Z', + }, + 'gemini-flash': { + remainingFraction: 0.3, + resetTime: '2026-08-01T00:00:00Z', + }, + }) + expect(result?.gemini?.remainingFraction).toBe(0.3) // gemini-flash wins fraction + // M4 regression: the winner\'s resetTime (Aug-01) is LATER than the loser\'s + // (Jul-26). The merged result must carry the earlier one (Jul-26). + expect(result?.gemini?.resetTime).toBe('2026-07-26T00:00:00Z') + // No windows must be invented for migrated entries. + expect(result?.gemini?.windows).toBeUndefined() + }) + + it('normalizeLegacyCachedQuota: M4 a-wins branch also merges resetTime', () => { + // Specifically tests the a-wins path (fa <= fb): gemini-pro has lower + // fraction than gemini-flash AND an earlier resetTime — both must survive. + const result = normalizeLegacyCachedQuota({ + 'gemini-pro': { + remainingFraction: 0.2, // wins + resetTime: '2026-07-26T00:00:00Z', // earlier + }, + 'gemini-flash': { + remainingFraction: 0.5, + resetTime: '2026-08-05T00:00:00Z', + }, + }) + expect(result?.gemini?.remainingFraction).toBe(0.2) + // Even though a wins, the merged resetTime must be the earlier of the two. + expect(result?.gemini?.resetTime).toBe('2026-07-26T00:00:00Z') + }) + + it('normalizeLegacyCachedQuota: a-wins branch with loser having earlier resetTime (M4 bug scenario)', () => { + // This is the exact M4 bug: a wins the fraction comparison but b has the + // earlier resetTime. Before the fix, b\'s resetTime was silently dropped. + const result = normalizeLegacyCachedQuota({ + 'gemini-pro': { + remainingFraction: 0.3, // wins (lower) + resetTime: '2026-08-01T00:00:00Z', // later + }, + 'gemini-flash': { + remainingFraction: 0.5, + resetTime: '2026-07-26T00:00:00Z', // earlier — must survive + }, + }) + expect(result?.gemini?.remainingFraction).toBe(0.3) + // M4 fix: loser\'s earlier resetTime must win over the winner\'s later one. + expect(result?.gemini?.resetTime).toBe('2026-07-26T00:00:00Z') + }) + + it('normalizeLegacyCachedQuota: claude + gpt-oss collapse to non-gemini with MIN fraction', () => { + const result = normalizeLegacyCachedQuota({ + claude: { remainingFraction: 0.7, resetTime: '2026-08-01T00:00:00Z' }, + 'gpt-oss': { remainingFraction: 0.4, resetTime: '2026-07-28T00:00:00Z' }, + }) + expect(result?.['non-gemini']?.remainingFraction).toBe(0.4) // gpt-oss wins + expect(result?.['non-gemini']?.resetTime).toBe('2026-07-28T00:00:00Z') // earlier + }) + + it('normalizeLegacyCachedQuota: no windows invented for migrated entries', () => { + // Legacy entries never had windows; the normalizer must not synthesize them. + const result = normalizeLegacyCachedQuota({ + 'gemini-pro': { remainingFraction: 0.6 }, + claude: { remainingFraction: 0.8 }, + }) + expect(result?.gemini?.windows).toBeUndefined() + expect(result?.['non-gemini']?.windows).toBeUndefined() + }) + + it('normalizeLegacyCachedQuota: canonical keys pass through unchanged', () => { + // Non-legacy input must be returned as-is without any transformation. + const canonical = { + gemini: { remainingFraction: 0.6, resetTime: '2026-08-01T00:00:00Z' }, + 'non-gemini': { remainingFraction: 0.4 }, + } + const result = normalizeLegacyCachedQuota(canonical) + expect(result).toBe(canonical) // same reference — fast path returned as-is + }) + it('identity mismatch drops the entire cachedQuota including windows', () => { const redacted = redactAccountForSidebar({ index: 0, @@ -732,3 +860,92 @@ describe('windows rework — producer seam tests', () => { expect(Object.keys(redacted.quota)).toHaveLength(0) }) }) + +describe('render boundary — current flag', () => { + let fixture: Fixture + + beforeEach(() => { + fixture = makeFixture() + }) + + afterEach(() => { + fixture.cleanup() + }) + + it('preserves current: true from producer input through to sidebar state', async () => { + const state = buildSidebarMachineStateFromAccounts([ + { index: 0, current: true }, + { index: 1, current: false }, + ]) + // The render boundary is `redactAccountForSidebar` inside + // `buildSidebarMachineStateFromAccounts`. A `current: true` from any + // producer must survive to the state the TUI reads so that + // `visibleAccounts` (when fallbackAccounts is off) renders the + // active account instead of an empty list. + expect(state.accounts).toHaveLength(2) + expect(state.accounts[0]?.current).toBe(true) + expect(state.accounts[1]?.current).toBe(false) + + // Also exercise the full write→read round-trip the TUI actually uses. + await setSidebarMachineState(state, { stateFile: fixture.stateFile }) + const disk = readSidebarState(fixture.stateFile) + expect(disk.accounts[0]?.current).toBe(true) + expect(disk.accounts[1]?.current).toBe(false) + }) + + it('without fallbackAccounts, current: false for all accounts → no accounts visible', () => { + // The filter the TUI runs is: + // state.accounts.filter(account => account.current) + // When every account has current: false, the result is empty — that + // was the live bug before this fix. + const state = buildSidebarMachineStateFromAccounts([ + { index: 0, current: false }, + { index: 1, current: false }, + ]) + const visible = state.accounts.filter((entry) => entry.current) + expect(visible).toHaveLength(0) + }) + + it('without fallbackAccounts, current: true on one account → that account is visible', () => { + const state = buildSidebarMachineStateFromAccounts([ + { index: 0, current: true }, + { index: 1, current: false }, + ]) + const visible = state.accounts.filter((entry) => entry.current) + expect(visible).toHaveLength(1) + expect(visible[0]?.id).toBe('acct-0') + }) + + // ── Fix 2: tier round-trip ────────────────────────────────────────────────── + + it('Fix2-tier-roundtrip: tier writes to disk and reads back with capturedAt', async () => { + // An account with a captured tier must survive a full write→read cycle. + const capturedAt = 1_785_005_949_000 + const state = buildSidebarMachineStateFromAccounts([ + { index: 0, tier: { id: 'free-tier', capturedAt } }, + ]) + await setSidebarMachineState(state, { stateFile: fixture.stateFile }) + const disk = readSidebarState(fixture.stateFile) + expect(disk.accounts[0]?.tier).toEqual({ id: 'free-tier', capturedAt }) + }) + + it('Fix2-tier-absent: account without tier yields no tier key', () => { + // An account with no tier in the redaction input must not emit tier in the + // rendered state — absent ≠ free, absent ≠ {} . + const state = buildSidebarMachineStateFromAccounts([{ index: 0 }]) + expect(state.accounts[0]?.tier).toBeUndefined() + // Confirm `tier` is not a key at all in serialized form. + const serialized = JSON.stringify(state.accounts[0]) + expect(serialized).not.toContain('tier') + }) + + it('Fix2-tier-redaction: tier survives redactAccountForSidebar (not PII)', () => { + // Plan tier is metadata, not a credential — it must pass the redaction boundary. + const capturedAt = 1_785_000_000_000 + const result = redactAccountForSidebar({ + index: 0, + tier: { id: 'paid-tier', capturedAt }, + }) + expect(result.tier).toEqual({ id: 'paid-tier', capturedAt }) + }) +}) diff --git a/packages/opencode/src/sidebar-state.ts b/packages/opencode/src/sidebar-state.ts index 8eedca4..2dbcdfc 100644 --- a/packages/opencode/src/sidebar-state.ts +++ b/packages/opencode/src/sidebar-state.ts @@ -76,6 +76,11 @@ export interface SidebarAccountState { current: boolean cooldownUntil?: number quota: Partial> + /** + * Captured plan tier. Absent when unknown — do NOT default to `free-tier`. + * `id` is the raw upstream string (e.g. `"free-tier"`); never normalised. + */ + tier?: { id: string; paidId?: string; capturedAt: number } } export interface SidebarRoutingEntry { @@ -305,6 +310,7 @@ function normalizeAccount(input: unknown): SidebarAccountState | null { if (normalized) quota[key] = normalized } } + const tier = normalizeTier(record.tier) return { id, label, @@ -313,9 +319,26 @@ function normalizeAccount(input: unknown): SidebarAccountState | null { current, cooldownUntil, quota, + ...(tier !== undefined ? { tier } : {}), } } +function normalizeTier( + input: unknown, +): { id: string; paidId?: string; capturedAt: number } | undefined { + if (!isObject(input)) return undefined + const record = input as Record + const id = + typeof record.id === 'string' && record.id.length > 0 ? record.id : null + const capturedAt = toFiniteNumber(record.capturedAt) + if (!id || capturedAt === null) return undefined + const paidId = + typeof record.paidId === 'string' && record.paidId.length > 0 + ? record.paidId + : undefined + return { id, ...(paidId ? { paidId } : {}), capturedAt } +} + function normalizeQuota(input: unknown): SidebarQuotaEntry | null { if (!isObject(input)) return null const record = input as Record @@ -484,6 +507,21 @@ export interface SidebarAccountRedactionInput { resetTime: string }> } + // Legacy per-pool keys written by older versions of the plugin. + // Accepted here so the normalizer can map them to the current pool + // keys at read time without requiring a disk migration. + [legacyKey: string]: + | { + remainingFraction?: number + resetTime?: string + modelCount?: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } + | undefined } /** * Opaque identity stamp that was attached to the persisted quota snapshot. @@ -499,6 +537,119 @@ export interface SidebarAccountRedactionInput { * command-data service. Omitted in the persisted sidebar state. */ currentQuotaAccountId?: string + /** + * Captured plan tier from `loadCodeAssist`. Absent when unknown. The `id` + * is the raw upstream string; `capturedAt` is epoch ms. Not PII — tier + * metadata survives the redaction boundary unchanged. + */ + tier?: { id: string; paidId?: string; capturedAt: number } +} + +/** + * Build a `{ id, capturedAt }` tier object from account fields, or `undefined` + * when either field is absent. Centralised so every producer (index.ts, + * background-quota-refresh.ts, command-data.ts) uses the same guard and the + * same shape -- the same pattern shipped as inline literals across three files + * and produced six missed-field bugs; one helper ends that. + */ +export function toCapturedTier(account: { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number +}): { id: string; paidId?: string; capturedAt: number } | undefined { + return account.capturedTierId !== undefined && + account.capturedTierAt !== undefined + ? { + id: account.capturedTierId, + ...(account.capturedPaidTierId !== undefined + ? { paidId: account.capturedPaidTierId } + : {}), + capturedAt: account.capturedTierAt, + } + : undefined +} + +/** + * Map pre-pool legacy quota keys to the current two-pool schema at read time. + * + * Legacy mappings (empirically settled from burn tests): + * `gemini-pro` + `gemini-flash` → `gemini` + * `claude` + `gpt-oss` → `non-gemini` + * + * Where multiple legacy keys collapse into one pool the MIN remainingFraction + * and earliest resetTime are used (most-constrained-first, matching the + * window-level rule). No `windows` arrays are invented for migrated entries; + * the single aggregate bar is the correct render for pre-window snapshots. + * + * Non-legacy keys that are already canonical (`gemini`, `non-gemini`) are + * carried through unchanged. The next real quota refresh overwrites any + * migrated snapshot with authoritative data. + */ +export function normalizeLegacyCachedQuota( + raw: SidebarAccountRedactionInput['cachedQuota'], +): SidebarAccountRedactionInput['cachedQuota'] { + if (!raw) return raw + + // Pool already carries current keys — fast path when no legacy keys present. + const hasLegacy = + 'gemini-pro' in raw || + 'gemini-flash' in raw || + 'claude' in raw || + 'gpt-oss' in raw + if (!hasLegacy) return raw + + type PoolEntry = NonNullable< + SidebarAccountRedactionInput['cachedQuota'] + >['gemini'] + + // Pick the earlier of two ISO reset timestamps, preferring a defined + // value over undefined. Extracted so both branches of minFraction use + // the same merge rather than the previous asymmetry where the a-wins + // branch dropped b's possibly-earlier resetTime. + const earlierResetTime = ( + a: string | undefined, + b: string | undefined, + ): string | undefined => { + if (!a) return b + if (!b) return a + return a < b ? a : b + } + + const minFraction = (a: PoolEntry, b: PoolEntry): PoolEntry => { + if (!a) return b + if (!b) return a + const fa = a.remainingFraction ?? 1 + const fb = b.remainingFraction ?? 1 + const winner = fa <= fb ? a : b + const loser = fa <= fb ? b : a + // Always merge to the earliest resetTime regardless of which pool wins + // the fraction comparison — an earlier window expiry from the loser + // pool must not be silently dropped. + return { + ...winner, + resetTime: earlierResetTime(winner.resetTime, loser.resetTime), + } + } + + const gemini: PoolEntry = + minFraction( + raw.gemini, + minFraction( + raw['gemini-pro'] as PoolEntry, + raw['gemini-flash'] as PoolEntry, + ), + ) ?? raw.gemini + + const nonGemini: PoolEntry = + minFraction( + raw['non-gemini'], + minFraction(raw.claude as PoolEntry, raw['gpt-oss'] as PoolEntry), + ) ?? raw['non-gemini'] + + return { + ...(gemini !== undefined ? { gemini } : {}), + ...(nonGemini !== undefined ? { 'non-gemini': nonGemini } : {}), + } } /** @@ -554,6 +705,20 @@ export function projectQuotaPoolForSidebar(source: { return { remainingPercent, resetAt, windows } } +/** + * Returns `true` when an account at `index` is the active account for + * at least one model family. Two accounts CAN both be current — one + * serving claude, one serving gemini — so this must not collapse to one. + */ +export function isAccountCurrent( + index: number, + activeIndexByFamily: { claude: number; gemini: number }, +): boolean { + return ( + index === activeIndexByFamily.claude || index === activeIndexByFamily.gemini + ) +} + /** * Convert a live account snapshot into the redacted shape the TUI renders. * The redacted `SidebarAccountState` carries no email, refresh token, access @@ -573,7 +738,7 @@ export function redactAccountForSidebar( ? source.coolingDownUntil : undefined const health = clampNumber( - typeof source.healthScore === 'number' ? source.healthScore : null, + typeof source.healthScore === 'number' ? source.healthScore : 100, 0, 100, ) @@ -587,7 +752,12 @@ export function redactAccountForSidebar( typeof source.cachedQuotaAccountId === 'string' && typeof source.currentQuotaAccountId === 'string' && source.cachedQuotaAccountId !== source.currentQuotaAccountId - const cached = staleCachedQuota ? undefined : source.cachedQuota + // Normalize legacy per-pool keys to the current two-pool schema before + // reading. This is non-destructive: the next real quota refresh will + // overwrite the migrated snapshot with authoritative data. + const cached = staleCachedQuota + ? undefined + : normalizeLegacyCachedQuota(source.cachedQuota) if (cached) { for (const key of ['gemini', 'non-gemini'] as const) { const entry = cached[key] @@ -605,6 +775,8 @@ export function redactAccountForSidebar( current, cooldownUntil, quota, + // Tier is plan metadata, not PII — it passes the redaction boundary. + ...(source.tier !== undefined ? { tier: source.tier } : {}), } } diff --git a/packages/opencode/src/tui-compiled/plugin/command-data.ts b/packages/opencode/src/tui-compiled/plugin/command-data.ts index e517ca4..eb72733 100644 --- a/packages/opencode/src/tui-compiled/plugin/command-data.ts +++ b/packages/opencode/src/tui-compiled/plugin/command-data.ts @@ -44,8 +44,10 @@ import { createHash } from 'node:crypto' import { buildSidebarMachineStateFromAccounts, + normalizeLegacyCachedQuota, type SidebarAccountRedactionInput, setSidebarMachineState, + toCapturedTier, } from '../sidebar-state' /** @@ -143,6 +145,8 @@ export interface CommandAccountRow { * without relying on unstable numeric indices. */ coolingDownUntil?: number + /** Health score in [0, 100], forwarded from the live tracker. */ + healthScore?: number quota: Array<{ key: 'gemini' | 'non-gemini' label: string @@ -154,6 +158,8 @@ export interface CommandAccountRow { resetAt?: number }> }> + /** Captured plan tier. Absent when unknown; never defaulted. */ + tier?: { id: string; capturedAt: number } } /** @@ -186,6 +192,13 @@ interface LiveAccountSnapshot { cachedQuotaAccountId?: string accountIneligible?: boolean coolingDownUntil?: number + healthScore?: number + /** Captured plan tier from the most recent loadCodeAssist response. */ + capturedTierId?: string + /** Captured paid-tier ID from the most recent loadCodeAssist response. */ + capturedPaidTierId?: string + /** Epoch ms when capturedTierId was last recorded. */ + capturedTierAt?: number } function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { @@ -193,11 +206,19 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { // (the refresh token changed, or an index shift placed another account's // snapshot at this position). Drop the stale cache rather than rendering // the wrong account's quota percentages. - const cached = + const rawCached = entry.cachedQuotaAccountId && entry.cachedQuotaAccountId !== quotaAccountIdentity(entry.refreshToken) ? undefined : entry.cachedQuota + // Normalize legacy pool keys (gemini-pro, gemini-flash, claude, gpt-oss) + // to canonical keys before the SUPPORTED_QUOTA_KEYS loop. Without this + // a legacy snapshot produces empty quota rows, and writeSidebar would + // permanently re-emit legacy keys on every dialog open until the next + // real quota refresh rewrites the on-disk snapshot. + const cached = normalizeLegacyCachedQuota( + rawCached as Parameters[0], + ) const quota: CommandAccountRow['quota'] = [] for (const key of SUPPORTED_QUOTA_KEYS) { const cachedEntry = cached?.[key] @@ -246,7 +267,9 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { enabled: entry.enabled, current: entry.active, coolingDownUntil: entry.coolingDownUntil, + healthScore: entry.healthScore, quota, + tier: toCapturedTier(entry), } } @@ -399,8 +422,20 @@ export interface CommandDataService { * refresh token, bump `cachedQuotaUpdatedAt`, and push a label-only * sidebar snapshot. Returns the freshly persisted rows so the * dialog can re-render in place. + * + * Bypass flag is intentional: explicit user-triggered refreshes + * ("Refresh" button, direct slash-command with `refresh` argument) + * must always fetch, even if the quota manager has backed off. */ refreshQuota(): Promise + /** + * Non-forced quota check that RESPECTS the quota manager's per-account + * backoff. Used by the dialog OPEN path — which is a VIEW, not a user + * request to fetch. Accounts already fresh (or in backoff) are skipped; + * the result is discarded (the sidebar poller carries the freshness + * guarantee for idle accounts). + */ + refreshQuotaRespectingBackoff(): Promise /** * Pin `index` as the active account for every family the dialog * tracks (claude + gemini). Mutates the live AccountManager AND the @@ -452,54 +487,63 @@ export function createCommandDataService( const projectRows = (): CommandAccountRow[] => accountManagerView.getAccounts().map(toCommandAccountRow) - const writeSidebar = (rows: CommandAccountRow[]): void => { - const accounts: SidebarAccountRedactionInput[] = rows.map((row) => { - const gemini = row.quota.find((q) => q.key === 'gemini') - const nonGemini = row.quota.find((q) => q.key === 'non-gemini') - // Map a CommandAccountRow quota entry → cachedQuota pool shape - // so projectQuotaPoolForSidebar (the canonical projection seam) - // carries the per-window breakdown into the sidebar state. - const toPool = ( - q: - | { - remainingPercent: number | null - resetAt?: number - windows?: CommandAccountRow['quota'][number]['windows'] - } - | undefined, - ) => { - if (!q || q.remainingPercent == null) return undefined - return { - remainingFraction: q.remainingPercent / 100, - resetTime: - typeof q.resetAt === 'number' && Number.isFinite(q.resetAt) - ? new Date(q.resetAt).toISOString() - : undefined, - windows: q.windows?.map((w) => ({ - window: w.window, - remainingFraction: w.remainingPercent / 100, - resetTime: - typeof w.resetAt === 'number' && Number.isFinite(w.resetAt) - ? new Date(w.resetAt).toISOString() - : '', - })), - } - } + // Map a CommandAccountRow back to the SidebarAccountRedactionInput shape. + // Extracted to prevent the "sixth dropped field" pattern: every field in + // CommandAccountRow that must reach the sidebar lives here once, not as an + // inline literal that a future author has to remember to update in two places. + const toRedactionInput = ( + row: CommandAccountRow, + ): SidebarAccountRedactionInput => { + const gemini = row.quota.find((q) => q.key === 'gemini') + const nonGemini = row.quota.find((q) => q.key === 'non-gemini') + // Map a CommandAccountRow quota entry to cachedQuota pool shape + // so projectQuotaPoolForSidebar (the canonical projection seam) + // carries the per-window breakdown into the sidebar state. + const toPool = ( + q: + | { + remainingPercent: number | null + resetAt?: number + windows?: CommandAccountRow['quota'][number]['windows'] + } + | undefined, + ) => { + if (!q || q.remainingPercent == null) return undefined return { - index: row.index, - label: row.label, - enabled: row.enabled, - current: row.current, - cachedQuota: { - gemini: toPool(gemini), - 'non-gemini': toPool(nonGemini), - }, - // The stamp check has already been done by `toCommandAccountRow` - // before the rows reach this writer; nothing further for the - // sidebar projection to validate here. + remainingFraction: q.remainingPercent / 100, + resetTime: + typeof q.resetAt === 'number' && Number.isFinite(q.resetAt) + ? new Date(q.resetAt).toISOString() + : undefined, + windows: q.windows?.map((w) => ({ + window: w.window, + remainingFraction: w.remainingPercent / 100, + resetTime: + typeof w.resetAt === 'number' && Number.isFinite(w.resetAt) + ? new Date(w.resetAt).toISOString() + : '', + })), } - }) - // Fire-and-forget — the sidebar writer is fenced by its own queue, + } + return { + index: row.index, + label: row.label, + enabled: row.enabled, + current: row.current, + healthScore: row.healthScore, + // Tier passes the redaction boundary unchanged (plan metadata, not PII). + // The stamp check was already done by toCommandAccountRow. + tier: row.tier, + cachedQuota: { + gemini: toPool(gemini), + 'non-gemini': toPool(nonGemini), + }, + } + } + + const writeSidebar = (rows: CommandAccountRow[]): void => { + const accounts = rows.map(toRedactionInput) + // Fire-and-forget -- the sidebar writer is fenced by its own queue, // so a transient lock contention cannot block the dialog response. void setSidebarMachineState( buildSidebarMachineStateFromAccounts(accounts, { checkedAt: now() }), @@ -611,6 +655,77 @@ export function createCommandDataService( return rows }, + async refreshQuotaRespectingBackoff() { + const accountsForQuota = accountManagerView.getAccountsForQuotaCheck() + if (accountsForQuota.length === 0) return + + // Non-forced: the quota manager's per-account backoff and dedup + // apply. Accounts that are fresh or in backoff are skipped. Unlike + // the earlier fire-and-forget path, successful results ARE folded + // into the live cache so the sidebar snapshot reflects fresh data + // immediately — the dialog renders ONLY from the sidebar file, and + // the quota manager's push path does not update AccountManager + // entries for this caller. + let results: Awaited> + try { + results = await quotaManager.refreshAccounts(accountsForQuota, { + indexFor: (account) => accountsForQuota.indexOf(account), + force: false, + }) + } catch { + // Non-critical: dialog open must never fail because a background + // refresh check encountered backoff or a transient error. + return + } + + // Build updates akin to the forced refresh path — resolve live + // indexes after the network call (a concurrent add/remove may have + // shifted positions), fold successful results into the live cache, + // then push a sidebar snapshot. + const updates: Array<{ + refreshToken: string + groups?: Partial< + Record + > + }> = [] + for (const result of results) { + const refreshToken = + result.updatedAccount?.refreshToken ?? + accountsForQuota[result.index]?.refreshToken + if (!refreshToken) continue + const groups = + result.status === 'ok' && result.quota?.groups + ? result.quota.groups + : undefined + updates.push({ refreshToken, groups }) + } + + if (updates.length === 0) return + + const liveIndexByRefreshToken = new Map() + for (const entry of accountManagerView.getAccounts()) { + liveIndexByRefreshToken.set(entry.refreshToken, entry.index) + } + let liveQuotaChanged = false + for (const update of updates) { + const liveIndex = liveIndexByRefreshToken.get(update.refreshToken) + if (liveIndex === undefined || !update.groups) continue + accountManagerView.updateQuotaCache( + liveIndex, + update.groups, + update.refreshToken, + ) + liveQuotaChanged = true + } + if (liveQuotaChanged) accountManagerView.requestSaveToDisk() + + // Push a fresh sidebar snapshot — the live view now carries the + // just-refreshed quota plus any skipped (backoff/fresh) accounts + // whose cached percentages are unchanged. + const rows = projectRows() + writeSidebar(rows) + }, + async setCurrentAccount(index) { return mutateLiveAndStorage({ action: 'setCurrent', index }) }, @@ -738,7 +853,12 @@ export function createCommandDataService( // (the core's buildStorageSnapshot clamps negative sentinels // to 0, and auth-doctor treats a negative activeIndex as // corruption — cf. auth-doctor.ts:149-156). - return index < nextAccounts.length ? index : 0 + // + // Use tokenIdx (the storage-snapshot position) rather than the + // outer-scope `index` (live-read position): a concurrent mutation + // between the live read and the lock can diverge the two, and the + // storage-level cursor must match the storage-level array. + return tokenIdx < nextAccounts.length ? tokenIdx : 0 } const legacyClaude = current.activeIndexByFamily?.claude ?? current.activeIndex diff --git a/packages/opencode/src/tui-compiled/sidebar-state.ts b/packages/opencode/src/tui-compiled/sidebar-state.ts index 8eedca4..2dbcdfc 100644 --- a/packages/opencode/src/tui-compiled/sidebar-state.ts +++ b/packages/opencode/src/tui-compiled/sidebar-state.ts @@ -76,6 +76,11 @@ export interface SidebarAccountState { current: boolean cooldownUntil?: number quota: Partial> + /** + * Captured plan tier. Absent when unknown — do NOT default to `free-tier`. + * `id` is the raw upstream string (e.g. `"free-tier"`); never normalised. + */ + tier?: { id: string; paidId?: string; capturedAt: number } } export interface SidebarRoutingEntry { @@ -305,6 +310,7 @@ function normalizeAccount(input: unknown): SidebarAccountState | null { if (normalized) quota[key] = normalized } } + const tier = normalizeTier(record.tier) return { id, label, @@ -313,9 +319,26 @@ function normalizeAccount(input: unknown): SidebarAccountState | null { current, cooldownUntil, quota, + ...(tier !== undefined ? { tier } : {}), } } +function normalizeTier( + input: unknown, +): { id: string; paidId?: string; capturedAt: number } | undefined { + if (!isObject(input)) return undefined + const record = input as Record + const id = + typeof record.id === 'string' && record.id.length > 0 ? record.id : null + const capturedAt = toFiniteNumber(record.capturedAt) + if (!id || capturedAt === null) return undefined + const paidId = + typeof record.paidId === 'string' && record.paidId.length > 0 + ? record.paidId + : undefined + return { id, ...(paidId ? { paidId } : {}), capturedAt } +} + function normalizeQuota(input: unknown): SidebarQuotaEntry | null { if (!isObject(input)) return null const record = input as Record @@ -484,6 +507,21 @@ export interface SidebarAccountRedactionInput { resetTime: string }> } + // Legacy per-pool keys written by older versions of the plugin. + // Accepted here so the normalizer can map them to the current pool + // keys at read time without requiring a disk migration. + [legacyKey: string]: + | { + remainingFraction?: number + resetTime?: string + modelCount?: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } + | undefined } /** * Opaque identity stamp that was attached to the persisted quota snapshot. @@ -499,6 +537,119 @@ export interface SidebarAccountRedactionInput { * command-data service. Omitted in the persisted sidebar state. */ currentQuotaAccountId?: string + /** + * Captured plan tier from `loadCodeAssist`. Absent when unknown. The `id` + * is the raw upstream string; `capturedAt` is epoch ms. Not PII — tier + * metadata survives the redaction boundary unchanged. + */ + tier?: { id: string; paidId?: string; capturedAt: number } +} + +/** + * Build a `{ id, capturedAt }` tier object from account fields, or `undefined` + * when either field is absent. Centralised so every producer (index.ts, + * background-quota-refresh.ts, command-data.ts) uses the same guard and the + * same shape -- the same pattern shipped as inline literals across three files + * and produced six missed-field bugs; one helper ends that. + */ +export function toCapturedTier(account: { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number +}): { id: string; paidId?: string; capturedAt: number } | undefined { + return account.capturedTierId !== undefined && + account.capturedTierAt !== undefined + ? { + id: account.capturedTierId, + ...(account.capturedPaidTierId !== undefined + ? { paidId: account.capturedPaidTierId } + : {}), + capturedAt: account.capturedTierAt, + } + : undefined +} + +/** + * Map pre-pool legacy quota keys to the current two-pool schema at read time. + * + * Legacy mappings (empirically settled from burn tests): + * `gemini-pro` + `gemini-flash` → `gemini` + * `claude` + `gpt-oss` → `non-gemini` + * + * Where multiple legacy keys collapse into one pool the MIN remainingFraction + * and earliest resetTime are used (most-constrained-first, matching the + * window-level rule). No `windows` arrays are invented for migrated entries; + * the single aggregate bar is the correct render for pre-window snapshots. + * + * Non-legacy keys that are already canonical (`gemini`, `non-gemini`) are + * carried through unchanged. The next real quota refresh overwrites any + * migrated snapshot with authoritative data. + */ +export function normalizeLegacyCachedQuota( + raw: SidebarAccountRedactionInput['cachedQuota'], +): SidebarAccountRedactionInput['cachedQuota'] { + if (!raw) return raw + + // Pool already carries current keys — fast path when no legacy keys present. + const hasLegacy = + 'gemini-pro' in raw || + 'gemini-flash' in raw || + 'claude' in raw || + 'gpt-oss' in raw + if (!hasLegacy) return raw + + type PoolEntry = NonNullable< + SidebarAccountRedactionInput['cachedQuota'] + >['gemini'] + + // Pick the earlier of two ISO reset timestamps, preferring a defined + // value over undefined. Extracted so both branches of minFraction use + // the same merge rather than the previous asymmetry where the a-wins + // branch dropped b's possibly-earlier resetTime. + const earlierResetTime = ( + a: string | undefined, + b: string | undefined, + ): string | undefined => { + if (!a) return b + if (!b) return a + return a < b ? a : b + } + + const minFraction = (a: PoolEntry, b: PoolEntry): PoolEntry => { + if (!a) return b + if (!b) return a + const fa = a.remainingFraction ?? 1 + const fb = b.remainingFraction ?? 1 + const winner = fa <= fb ? a : b + const loser = fa <= fb ? b : a + // Always merge to the earliest resetTime regardless of which pool wins + // the fraction comparison — an earlier window expiry from the loser + // pool must not be silently dropped. + return { + ...winner, + resetTime: earlierResetTime(winner.resetTime, loser.resetTime), + } + } + + const gemini: PoolEntry = + minFraction( + raw.gemini, + minFraction( + raw['gemini-pro'] as PoolEntry, + raw['gemini-flash'] as PoolEntry, + ), + ) ?? raw.gemini + + const nonGemini: PoolEntry = + minFraction( + raw['non-gemini'], + minFraction(raw.claude as PoolEntry, raw['gpt-oss'] as PoolEntry), + ) ?? raw['non-gemini'] + + return { + ...(gemini !== undefined ? { gemini } : {}), + ...(nonGemini !== undefined ? { 'non-gemini': nonGemini } : {}), + } } /** @@ -554,6 +705,20 @@ export function projectQuotaPoolForSidebar(source: { return { remainingPercent, resetAt, windows } } +/** + * Returns `true` when an account at `index` is the active account for + * at least one model family. Two accounts CAN both be current — one + * serving claude, one serving gemini — so this must not collapse to one. + */ +export function isAccountCurrent( + index: number, + activeIndexByFamily: { claude: number; gemini: number }, +): boolean { + return ( + index === activeIndexByFamily.claude || index === activeIndexByFamily.gemini + ) +} + /** * Convert a live account snapshot into the redacted shape the TUI renders. * The redacted `SidebarAccountState` carries no email, refresh token, access @@ -573,7 +738,7 @@ export function redactAccountForSidebar( ? source.coolingDownUntil : undefined const health = clampNumber( - typeof source.healthScore === 'number' ? source.healthScore : null, + typeof source.healthScore === 'number' ? source.healthScore : 100, 0, 100, ) @@ -587,7 +752,12 @@ export function redactAccountForSidebar( typeof source.cachedQuotaAccountId === 'string' && typeof source.currentQuotaAccountId === 'string' && source.cachedQuotaAccountId !== source.currentQuotaAccountId - const cached = staleCachedQuota ? undefined : source.cachedQuota + // Normalize legacy per-pool keys to the current two-pool schema before + // reading. This is non-destructive: the next real quota refresh will + // overwrite the migrated snapshot with authoritative data. + const cached = staleCachedQuota + ? undefined + : normalizeLegacyCachedQuota(source.cachedQuota) if (cached) { for (const key of ['gemini', 'non-gemini'] as const) { const entry = cached[key] @@ -605,6 +775,8 @@ export function redactAccountForSidebar( current, cooldownUntil, quota, + // Tier is plan metadata, not PII — it passes the redaction boundary. + ...(source.tier !== undefined ? { tier: source.tier } : {}), } } diff --git a/packages/opencode/src/tui-compiled/tui.tsx b/packages/opencode/src/tui-compiled/tui.tsx index 710f5dd..5e06ed5 100644 --- a/packages/opencode/src/tui-compiled/tui.tsx +++ b/packages/opencode/src/tui-compiled/tui.tsx @@ -866,8 +866,13 @@ export function SidebarPanel(props) { const account = () => visibleAccounts().find(entry => entry.current) ?? visibleAccounts()[0]; const used = () => { const q = account()?.quota; - const entry = q?.gemini ?? q?.['non-gemini']; - return entry ? 100 - clamp(entry.remainingPercent, 0, 100) : null; + // Derive tone from the WORST (most-constrained) pool so a + // critical non-gemini pool doesn't hide behind a healthy + // gemini tone. Worst = highest used-percent across available pools. + const geminiUsed = q?.gemini != null ? 100 - clamp(q.gemini.remainingPercent, 0, 100) : null; + const nonGeminiUsed = q?.['non-gemini'] != null ? 100 - clamp(q['non-gemini'].remainingPercent, 0, 100) : null; + if (geminiUsed === null && nonGeminiUsed === null) return null; + return Math.max(geminiUsed ?? 0, nonGeminiUsed ?? 0); }; const poolText = key => { const entry = account()?.quota[key]; diff --git a/packages/opencode/src/tui-windows-frames.test.tsx b/packages/opencode/src/tui-windows-frames.test.tsx index 780907b..5b530a2 100644 --- a/packages/opencode/src/tui-windows-frames.test.tsx +++ b/packages/opencode/src/tui-windows-frames.test.tsx @@ -42,6 +42,11 @@ function makeFixture(): Fixture { const root = mkdtempSync(join(tmpdir(), 'agy-tui-windows-')) const statePath = join(root, 'sidebar-state.json') const prefsPath = join(root, 'tui-preferences.jsonc') + // Save preload-pinned values to restore on cleanup instead of deleting — + // a delete drops resolution to the operator's real state/config dirs. + const savedSidebar = process.env[SIDEBAR_STATE_ENV] + const savedTuiLog = process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE + const savedPrefs = process.env[TUI_PREFS_FILE_ENV] process.env[SIDEBAR_STATE_ENV] = statePath process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE = join(root, 'tui.log') process.env[TUI_PREFS_FILE_ENV] = prefsPath @@ -49,9 +54,14 @@ function makeFixture(): Fixture { statePath, prefsPath, cleanup: () => { - delete process.env[SIDEBAR_STATE_ENV] - delete process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE - delete process.env[TUI_PREFS_FILE_ENV] + if (savedSidebar !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebar + else delete process.env[SIDEBAR_STATE_ENV] + if (savedTuiLog !== undefined) + process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE = savedTuiLog + else delete process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE + if (savedPrefs !== undefined) process.env[TUI_PREFS_FILE_ENV] = savedPrefs + else delete process.env[TUI_PREFS_FILE_ENV] rmSync(root, { recursive: true, force: true }) }, } @@ -132,16 +142,16 @@ describe('windows rework — reviewer frames', () => { remainingPercent: 92, resetAt: future2, windows: [ - { window: 'weekly', remainingPercent: 92, resetAt: future2 }, { window: '5h', remainingPercent: 99, resetAt: future }, + { window: 'weekly', remainingPercent: 92, resetAt: future2 }, ], }, 'non-gemini': { remainingPercent: 96, resetAt: future2, windows: [ - { window: 'weekly', remainingPercent: 99, resetAt: future2 }, { window: '5h', remainingPercent: 96, resetAt: future }, + { window: 'weekly', remainingPercent: 99, resetAt: future2 }, ], }, }, @@ -167,6 +177,12 @@ describe('windows rework — reviewer frames', () => { expect(frame).toContain('Gm 5h') expect(frame).toContain('NG 7d') expect(frame).toContain('NG 5h') + // Window order: aggregator now sorts shortest-first so 5h must + // render before 7d. The fixture carries windows in [5h, weekly] + // order (matching aggregateQuotaSummary output) and the renderer + // preserves input order, so the assertion is deterministic. + expect(frame.indexOf('Gm 5h')).toBeLessThan(frame.indexOf('Gm 7d')) + expect(frame.indexOf('NG 5h')).toBeLessThan(frame.indexOf('NG 7d')) testSetup.renderer.destroy() }) @@ -304,16 +320,16 @@ describe('windows rework — reviewer frames', () => { remainingPercent: 92, resetAt: future2, windows: [ - { window: 'weekly', remainingPercent: 92, resetAt: future2 }, { window: '5h', remainingPercent: 99, resetAt: future }, + { window: 'weekly', remainingPercent: 92, resetAt: future2 }, ], }, 'non-gemini': { remainingPercent: 96, resetAt: future2, windows: [ - { window: 'weekly', remainingPercent: 99, resetAt: future2 }, { window: '5h', remainingPercent: 96, resetAt: future }, + { window: 'weekly', remainingPercent: 99, resetAt: future2 }, ], }, }, @@ -347,6 +363,9 @@ describe('windows rework — reviewer frames', () => { expect(frame).toContain('Gm 5h') expect(frame).toContain('NG 7d') expect(frame).toContain('NG 5h') + // Same order assertion as (a): shortest window first. + expect(frame.indexOf('Gm 5h')).toBeLessThan(frame.indexOf('Gm 7d')) + expect(frame.indexOf('NG 5h')).toBeLessThan(frame.indexOf('NG 7d')) testSetup.renderer.destroy() }) }) diff --git a/packages/opencode/src/tui.test.tsx b/packages/opencode/src/tui.test.tsx index 67c4d30..455a702 100644 --- a/packages/opencode/src/tui.test.tsx +++ b/packages/opencode/src/tui.test.tsx @@ -80,6 +80,11 @@ function makeFixture(): Fixture { const statePath = join(root, 'sidebar-state.json') const logPath = join(root, 'tui.log') const prefsPath = join(root, 'tui-preferences.jsonc') + // Save preload-pinned values to restore on cleanup instead of deleting — + // a delete drops resolution to the operator's real state/config dirs. + const savedSidebar = process.env[SIDEBAR_STATE_ENV] + const savedTuiLog = process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE + const savedPrefs = process.env[TUI_PREFS_FILE_ENV] process.env[SIDEBAR_STATE_ENV] = statePath process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE = logPath process.env[TUI_PREFS_FILE_ENV] = prefsPath @@ -88,9 +93,14 @@ function makeFixture(): Fixture { logPath, prefsPath, cleanup: () => { - delete process.env[SIDEBAR_STATE_ENV] - delete process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE - delete process.env[TUI_PREFS_FILE_ENV] + if (savedSidebar !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebar + else delete process.env[SIDEBAR_STATE_ENV] + if (savedTuiLog !== undefined) + process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE = savedTuiLog + else delete process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE + if (savedPrefs !== undefined) process.env[TUI_PREFS_FILE_ENV] = savedPrefs + else delete process.env[TUI_PREFS_FILE_ENV] rmSync(root, { recursive: true, force: true }) }, } diff --git a/packages/opencode/src/tui.tsx b/packages/opencode/src/tui.tsx index b8e0deb..3fb397b 100644 --- a/packages/opencode/src/tui.tsx +++ b/packages/opencode/src/tui.tsx @@ -858,8 +858,19 @@ export function SidebarPanel(props: SidebarPanelProps): JSX.Element { visibleAccounts()[0] const used = () => { const q = account()?.quota - const entry = q?.gemini ?? q?.['non-gemini'] - return entry ? 100 - clamp(entry.remainingPercent, 0, 100) : null + // Derive tone from the WORST (most-constrained) pool so a + // critical non-gemini pool doesn't hide behind a healthy + // gemini tone. Worst = highest used-percent across available pools. + const geminiUsed = + q?.gemini != null + ? 100 - clamp(q.gemini.remainingPercent, 0, 100) + : null + const nonGeminiUsed = + q?.['non-gemini'] != null + ? 100 - clamp(q['non-gemini'].remainingPercent, 0, 100) + : null + if (geminiUsed === null && nonGeminiUsed === null) return null + return Math.max(geminiUsed ?? 0, nonGeminiUsed ?? 0) } const poolText = (key: 'gemini' | 'non-gemini') => { const entry = account()?.quota[key] diff --git a/test/environment.test.ts b/test/environment.test.ts index ddd32ec..def61ec 100644 --- a/test/environment.test.ts +++ b/test/environment.test.ts @@ -1,6 +1,14 @@ import { describe, expect, it } from 'bun:test' +import { homedir } from 'node:os' +import { join } from 'node:path' +import { getUserConfigPath } from '../packages/opencode/src/plugin/config/loader.ts' +import { getOpencodeConfigDir } from '../packages/opencode/src/plugin/config/updater.ts' +import { defaultFilesystemRoots } from '../packages/opencode/src/plugin/dependencies.ts' import { getConfigDir as getStorageConfigDir } from '../packages/opencode/src/plugin/storage.ts' +import { getRpcDir } from '../packages/opencode/src/rpc/rpc-dir.ts' +import { getSidebarStateFile } from '../packages/opencode/src/sidebar-state.ts' +import { resolveTuiLogPath } from '../packages/opencode/src/tui/file-logger.ts' import { getPiAntigravityAuthFile, getPiConfigDir, @@ -63,3 +71,227 @@ describe('test environment isolation', () => { expect(recovery.PART_STORAGE.startsWith(root)).toBe(true) }) }) + +/** + * Guard: every writable-path resolver must resolve inside the test root. + * + * Two assertions per resolver: + * (a) under-test: the resolved path IS inside the test root + * (b) production default: the resolved path is NOT inside the test root + * (i.e. it would write to the real user dir in production, proving + * the test-env override is actually doing work — a resolver + * accidentally hardcoded to a temp path would pass (a) but fail (b)) + * + * If any assertion fails the suite is silently writing to the operator's + * live state. The error message names the offending resolver and the + * path it resolved to so the root cause is immediately visible. + */ +describe('state leak guard — every path resolver must stay inside the test root', () => { + const root = process.env.ANTIGRAVITY_TEST_ROOT + if (!root) throw new Error('ANTIGRAVITY_TEST_ROOT not set by preload') + + // The operator's real home — used to assert production defaults differ. + // homedir() reads HOME (pinned by preload), so the production default is + // computed by stripping the test indirection. + const realHome = homedir() + // Production xdg-state default that the real user sees outside of tests. + const productionXdgState = join(realHome, '.local', 'state') + + // ── Sidebar state ───────────────────────────────────────────────────────── + + it('getSidebarStateFile() resolves under the test root', () => { + const resolved = getSidebarStateFile() + expect( + resolved.startsWith(root), + `getSidebarStateFile() resolved OUTSIDE test root: ${resolved}\n` + + "The suite would corrupt the operator's live sidebar state file. " + + 'Check that ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE is pinned in test/setup.ts.', + ).toBe(true) + }) + + it('getSidebarStateFile() production default is NOT under the test root', () => { + // Verifies the test-env override is doing real work: without it, the + // resolver would fall back to this production path. + const productionDefault = join( + productionXdgState, + 'cortexkit', + 'antigravity-auth', + 'sidebar-state.json', + ) + expect(productionDefault.startsWith(root)).toBe(false) + }) + + // ── RPC dir ─────────────────────────────────────────────────────────────── + + it('getRpcDir() resolves under the test root', () => { + // Use an arbitrary project directory; what matters is that the root dir + // is isolated, not the per-project hash suffix. + const resolved = getRpcDir('/some/project') + expect( + resolved.startsWith(root), + `getRpcDir() resolved OUTSIDE test root: ${resolved}\n` + + "The suite would write port files into the operator's real state dir " + + '(2119 leaked project dirs observed). Check that ANTIGRAVITY_AUTH_RPC_DIR ' + + 'is pinned in test/setup.ts.', + ).toBe(true) + }) + + it('getRpcDir() production default is NOT under the test root', () => { + const productionDefault = join( + productionXdgState, + 'cortexkit', + 'antigravity-auth', + 'rpc', + ) + expect(productionDefault.startsWith(root)).toBe(false) + }) + + // ── defaultFilesystemRoots() — same XDG_STATE_HOME gap ─────────────────── + + it('defaultFilesystemRoots() sidebarStateRoot resolves under the test root', () => { + const { sidebarStateRoot } = defaultFilesystemRoots() + expect( + sidebarStateRoot.startsWith(root), + `defaultFilesystemRoots().sidebarStateRoot resolved OUTSIDE test root: ${sidebarStateRoot}`, + ).toBe(true) + }) + + it('defaultFilesystemRoots() rpcRoot resolves under the test root', () => { + const { rpcRoot } = defaultFilesystemRoots() + expect( + rpcRoot.startsWith(root), + `defaultFilesystemRoots().rpcRoot resolved OUTSIDE test root: ${rpcRoot}`, + ).toBe(true) + }) + + it('defaultFilesystemRoots() production xdg-state defaults are NOT under the test root', () => { + const productionSidebar = join( + productionXdgState, + 'cortexkit', + 'antigravity-auth', + ) + const productionRpc = join( + productionXdgState, + 'cortexkit', + 'antigravity-auth', + 'rpc', + ) + expect(productionSidebar.startsWith(root)).toBe(false) + expect(productionRpc.startsWith(root)).toBe(false) + }) + + // ── TUI log ─────────────────────────────────────────────────────────────── + + it('resolveTuiLogPath() resolves under the test root', () => { + const resolved = resolveTuiLogPath() + expect( + resolved.startsWith(root), + `resolveTuiLogPath() resolved OUTSIDE test root: ${resolved}\n` + + 'Check that ANTIGRAVITY_AUTH_TUI_LOG_FILE or XDG_STATE_HOME is pinned.', + ).toBe(true) + }) + + it('resolveTuiLogPath() production default is NOT under the test root', () => { + const productionDefault = join( + productionXdgState, + 'cortexkit', + 'antigravity-auth', + 'tui.log', + ) + expect(productionDefault.startsWith(root)).toBe(false) + }) + + // ── Plugin config (getUserConfigPath, getOpencodeConfigDir) ─────────────── + + it('getUserConfigPath() resolves under the test root', () => { + const resolved = getUserConfigPath() + expect( + resolved.startsWith(root), + `getUserConfigPath() resolved OUTSIDE test root: ${resolved}\n` + + 'Check that OPENCODE_CONFIG_DIR is pinned in test/setup.ts.', + ).toBe(true) + }) + + it('getOpencodeConfigDir() resolves under the test root', () => { + // getOpencodeConfigDir() uses XDG_CONFIG_HOME (pinned), not OPENCODE_CONFIG_DIR. + const resolved = getOpencodeConfigDir() + expect( + resolved.startsWith(root), + `getOpencodeConfigDir() resolved OUTSIDE test root: ${resolved}`, + ).toBe(true) + }) + + it('config resolver production defaults are NOT under the test root', () => { + const productionConfig = join(realHome, '.config', 'opencode') + expect(productionConfig.startsWith(root)).toBe(false) + }) + + // ── Storage config dir ──────────────────────────────────────────────────── + + it('storage getConfigDir() resolves under the test root', () => { + const resolved = getStorageConfigDir() + expect( + resolved.startsWith(root), + `storage getConfigDir() resolved OUTSIDE test root: ${resolved}`, + ).toBe(true) + }) + + // ── auto-update-checker constants (module-level captures, separate class) ──── + // + // auto-update-checker/constants.ts uses os.homedir() directly (not HOME + // env) for CACHE_DIR. Bun's os.homedir() is an OS-level call that does NOT + // read process.env.HOME — it returns the passwd entry. This means CACHE_DIR + // resolves to the real ~/.cache/opencode regardless of env pinning. + // + // This is a named finding, not a failure to fix: the test suite does NOT + // write to CACHE_DIR. The write paths (installPackage, invalidatePackage) + // only fire from the plugin's session.created event, which unit tests never + // trigger. The only test-time read is an existsSync check in checker.ts:149 + // which is a no-op when the directory is absent. CACHE_DIR is therefore + // UNREACHABLE FOR WRITES in the unit test suite. + // + // USER_OPENCODE_CONFIG uses XDG_CONFIG_HOME (pinned) and resolves correctly. + // It is asserted here separately because its failure mode (import-order + // sensitivity) differs from the XDG_STATE_HOME gap above. + + it('auto-update-checker USER_OPENCODE_CONFIG resolves under the test root', async () => { + const { USER_OPENCODE_CONFIG } = await import( + `../packages/opencode/src/hooks/auto-update-checker/constants.ts?bust=${Date.now()}` + ) + expect( + (USER_OPENCODE_CONFIG as string).startsWith(root), + `auto-update-checker USER_OPENCODE_CONFIG resolved OUTSIDE test root: ${USER_OPENCODE_CONFIG as string}`, + ).toBe(true) + }) + + it('auto-update-checker USER_OPENCODE_CONFIG production default is NOT under the test root', () => { + const productionUserConfig = join( + realHome, + '.config', + 'opencode', + 'opencode.json', + ) + expect(productionUserConfig.startsWith(root)).toBe(false) + }) + + it('auto-update-checker CACHE_DIR resolves OUTSIDE the test root (write-unreachable)', async () => { + // CACHE_DIR uses os.homedir() (OS-level, not env-based) and resolves + // to the real ~/.cache/opencode even under test — env pinning cannot + // redirect it. This is the (b)-half guard: CACHE_DIR must be outside + // the test root. If it ever drifts inside, env isolation has broken + // and the situation must be re-evaluated. + const { CACHE_DIR } = await import( + `../packages/opencode/src/hooks/auto-update-checker/constants.ts?bust=${Date.now()}` + ) + expect( + (CACHE_DIR as string).startsWith(root), + `CACHE_DIR resolved INSIDE test root: ${CACHE_DIR as string}\n` + + 'os.homedir() now returns the test home — the isolation rationale no longer holds. ' + + 'Audit write paths before the suite is considered safe.', + ).toBe(false) + // The real protection: none of the writing callsites are reachable from + // bun test. cache.ts::installPackage and invalidatePackage are only called + // from runBackgroundUpdateCheck which fires on session.created events that + // unit tests never emit. Verified by grep: no test imports cache.ts. + }) +}) diff --git a/test/setup.ts b/test/setup.ts index a39fc81..8dbe06c 100644 --- a/test/setup.ts +++ b/test/setup.ts @@ -8,9 +8,10 @@ const home = join(root, 'home') const config = join(root, 'config') const cache = join(root, 'cache') const data = join(root, 'data') +const state = join(root, 'state') const pi = join(root, 'pi-agent') -for (const path of [home, config, cache, data, pi]) +for (const path of [home, config, cache, data, state, pi]) mkdirSync(path, { recursive: true }) process.env.ANTIGRAVITY_TEST_ROOT = root @@ -19,12 +20,32 @@ process.env.USERPROFILE = home process.env.XDG_CONFIG_HOME = config process.env.XDG_CACHE_HOME = cache process.env.XDG_DATA_HOME = data +// XDG_STATE_HOME must be pinned before sidebar-state.ts and rpc-dir.ts import, +// because xdg-basedir reads state from env once at module load, and both +// resolvers fall back to when their env +// override is unset. Without this pin, any test that deletes the per-test +// override drops resolution to the operator's real ~/.local/state dir. +process.env.XDG_STATE_HOME = state process.env.APPDATA = config process.env.LOCALAPPDATA = cache process.env.OPENCODE_CONFIG_DIR = join(config, 'opencode') process.env.PI_AGENT_DIR = pi process.env.PI_ANTIGRAVITY_AUTH_FILE = join(pi, 'antigravity-accounts.json') process.env.OPENCODE_ANTIGRAVITY_GEMINI_DUMP_DIR = join(root, 'gemini-dumps') +// Pin the two state-dir overrides so tests that reset them restore to an +// in-root path rather than dropping to the real default. +process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = join( + state, + 'cortexkit', + 'antigravity-auth', + 'sidebar-state.json', +) +process.env.ANTIGRAVITY_AUTH_RPC_DIR = join( + state, + 'cortexkit', + 'antigravity-auth', + 'rpc', +) const stubbedGlobals = new Map() const dateNowSpyState: { active: boolean } = { active: false }