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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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,