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/__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/account-manager.test.ts b/packages/core/src/account-manager.test.ts index 093bafa..3146ed9 100644 --- a/packages/core/src/account-manager.test.ts +++ b/packages/core/src/account-manager.test.ts @@ -54,6 +54,49 @@ describe('core AccountManager', () => { ).toEqual(['r1', 'r2', 'r3']) }) + it('normalizes persisted legacy quota keys for soft-quota and proactive-rotation reads', () => { + const now = 1_700_000_000_000 + const legacy: AccountStorageV4 = { + version: 4, + accounts: [ + { + refreshToken: 'legacy-token', + addedAt: 1, + lastUsed: 0, + cachedQuota: { + claude: { remainingFraction: 0.4, modelCount: 1 }, + }, + cachedQuotaUpdatedAt: now, + }, + { + refreshToken: 'other-token', + addedAt: 1, + lastUsed: 0, + }, + ], + activeIndex: 0, + } + const memory = createStore(legacy) + const manager = new AccountManager(undefined, legacy, { + store: memory.store, + now: () => now, + }) + const account = manager.getAccounts()[0]! + + expect( + manager.isAccountOverSoftQuota( + account, + 'claude', + 50, + 60_000, + 'claude-sonnet', + ), + ).toBe(true) + expect( + manager.shouldProactivelyRotate('claude', 'claude-sonnet', 50, 60_000), + ).toBe(true) + }) + it.each([ 'sticky', 'round-robin', @@ -69,6 +112,83 @@ describe('core AccountManager', () => { ).not.toBeNull() }) + it('hybrid skips the active Gemini account limited on the antigravity header style', () => { + const now = 1_700_000_000_000 + const hybridStored: AccountStorageV4 = { + version: 4, + accounts: [ + { refreshToken: 'r1', projectId: 'p1', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r2', projectId: 'p2', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r3', projectId: 'p3', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r4', projectId: 'p4', addedAt: 1, lastUsed: 0 }, + ], + activeIndex: 1, + activeIndexByFamily: { gemini: 1 }, + } + const memory = createStore(hybridStored) + const manager = new AccountManager(undefined, hybridStored, { + store: memory.store, + now: () => now, + random: () => 0.5, + }) + const limited = manager.getAccounts()[1]! + manager.markRateLimitedWithReason( + limited, + 'gemini', + 'antigravity', + 'antigravity-gemini-3.6-flash', + 'RATE_LIMIT_EXCEEDED', + ) + + const selected = manager.getCurrentOrNextForFamily( + 'gemini', + 'antigravity-gemini-3.6-flash', + 'hybrid', + 'antigravity', + ) + + expect(selected?.index).toBe(0) + }) + + it('hybrid returns null when every Gemini account is limited on the antigravity header style', () => { + const now = 1_700_000_000_000 + const hybridStored: AccountStorageV4 = { + version: 4, + accounts: [ + { refreshToken: 'r1', projectId: 'p1', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r2', projectId: 'p2', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r3', projectId: 'p3', addedAt: 1, lastUsed: 0 }, + { refreshToken: 'r4', projectId: 'p4', addedAt: 1, lastUsed: 0 }, + ], + activeIndex: 1, + activeIndexByFamily: { gemini: 1 }, + } + const memory = createStore(hybridStored) + const manager = new AccountManager(undefined, hybridStored, { + store: memory.store, + now: () => now, + random: () => 0.5, + }) + for (const account of manager.getAccounts()) { + manager.markRateLimitedWithReason( + account, + 'gemini', + 'antigravity', + 'antigravity-gemini-3.6-flash', + 'RATE_LIMIT_EXCEEDED', + ) + } + + const selected = manager.getCurrentOrNextForFamily( + 'gemini', + 'antigravity-gemini-3.6-flash', + 'hybrid', + 'antigravity', + ) + + expect(selected).toBeNull() + }) + it('tracks model-specific limits independently', () => { let now = 1_000 const memory = createStore(stored) @@ -143,6 +263,172 @@ 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('persists and restores the captured tier schema marker across save→loadFromDisk', async () => { + const seeded = { + version: 4, + accounts: [ + { + refreshToken: 'r1', + projectId: 'p1', + addedAt: 1, + lastUsed: 0, + capturedTierId: 'free-tier', + capturedTierAt: 1_700_000_000_000, + capturedTierSchemaVersion: 1, + }, + ], + activeIndex: 0, + } as AccountStorageV4 & { + accounts: Array<{ capturedTierSchemaVersion?: number }> + } + const memory = createStore(seeded) + const manager = new AccountManager(undefined, seeded, { + store: memory.store, + }) + + await manager.saveToDiskReplace() + + expect(memory.state()?.accounts[0]).toMatchObject({ + capturedTierSchemaVersion: 1, + }) + const reloaded = new AccountManager( + undefined, + memory.state() ?? undefined, + { store: memory.store }, + ) + expect( + (reloaded.getAccounts()[0] as { capturedTierSchemaVersion?: number }) + ?.capturedTierSchemaVersion, + ).toBe(1) + }) + + it('drops the quota write when the refresh token captured at refresh time is gone (remove-during-refresh race)', () => { + // Race: an async quota refresh is in flight for account A while the + // user removes account A from the pool. When the refresh resolves, + // 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 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() + 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, { @@ -175,3 +461,70 @@ describe('AccountManager instance dependencies', () => { expect(second.getAccounts()[0]?.coolingDownUntil).toBe(9_500) }) }) + +describe('managedProjectId projection', () => { + it('getAccountsForQuotaCheck falls back to record managedProjectId when parts lack it', () => { + 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, + }) + // 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') + 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 — dispose clears the debounce and forces it immediately. + manager.requestSaveToDisk() + await manager.dispose() + 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 2ad6fa5..9be86b4 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,14 +19,18 @@ import { MAX_FINGERPRINT_HISTORY, updateFingerprintVersion, } from './fingerprint.ts' -import type { QuotaGroup, QuotaGroupSummary } from './quota-types.ts' +import { getQuotaGroupForModel } from './model-registry.ts' +import { + normalizeLegacyCachedQuota, + type QuotaGroup, + type QuotaGroupSummary, +} from './quota-types.ts' import { type AccountWithMetrics, getHealthTracker, getTokenTracker, selectHybridAccount, } from './rotation.ts' -import { getModelFamily } from './transform/model-resolver.ts' export type { AccountSelectionStrategy, @@ -72,6 +77,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 @@ -89,7 +101,20 @@ 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 + /** + * Captured plan tier ID from the most recent `loadCodeAssist` response. + * Raw upstream string (e.g. `"free-tier"`) — never normalised. + */ + capturedTierId?: string + /** Raw paid-tier ID from the most recent `loadCodeAssist` response. */ + capturedPaidTierId?: string + /** Epoch ms when `capturedTierId` was last recorded. */ + capturedTierAt?: number + /** Schema version of the most recent tier capture, even when paid tier is absent. */ + capturedTierSchemaVersion?: number verificationRequired?: boolean verificationRequiredAt?: number verificationRequiredReason?: string @@ -113,6 +138,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 +245,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 +260,20 @@ export function resolveQuotaGroup( model?: string | null, ): QuotaGroup { if (model) { - return getModelFamily(model) - } - return family === 'claude' ? 'claude' : 'gemini-pro' + const registryGroup = getQuotaGroupForModel(model) + if (registryGroup) return registryGroup + const lower = model.toLowerCase() + // 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' } function isOverSoftQuotaThreshold( @@ -368,6 +416,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, @@ -378,10 +430,16 @@ export class AccountManager { touchedForQuota: {}, fingerprint: acc.fingerprint ?? generateFingerprint(), fingerprintHistory: acc.fingerprintHistory ?? [], - cachedQuota: acc.cachedQuota as - | Partial> - | undefined, + cachedQuota: normalizeLegacyCachedQuota(acc.cachedQuota), + // Restore the opaque identity stamp alongside the quota so the + // post-load projection can detect a stale snapshot captured + // for a different account after an index shift. + cachedQuotaAccountId: acc.cachedQuotaAccountId, cachedQuotaUpdatedAt: acc.cachedQuotaUpdatedAt, + capturedTierId: acc.capturedTierId, + capturedPaidTierId: acc.capturedPaidTierId, + capturedTierAt: acc.capturedTierAt, + capturedTierSchemaVersion: acc.capturedTierSchemaVersion, dailyRequestCounts: acc.dailyRequestCounts, verificationRequired: acc.verificationRequired, verificationRequiredAt: acc.verificationRequiredAt, @@ -656,6 +714,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', @@ -774,7 +847,13 @@ export class AccountManager { lastUsed: acc.lastUsed, healthScore: healthTracker.getScore(acc.index), isRateLimited: - isRateLimitedForFamily(acc, family, this.now, model) || + isRateLimitedForHeaderStyle( + acc, + family, + headerStyle, + this.now, + model, + ) || isOverSoftQuotaThreshold( acc, family, @@ -1518,6 +1597,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 } @@ -1598,8 +1681,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, @@ -1615,7 +1698,15 @@ 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, + capturedTierId: a.capturedTierId, + capturedPaidTierId: a.capturedPaidTierId, + capturedTierAt: a.capturedTierAt, + capturedTierSchemaVersion: a.capturedTierSchemaVersion, dailyRequestCounts: a.dailyRequestCounts, verificationRequired: a.verificationRequired, verificationRequiredAt: a.verificationRequiredAt, @@ -1822,11 +1913,71 @@ 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() + } + + /** + * Apply a subset of fields from a quota-fetch `updatedAccount` result onto + * the live in-memory record for the given index. Only patches fields that + * are present and non-empty in `patch` to avoid overwriting valid state + * with stale or missing values. + * + * Identity guard: if `expectedRefreshToken` is provided and the account at + * `accountIndex` no longer carries that token (concurrent reorder/replace), + * the patch is silently dropped. + * + * `managedProjectId` is intentionally absent from the patch type: the only + * caller (`BackgroundQuotaRefresh`) routes through `PollerAccountView` which + * exposes only `capturedTierId`/`capturedTierAt`; project-context updates + * happen via `ensureProjectContext`, not through this method. + */ + applyUpdatedAccount( + accountIndex: number, + patch: Partial< + Pick< + AccountMetadataV3, + | 'capturedTierId' + | 'capturedPaidTierId' + | 'capturedTierAt' + | 'capturedTierSchemaVersion' + > + >, + expectedRefreshToken?: string, + ): void { + const account = this.accounts[accountIndex] + if ( + !account || + (expectedRefreshToken !== undefined && + account.parts.refreshToken !== expectedRefreshToken) + ) + return + if (patch.capturedTierId !== undefined) { + account.capturedTierId = patch.capturedTierId + } + if (patch.capturedPaidTierId !== undefined) { + account.capturedPaidTierId = patch.capturedPaidTierId + } + if (patch.capturedTierAt !== undefined) { + account.capturedTierAt = patch.capturedTierAt + } + if (patch.capturedTierSchemaVersion !== undefined) { + account.capturedTierSchemaVersion = patch.capturedTierSchemaVersion } } @@ -2000,8 +2151,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, diff --git a/packages/core/src/account-types.ts b/packages/core/src/account-types.ts index 3c7048b..47ae022 100644 --- a/packages/core/src/account-types.ts +++ b/packages/core/src/account-types.ts @@ -64,10 +64,22 @@ 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, - { 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?: { @@ -78,6 +90,17 @@ export interface AccountMetadataV3 { resetTime?: string }[] cachedQuotaUpdatedAt?: number + /** + * Captured plan tier ID from the most recent `loadCodeAssist` response. + * Raw upstream string (e.g. `"free-tier"`) — never normalised. + */ + capturedTierId?: string + /** Raw paid-tier ID from the most recent `loadCodeAssist` response. */ + capturedPaidTierId?: string + /** Epoch ms when `capturedTierId` was last recorded. */ + capturedTierAt?: number + /** Schema version of the most recent tier capture, even when paid tier is absent. */ + capturedTierSchemaVersion?: number /** Daily request counts per model family, resets when date changes */ dailyRequestCounts?: { date: string diff --git a/packages/core/src/auth-types.ts b/packages/core/src/auth-types.ts index 10b2b90..a60e40f 100644 --- a/packages/core/src/auth-types.ts +++ b/packages/core/src/auth-types.ts @@ -33,4 +33,10 @@ export interface RefreshParts { export interface ProjectContextResult { auth: OAuthAuthDetails effectiveProjectId: string + /** + * Plan tier captured from the `loadCodeAssist` payload at project-context + * resolution time. Present only when the upstream returned a non-empty + * `currentTier.id`; absent when the payload lacked tier info. + */ + capturedTier?: { id: string; paidId?: string; capturedAt: number } } diff --git a/packages/core/src/model-registry.ts b/packages/core/src/model-registry.ts index b712fbb..029d20f 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,18 @@ 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] ?? + // 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/project.test.ts b/packages/core/src/project.test.ts index 9f1224c..bd6fdc1 100644 --- a/packages/core/src/project.test.ts +++ b/packages/core/src/project.test.ts @@ -116,6 +116,55 @@ describe('project bootstrap', () => { expect(fetchSpy).toHaveBeenCalledTimes(1) }) + it('Fix2-tier-capture: capturedTier is returned from ensureProjectContext when loadCodeAssist returns currentTier', async () => { + const capturedNow = 1_785_000_000_000 + spyOn(Date, 'now').mockImplementation(() => capturedNow) + const fetchSpy = mock().mockResolvedValue( + mockResponse({ + cloudaicompanionProject: { id: 'my-project' }, + currentTier: { id: 'pro-tier' }, + paidTier: { id: 'g1-pro-tier' }, + }), + ) + const { fetchWithAgyCliTransport } = await import('./agy-transport.ts') + ;(fetchWithAgyCliTransport as any).mockImplementation(fetchSpy) + + const auth = { + type: 'oauth' as const, + access: 'access-token', + refresh: 'bare-token', + expires: Date.now() + 60_000, + } + + const result = await ensureProjectContext(auth) + + expect(result.capturedTier).toEqual({ + id: 'pro-tier', + paidId: 'g1-pro-tier', + capturedAt: capturedNow, + }) + }) + + it('Fix2-tier-absent: capturedTier is absent when loadCodeAssist returns no currentTier', async () => { + const fetchSpy = mock().mockResolvedValue( + mockResponse({ cloudaicompanionProject: { id: 'my-project' } }), + ) + const { fetchWithAgyCliTransport } = await import('./agy-transport.ts') + ;(fetchWithAgyCliTransport as any).mockImplementation(fetchSpy) + + const auth = { + type: 'oauth' as const, + access: 'access-token', + refresh: 'bare-token-2', + expires: Date.now() + 60_000, + } + + const result = await ensureProjectContext(auth) + + // No tier in payload — must be absent, not defaulted. + expect(result.capturedTier).toBeUndefined() + }) + it('does not retry managed-project provisioning after a cached failure expires', async () => { let now = 1_000 spyOn(Date, 'now').mockImplementation(() => now) diff --git a/packages/core/src/project.ts b/packages/core/src/project.ts index f2645db..02a03be 100644 --- a/packages/core/src/project.ts +++ b/packages/core/src/project.ts @@ -40,6 +40,11 @@ interface LoadCodeAssistPayload { currentTier?: { id?: string } + paidTier?: + | string + | { + id?: string + } allowedTiers?: AntigravityUserTier[] } @@ -295,6 +300,7 @@ export async function ensureProjectContext( const persistManagedProject = async ( managedProjectId: string, + capturedTier?: ProjectContextResult['capturedTier'], ): Promise => { const updatedAuth: OAuthAuthDetails = { ...auth, @@ -305,7 +311,11 @@ export async function ensureProjectContext( }), } - return { auth: updatedAuth, effectiveProjectId: managedProjectId } + return { + auth: updatedAuth, + effectiveProjectId: managedProjectId, + capturedTier, + } } // Try to resolve a managed project from Antigravity if possible. @@ -313,10 +323,24 @@ export async function ensureProjectContext( accessToken, parts.projectId ?? fallbackProjectId, ) + // Capture tier from the loadCodeAssist payload. The raw id is stored + // as-is; absent payload or missing tier leaves capturedTier undefined. + const capturedTierId = loadPayload?.currentTier?.id + const paidTierId = + typeof loadPayload?.paidTier === 'string' + ? loadPayload.paidTier + : loadPayload?.paidTier?.id + const tierFromPayload: ProjectContextResult['capturedTier'] = capturedTierId + ? { + id: capturedTierId, + ...(paidTierId ? { paidId: paidTierId } : {}), + capturedAt: Date.now(), + } + : undefined const resolvedManagedProjectId = extractManagedProjectId(loadPayload) if (resolvedManagedProjectId) { - return persistManagedProject(resolvedManagedProjectId) + return persistManagedProject(resolvedManagedProjectId, tierFromPayload) } // No managed project found - try to auto-provision one via onboarding. @@ -337,7 +361,7 @@ export async function ensureProjectContext( log.debug('Successfully provisioned managed project', { provisionedProjectId, }) - return persistManagedProject(provisionedProjectId) + return persistManagedProject(provisionedProjectId, tierFromPayload) } log.warn( @@ -352,11 +376,19 @@ export async function ensureProjectContext( } if (parts.projectId) { - return { auth, effectiveProjectId: parts.projectId } + return { + auth, + effectiveProjectId: parts.projectId, + capturedTier: tierFromPayload, + } } // No project id present in auth; fall back to the hardcoded id for requests. - return { auth, effectiveProjectId: fallbackProjectId } + return { + auth, + effectiveProjectId: fallbackProjectId, + capturedTier: tierFromPayload, + } } if (!cacheKey) { diff --git a/packages/core/src/quota-manager.test.ts b/packages/core/src/quota-manager.test.ts index c3287f6..04efe4f 100644 --- a/packages/core/src/quota-manager.test.ts +++ b/packages/core/src/quota-manager.test.ts @@ -1,6 +1,15 @@ 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 RetrieveUserQuotaSummaryBucket, + type RetrieveUserQuotaSummaryResponse, +} from './quota-manager.ts' import type { AccountQuotaResult } from './quota-types.ts' function makeAccount( @@ -44,7 +53,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 +81,69 @@ 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 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-flash') - expect(classifyQuotaGroup('gemini-3.1-pro', 'Gemini 3.1 Pro')).toBe( - 'gemini-pro', - ) + // 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', () => { + const { classifyQuotaGroup } = createQuotaManager({ + fetchAccountQuota: makeHarness().fetch, + keyOf: keyOfAccount, + }) + + expect(classifyQuotaGroup('tab_12345', 'Autocomplete')).toBe('gemini') }) - it('classifies gpt-oss variants', () => { + 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', + ) + }) + + 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', ) }) @@ -140,9 +182,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 +206,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) }) }) @@ -743,3 +785,490 @@ 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) + // shortest-first: 5h before weekly + expect(gemini.windows![0]!.window).toBe('5h') + expect(gemini.windows![0]!.remainingFraction).toBeCloseTo(0.9886, 3) + expect(gemini.windows![1]!.window).toBe('weekly') + expect(gemini.windows![1]!.remainingFraction).toBeCloseTo(0.9214, 3) + + const nonGemini = summary.groups['non-gemini']! + expect(nonGemini.windows).toHaveLength(2) + expect(nonGemini.windows![0]!.window).toBe('5h') + expect(nonGemini.windows![0]!.remainingFraction).toBeCloseTo(0.9556, 3) + expect(nonGemini.windows![1]!.window).toBe('weekly') + expect(nonGemini.windows![1]!.remainingFraction).toBeCloseTo(0.9852, 3) + }) + + it('maps Free (weekly-only) groups to pools', () => { + 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: shortest first', () => { + const response = loadFixture('pro-ruqs.json') + const summary = aggregateQuotaSummary(response) + + for (const group of Object.values(summary.groups)) { + if (!group?.windows) continue + expect(group.windows).toHaveLength(2) + expect(group.windows![0]!.window).toBe('5h') + expect(group.windows![1]!.window).toBe('weekly') + } + }) + + it('sorts unknown window kinds last, deterministically', () => { + const buckets: RetrieveUserQuotaSummaryBucket[] = [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly', + window: 'weekly', + resetTime: '2026-07-28T00:00:00Z', + remainingFraction: 0.9, + }, + { + bucketId: 'gemini-daily', + displayName: 'Daily', + window: 'daily', + resetTime: '2026-07-25T00:00:00Z', + remainingFraction: 0.5, + }, + { + bucketId: 'gemini-5h', + displayName: '5h', + window: '5h', + resetTime: '2026-07-24T12:00:00Z', + remainingFraction: 0.8, + }, + ] + const summary = aggregateQuotaSummary({ + groups: [ + { displayName: 'Gemini', buckets, description: 'Models: A, B, C' }, + ], + }) + const windows = summary.groups.gemini!.windows! + expect(windows).toHaveLength(3) + // Known windows sorted shortest-first: 5h → weekly + expect(windows[0]!.window).toBe('5h') + expect(windows[1]!.window).toBe('weekly') + // Unknown window ('daily') sorts last + expect(windows[2]!.window).toBe('daily') + }) + + it('counts models from the description minus the prefix label', () => { + const response: RetrieveUserQuotaSummaryResponse = { + groups: [ + { + 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. 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') + }) + + 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) + }) +}) + +// ============================================================================ +// 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') + }) + + 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 ( + _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-only', + endpoints: ENDPOINTS, + fetchVia: fetchVia as any, + }) + expect(receivedProjectId).toBe('managed-only') + }) + + 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) + }) + + it('N1: a transient 500 on the managed attempt does NOT fall through to the projectId fallback', async () => { + // N1 regression: before the fix, tryBody returned null for ANY failure + // (403, 500, network). The fallback was entered on all of them, which + // could return a different project\'s quota data on transient errors. + // Now the fallback is gated on 403 only. + const triedProjects: string[] = [] + const fetchVia = async ( + _url: string, + init: RequestInit, + ): Promise => { + const body = JSON.parse((init as any).body ?? '{}') + triedProjects.push(body.project as string) + // Return 500 for the managed project attempt (transient error). + return new Response('internal error', { status: 500 }) + } + // Both managed and regular project IDs are provided. The 500 on the + // managed attempt must NOT trigger the projectId fallback. + await expect( + fetchQuotaSummary({ + accessToken: 'tok', + managedProjectId: 'managed-proj', + projectId: 'regular-proj', + endpoints: ENDPOINTS, + fetchVia: fetchVia as any, + }), + ).rejects.toThrow() + // Only the managed project should have been tried — the regular-proj + // fallback must not run on a transient error. + expect(triedProjects).toHaveLength(1) + expect(triedProjects[0]).toBe('managed-proj') + }) + + it('N1: a 500 on the managed attempt does not clobber the 403-fallback path', async () => { + // Complement: a 403 on the managed project DOES enter the fallback. + // This ensures the gating change did not accidentally disable the + // existing 403-fallback behavior. + const triedProjects: string[] = [] + const summary: RetrieveUserQuotaSummaryResponse = { groups: [] } + const fetchVia = async ( + _url: string, + init: RequestInit, + ): Promise => { + const body = JSON.parse((init as any).body ?? '{}') + triedProjects.push(body.project as string) + if (body.project === 'managed-proj') { + return new Response('{}', { status: 403 }) + } + return new Response(JSON.stringify(summary), { status: 200 }) + } + const result = await fetchQuotaSummary({ + accessToken: 'tok', + managedProjectId: 'managed-proj', + projectId: 'regular-proj', + endpoints: ENDPOINTS, + fetchVia: fetchVia as any, + }) + // The 403 on managed-proj must trigger the fallback to regular-proj. + expect(triedProjects).toHaveLength(2) + expect(triedProjects[0]).toBe('managed-proj') + expect(triedProjects[1]).toBe('regular-proj') + expect(result.summary).toEqual(summary) + }) +}) diff --git a/packages/core/src/quota-manager.ts b/packages/core/src/quota-manager.ts index b88aafb..5d7f864 100644 --- a/packages/core/src/quota-manager.ts +++ b/packages/core/src/quota-manager.ts @@ -26,8 +26,8 @@ import type { QuotaGroup, QuotaGroupSummary, QuotaSummary, + QuotaWindowEntry, } from './quota-types.ts' -import { getModelFamily } from './transform/model-resolver.ts' const log = createLogger('quota-manager') @@ -418,16 +418,19 @@ export function classifyQuotaGroup( } const combined = `${modelName} ${displayName ?? ''}`.toLowerCase() - if (combined.includes('claude')) { - return 'claude' + // 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' } - const isGemini3 = - combined.includes('gemini-3') || combined.includes('gemini 3') - if (!isGemini3) { - return null + if (combined.includes('gemini')) { + return 'gemini' } - const family = getModelFamily(modelName) - return family === 'gemini-flash' ? 'gemini-flash' : 'gemini-pro' + return null } function normalizeRemainingFraction(value: unknown): number { @@ -548,6 +551,315 @@ 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 +} + +/** + * 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 +} + +/** + * Map a quota window kind to its duration in milliseconds for + * deterministic sort order. Unknown windows get MAX_SAFE_INTEGER so + * they sort last without random reordering across runs. + */ +function windowDurationMs(window: 'weekly' | '5h' | string): number { + switch (window) { + case '5h': + return 5 * 60 * 60 * 1000 + case 'weekly': + return 7 * 24 * 60 * 60 * 1000 + default: + return Number.MAX_SAFE_INTEGER + } +} + +/** + * Aggregate a retrieveUserQuotaSummary response into a QuotaSummary. + * + * Each RUQS group maps to a pool via bucketId prefix. Within a pool, + * windows are stored shortest-first (5h before weekly, etc.). 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: shortest window first so the binding 5h window leads visually. + windows.sort( + (a, b) => windowDurationMs(a.window) - windowDurationMs(b.window), + ) + + const constrained = mostConstrainedWindow(windows) + // Pick the first RECOGNIZED bucket for pool derivation. Older + // 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 + ? parseDescriptionModelCount(group.description) + : 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 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, +): Promise { + const timeoutMs = options.timeoutMs ?? QUOTA_MANAGER_DEFAULT_TIMEOUT_MS + const userAgent = options.userAgent ?? buildAntigravityHarnessUserAgent() + const transport = options.fetchVia ?? defaultTransport + const errors: string[] = [] + + if (options.endpoints.length === 0) { + throw new Error('No endpoints configured for fetchQuotaSummary') + } + + // Discriminated result so the fallback is only entered on a 403, not on + // transient errors (429, 5xx, network). A transient failure on the managed + // project ID must not fall through to the regular projectId — that could + // return a DIFFERENT project's quota data instead of signalling a retry. + type TryBodyResult = + | { ok: true; summary: RetrieveUserQuotaSummaryResponse } + | { ok: false; reason: '403' | 'transient' } + + const tryBody = async ( + endpoint: string, + projectId: string, + ): Promise => { + 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 { + ok: true, + summary: (await response.json()) as RetrieveUserQuotaSummaryResponse, + } + } + + const status = response.status + if (status === 403) { + errors.push( + `retrieveUserQuotaSummary 403 at ${endpoint} (project=${projectId.slice(0, 12)}…)`, + ) + return { ok: false, reason: '403' } + } + + if (status === 429 || status >= 500) { + const message = await response.text().catch(() => '') + 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 { ok: false, reason: 'transient' } + } + + const message = await response.text().catch(() => '') + errors.push( + `retrieveUserQuotaSummary ${status} at ${endpoint}${message ? `: ${message.trim().slice(0, 200)}` : ''}`, + ) + return { ok: false, reason: 'transient' } + } catch (error) { + errors.push( + `retrieveUserQuotaSummary network error at ${endpoint}: ${error instanceof Error ? error.message : String(error)}`, + ) + return { ok: false, reason: 'transient' } + } + } + + // 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 + let primaryGot403 = false + if (primary) { + for (const endpoint of options.endpoints) { + const result = await tryBody(endpoint, primary) + if (result.ok) return { summary: result.summary } + if (result.reason === '403') primaryGot403 = true + } + } + + // Retry with the regular projectId ONLY when the managed-project attempt + // got a 403 (i.e. the managed project does not own the user). A transient + // failure (429, 5xx, network) must NOT fall through here — the regular + // projectId could return a different project's quota data. + const fallbackId = + primaryGot403 && + options.managedProjectId && + options.projectId && + options.managedProjectId !== options.projectId + ? options.projectId + : undefined + if (fallbackId) { + for (const endpoint of options.endpoints) { + const result = await tryBody(endpoint, fallbackId) + if (result.ok) return { summary: result.summary } + } + } + + // 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. */ @@ -698,6 +1010,15 @@ export interface FetchGeminiCliQuotaOptions { endpoints: readonly string[] timeoutMs?: number userAgent?: string + /** + * Optional transport override. Production callers omit this; the e2e + * harness and tests inject a stub. Mirrors `fetchQuotaSummary`'s seam. + */ + fetchVia?: ( + url: string, + options: RequestInit, + extra: { timeoutMs: number }, + ) => Promise } export async function fetchGeminiCliQuota( @@ -705,11 +1026,13 @@ export async function fetchGeminiCliQuota( ): Promise { const timeoutMs = options.timeoutMs ?? QUOTA_MANAGER_DEFAULT_TIMEOUT_MS const userAgent = options.userAgent ?? buildAntigravityHarnessUserAgent() + const transport = options.fetchVia ?? defaultTransport + const errors: string[] = [] for (const endpoint of options.endpoints) { const body = options.projectId ? { project: options.projectId } : {} try { - const response = await defaultTransport( + const response = await transport( `${endpoint}/v1internal:retrieveUserQuota`, { method: 'POST', @@ -729,10 +1052,25 @@ export async function fetchGeminiCliQuota( const status = response.status if (status === 429 || status >= 500) { + errors.push(`fetchGeminiCliQuota ${status} at ${endpoint}`) continue } + // Non-retryable server response (e.g. 403) — treat as no CLI quota. return { buckets: [] } - } catch {} + } catch (error) { + // Transport-level failure (network abort, DNS, timeout). Collect the + // error and try the next endpoint; if all endpoints fail, throw so + // the caller can distinguish a transient network failure from an + // authoritative 'no CLI quota configured' response. + errors.push(error instanceof Error ? error.message : String(error)) + } + } + + // All endpoints produced transport-level errors — propagate so the + // outer caller (quota.ts .catch) can surface the real failure reason + // rather than the generic 'No Gemini CLI quota available' message. + if (errors.length > 0) { + throw new Error(errors.join('; ') || 'fetchGeminiCliQuota failed') } return { buckets: [] } diff --git a/packages/core/src/quota-types.ts b/packages/core/src/quota-types.ts index b9aca3d..0e7df08 100644 --- a/packages/core/src/quota-types.ts +++ b/packages/core/src/quota-types.ts @@ -9,12 +9,90 @@ import type { AccountMetadataV3 } from './account-types.ts' -export type QuotaGroup = 'claude' | 'gemini-pro' | 'gemini-flash' | 'gpt-oss' +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. + * shortest-first (5h before weekly, etc.). Omitted in legacy cached shapes — + * consumers treat a single remainingFraction as one unlabeled window. + */ + windows?: QuotaWindowEntry[] +} + +/** + * Migrate pre-pool quota keys while loading persisted account data. Current + * keys pass through unchanged, making this safe for every storage read. + */ +export function normalizeLegacyCachedQuota( + raw: AccountMetadataV3['cachedQuota'], +): Partial> | undefined { + if (!raw) return raw + + const hasLegacy = + 'gemini-pro' in raw || + 'gemini-flash' in raw || + 'claude' in raw || + 'gpt-oss' in raw + if (!hasLegacy) return raw + + const earlierResetTime = ( + a: string | undefined, + b: string | undefined, + ): string | undefined => { + if (!a) return b + if (!b) return a + return a < b ? a : b + } + + const minFraction = ( + a: QuotaGroupSummary | undefined, + b: QuotaGroupSummary | undefined, + ): QuotaGroupSummary | undefined => { + if (!a) return b + if (!b) return a + const fa = a.remainingFraction ?? 1 + const fb = b.remainingFraction ?? 1 + const winner = fa <= fb ? a : b + const loser = fa <= fb ? b : a + return { + ...winner, + resetTime: earlierResetTime(winner.resetTime, loser.resetTime), + } + } + + const gemini = minFraction( + raw.gemini as QuotaGroupSummary | undefined, + minFraction( + raw['gemini-pro'] as QuotaGroupSummary | undefined, + raw['gemini-flash'] as QuotaGroupSummary | undefined, + ), + ) + const nonGemini = minFraction( + raw['non-gemini'] as QuotaGroupSummary | undefined, + minFraction( + raw.claude as QuotaGroupSummary | undefined, + raw['gpt-oss'] as QuotaGroupSummary | undefined, + ), + ) + + return { + ...(gemini !== undefined ? { gemini } : {}), + ...(nonGemini !== undefined ? { 'non-gemini': nonGemini } : {}), + } } export interface PerModelQuotaEntry { 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/e2e-tests/src/harness.ts b/packages/e2e-tests/src/harness.ts index a220266..4bba3d1 100644 --- a/packages/e2e-tests/src/harness.ts +++ b/packages/e2e-tests/src/harness.ts @@ -37,6 +37,8 @@ export interface E2eHarness { createPlugin(options?: { clientOverrides?: Record extraDependencies?: Partial + /** Override specific fields in the project-level config written to disk. */ + configOverrides?: Record }): Promise /** Tear down all built plugins + the mock server + temp dirs. */ dispose(): Promise @@ -109,7 +111,10 @@ export async function createE2eHarness( soft_quota_threshold_percent: 100, quota_refresh_interval_minutes: 0, proactive_rotation_threshold_percent: 0, + background_quota_refresh: false, auto_update: false, + // Allow per-test overrides (e.g. enabling background_quota_refresh). + ...(pluginOptions.configOverrides ?? {}), }), ) const client = createFakeClient(pluginOptions.clientOverrides) diff --git a/packages/e2e-tests/src/mock-antigravity-server.test.ts b/packages/e2e-tests/src/mock-antigravity-server.test.ts index c502906..1fa429d 100644 --- a/packages/e2e-tests/src/mock-antigravity-server.test.ts +++ b/packages/e2e-tests/src/mock-antigravity-server.test.ts @@ -164,3 +164,117 @@ 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) + }) + }) + + 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 15f475e..e84d3dd 100644 --- a/packages/e2e-tests/src/mock-antigravity-server.ts +++ b/packages/e2e-tests/src/mock-antigravity-server.ts @@ -87,6 +87,30 @@ 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' + /** + * 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 + 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 +154,7 @@ export type Fixture = | GenerateContentFixture | ProjectDiscoveryFixture | QuotaSummaryFixture + | QuotaSummaryWindowFixture | GeminiCliQuotaFixture | StreamChunkedFixture | TokenExpiryFixture @@ -241,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, { @@ -281,6 +306,7 @@ export async function startMockAntigravityServer(): Promise { fixture: Fixture, _request: IncomingMessage, response: ServerResponse, + reqBody: string, ): Promise { switch (fixture.kind) { case 'json': { @@ -312,6 +338,57 @@ 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: string | undefined + try { + postedProject = + (JSON.parse(reqBody) as { project?: string }).project ?? '' + } catch { + // 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, { + error: { + code: 7, + message: 'PERMISSION_DENIED', + status: 'PERMISSION_DENIED', + }, + }) + return + } + } + // 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': { const candidates = [ { diff --git a/packages/e2e-tests/src/plugin-flow.e2e.test.ts b/packages/e2e-tests/src/plugin-flow.e2e.test.ts index 1ce4698..ca54dbe 100644 --- a/packages/e2e-tests/src/plugin-flow.e2e.test.ts +++ b/packages/e2e-tests/src/plugin-flow.e2e.test.ts @@ -322,6 +322,94 @@ describe('plugin flow (e2e)', () => { }) }) + it('V3 poller: background_quota_refresh=true exercises the default-on path and refreshes idle accounts', async () => { + // The harness globally sets background_quota_refresh:false to prevent + // network calls outside the fetch interceptor. This test overrides that + // and verifies the poller path through the real factory: + // - plugin boots with background_quota_refresh:true + // - accounts are loaded (seeded by beforeEach) + // - the poller fires (stale freshness gate) and calls the mock server + // - the sidebar state file is updated with non-empty quota + await withHarness(async (h) => { + const sidebarPath = process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + if (!sidebarPath) throw new Error('sidebar path missing') + + // Use a longer startup jitter (350ms) so we have a reliable window + // to write a stale sidebar AFTER the auth loader (which writes a fresh + // checkedAt) but BEFORE the first poller tick. This makes the only + // write that can advance past `pollBefore` be the poller's own tick. + const plugin = await h.createPlugin({ + configOverrides: { + background_quota_refresh: true, + background_quota_refresh_interval_minutes: 1, + }, + extraDependencies: { + // random=0.012 → jitter = floor(0.012 * 30000) = 360ms. + // Auth loader runs in ~20ms; we write the stale sidebar at ~250ms + // — leaving ~110ms before the first tick fires. + clock: { random: () => 0.012 }, + }, + }) + + // Warm up the auth loader so the AccountManager is populated. + // This will write a fresh checkedAt to the sidebar (~T+20ms). + await plugin.auth.loader( + async () => ({ + type: 'oauth' as const, + refresh: 'refresh-a|project-a|managed-a', + access: 'access-a', + expires: Date.now() + 3_600_000, + }), + {} as Parameters[1], + ) + + // Give the auth loader\'s async sidebar write time to land, then + // forcibly overwrite with a stale checkedAt so the poller\'s freshness + // gate passes on its first tick (~360ms from start). Must bypass the + // merge logic (which rejects stale writes) — write directly to disk. + await new Promise((r) => setTimeout(r, 200)) + require('node:fs').writeFileSync( + sidebarPath, + JSON.stringify({ + version: 1, + checkedAt: Date.now() - 60_000, + accounts: [], + activeRouting: {}, + routingAuthoritative: false, + }), + ) + + // Capture pollBefore AFTER the stale write. Now the ONLY write that + // can advance checkedAt past this point is the poller\'s first tick. + const pollBefore = Date.now() + + const deadline = pollBefore + 5_000 + let sidebarCheckedAt = 0 + while (Date.now() < deadline) { + await new Promise((r) => setTimeout(r, 100)) + try { + const raw = require('node:fs').readFileSync( + sidebarPath, + 'utf8', + ) as string + const parsed = JSON.parse(raw) as { checkedAt?: number } + sidebarCheckedAt = parsed.checkedAt ?? 0 + if (sidebarCheckedAt > pollBefore) break + } catch { + // File not written yet — retry. + } + } + + await plugin.dispose() + + // The poller must have fired, passed the freshness gate on the stale + // sidebar, and written a fresh checkedAt. This assertion can ONLY be + // satisfied by the poller — not by the auth loader (which ran before + // pollBefore and wrote the now-overwritten stale checkedAt). + expect(sidebarCheckedAt).toBeGreaterThan(pollBefore) + }) + }) + it('dispose() releases every subsystem and clears the auth loader runtime', async () => { await withHarness(async (h) => { const plugin = await h.createPlugin() diff --git a/packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts b/packages/e2e-tests/src/rpc-tui-flow.e2e.test.ts index 5682455..a00d0a3 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,40 @@ describe('rpc / tui flow (e2e)', () => { kind: 'projectDiscovery', 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: 'quotaSummary', - models: [{ id: 'gemini-3-flash', displayName: 'Gemini 3 Flash' }], + 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, + }, + ], + }, + { + 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/assets/antigravity.schema.json b/packages/opencode/assets/antigravity.schema.json index e3ef5d1..77ba2d2 100644 --- a/packages/opencode/assets/antigravity.schema.json +++ b/packages/opencode/assets/antigravity.schema.json @@ -360,6 +360,16 @@ ], "additionalProperties": false }, + "background_quota_refresh": { + "default": true, + "type": "boolean" + }, + "background_quota_refresh_interval_minutes": { + "default": 5, + "type": "number", + "minimum": 1, + "maximum": 60 + }, "auto_update": { "default": true, "type": "boolean", diff --git a/packages/opencode/src/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/account-command-oauth.test.ts b/packages/opencode/src/plugin/account-command-oauth.test.ts new file mode 100644 index 0000000..31ad478 --- /dev/null +++ b/packages/opencode/src/plugin/account-command-oauth.test.ts @@ -0,0 +1,315 @@ +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: 'invalid_grant', + })) + const service = createAccountCommandOAuthService({ + authorize, + exchange, + persist: mock(async () => {}), + listAccounts: mock(async () => initialRows), + }) + + const started = await service.start('session-1') + 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 () => { + 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', + ) + + // 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, + }) + }) + + 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', + 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, + 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 start a new OAuth flow.', + ) + }) + + 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 start a new OAuth flow.', + ) + 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), + }) + + const result = await service.finish('missing-session', 'callback-code') + + expect(result).toEqual({ + 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 new file mode 100644 index 0000000..79495e5 --- /dev/null +++ b/packages/opencode/src/plugin/account-command-oauth.ts @@ -0,0 +1,217 @@ +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 + /** + * 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 +} + +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 + } + // 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 + } + + // Wrap listAccounts so a lock-contention or I/O failure in the error + // branches below does not replace the real OAuth failure text with a + // generic apply-error message. Mirrors the stage-4 guarded pattern. + const safeListAccounts = async (): Promise => { + try { + return await options.listAccounts() + } catch { + return [] + } + } + + return { + async start(sessionId) { + const authorization = await options.authorize() + 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 safeListAccounts() } + }, + + async finish(sessionId, callbackInput, label) { + const pending = takePending(sessionId) + if (!pending) { + return { + text: 'OAuth session expired. Please start again.', + accounts: await safeListAccounts(), + } + } + + // Stage 1: parse the callback URL/code. Failures here mean the + // user pasted a malformed callback — exchange has not started. + let callback: ReturnType + try { + callback = parseOAuthCallbackInput(callbackInput, pending.state) + } catch { + return { + text: 'OAuth authentication failed: could not parse the callback. Please start a new OAuth flow.', + accounts: await safeListAccounts(), + } + } + if ('error' in callback) { + return { + text: `OAuth authentication failed: ${callback.error}. Please start a new OAuth flow.`, + accounts: await safeListAccounts(), + } + } + + // 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 start a new OAuth flow.', + accounts: await safeListAccounts(), + } + } + if (result.type === 'failed') { + return { + text: 'OAuth authentication failed. Please start a new OAuth flow and try again.', + accounts: await safeListAccounts(), + } + } + + // 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 start a new OAuth flow.', + accounts: await safeListAccounts(), + } + } + + // 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, + } + }, + + dispose() { + pendingBySession.clear() + }, + } +} diff --git a/packages/opencode/src/plugin/accounts.test.ts b/packages/opencode/src/plugin/accounts.test.ts index 77b8ffa..8423130 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,34 @@ 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') + // 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/auth-loader.test.ts b/packages/opencode/src/plugin/auth-loader.test.ts index edde9b6..a357b13 100644 --- a/packages/opencode/src/plugin/auth-loader.test.ts +++ b/packages/opencode/src/plugin/auth-loader.test.ts @@ -54,9 +54,12 @@ describe('createAuthLoader', () => { index: 0, email: 'stored@example.com', enabled: true, + parts: { refreshToken: 'stored-refresh' }, + cachedQuota: undefined, }, ], requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), dispose: mock(async () => {}), } const { lifecycle, replacements } = createLifecycle() @@ -148,18 +151,32 @@ 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(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), dispose: mock(async () => {}), } const secondManager = { 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(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), dispose: mock(async () => {}), } const managers = [firstManager, secondManager] @@ -229,4 +246,220 @@ 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(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + 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(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + 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(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + 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(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + 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 eae2685..2297ed6 100644 --- a/packages/opencode/src/plugin/auth-loader.ts +++ b/packages/opencode/src/plugin/auth-loader.ts @@ -1,6 +1,10 @@ +import { createHash } from 'node:crypto' +import { getHealthTracker } from '@cortexkit/antigravity-auth-core' import { buildSidebarMachineStateFromAccounts, + isAccountCurrent, setSidebarMachineState, + toCapturedTier, } from '../sidebar-state' import { AccountManager } from './accounts' import { isOAuthAuth } from './auth' @@ -31,6 +35,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 +57,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 +107,7 @@ export function createAuthLoader({ createFetch, onGetAuth, dependencies, -}: CreateAuthLoaderOptions) { +}: CreateAuthLoaderOptions): LoadAndInstallRuntime & AuthLoaderHandle { const deps: AuthLoaderDependencies = { loadAccounts: dependencies?.loadAccounts ?? loadAccounts, clearAccounts: dependencies?.clearAccounts ?? clearAccounts, @@ -83,6 +120,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( { @@ -95,10 +139,80 @@ 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() + + // 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 }) + 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) => { + const activeByFamily = accountManager.getActiveIndexByFamily() + return { + index: entry.index, + label: entry.label, + enabled: entry.enabled, + current: isAccountCurrent(entry.index, activeByFamily), + coolingDownUntil: entry.coolingDownUntil, + healthScore: getHealthTracker().getScore(entry.index), + cachedQuota: entry.cachedQuota, + // Stamp the sidebar snapshot so the projection can detect a + // stale cache that landed on the wrong account (the manager's + // `cachedQuotaAccountId` is keyed to whatever account actually + // produced the snapshot — the live refresh-token hash is the + // expected identity at this slot). + cachedQuotaAccountId: entry.cachedQuotaAccountId, + currentQuotaAccountId: refreshTokenIdentity( + entry.parts.refreshToken, + ), + tier: toCapturedTier(entry), + } + }), + ), + ) + } + + async function runLoader( getAuth: GetAuth, provider: Provider, - ): Promise> => { + ): Promise> { onGetAuth?.(getAuth) let auth = await getAuth() @@ -164,45 +278,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 +299,47 @@ export function createAuthLoader({ return { apiKey: '', - 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), } } + + 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 { + const next = reloadChain.then(() => reloadRuntime(getAuth)) + reloadChain = next.catch(() => {}) + return next + } + // `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 } diff --git a/packages/opencode/src/plugin/background-quota-refresh.test.ts b/packages/opencode/src/plugin/background-quota-refresh.test.ts new file mode 100644 index 0000000..0d3034c --- /dev/null +++ b/packages/opencode/src/plugin/background-quota-refresh.test.ts @@ -0,0 +1,1981 @@ +import { afterEach, beforeEach, describe, expect, it, mock } from 'bun:test' +import { createHash } from 'node:crypto' +import { mkdirSync, mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' + +/** Mirror of the private quotaAccountIdentity used in production code. */ +function quotaAccountIdentity(refreshToken: string): string { + return createHash('sha256').update(refreshToken).digest('hex').slice(0, 16) +} + +import type { AccountMetadataV3 } from '@cortexkit/antigravity-auth-core' + +import { + readSidebarState, + SIDEBAR_STATE_ENV, + setSidebarMachineState, +} from '../sidebar-state' +import { + BackgroundQuotaRefresh, + type BackgroundQuotaRefreshOptions, + type PollerAccountView, +} from './background-quota-refresh' +import { createAntigravityPlugin } from './index' +import type { PluginClient, PluginInput } from './types' + +// ─── Minimal harness helpers ───────────────────────────────────────────────── + +function makeMinimalClient(): PluginClient { + return { + app: { log: mock(async () => {}) }, + auth: { set: mock(async () => {}) }, + session: { + abort: mock(async () => {}), + messages: mock(async () => ({ data: [] })), + prompt: mock(async () => {}), + }, + tui: { showToast: mock(async () => {}) }, + } as unknown as PluginClient +} + +function makePluginInput(client: PluginClient, directory: string): PluginInput { + return { + client, + project: {} as PluginInput['project'], + directory, + worktree: directory, + experimental_workspace: { register: mock(() => {}) }, + serverUrl: new URL('http://localhost:4096'), + $: (() => {}) as unknown as PluginInput['$'], + } +} + +// Build a fake AccountManager that satisfies BackgroundQuotaRefresh's usage. +type StubAccount = ReturnType[number] + +function makeStubAccount(refreshToken = 'tok-abc'): StubAccount { + return { + index: 0, + label: 'Account', + enabled: true, + parts: { refreshToken }, + cachedQuota: undefined, + cachedQuotaAccountId: undefined, + coolingDownUntil: undefined, + } +} + +function makeAccountManager( + accounts: StubAccount[] = [makeStubAccount()], + activeByFamily: { claude: number; gemini: number } = { claude: 0, gemini: 0 }, +): PollerAccountView { + return { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: (): AccountMetadataV3[] => + accounts.map((a) => ({ + refreshToken: a.parts.refreshToken, + enabled: a.enabled, + addedAt: 0, + lastUsed: 0, + })), + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => activeByFamily, + } +} + +// Build a QuotaManager stub that returns configurable results. +function makeQuotaManager( + opts: { + result?: 'ok' | 'error' | 'pending' + onRefreshAccounts?: () => void + delay?: number + } = {}, +) { + const { result = 'ok', delay = 0 } = opts + const refreshAccounts = mock( + async (accounts: Array<{ refreshToken: string }>, _options: unknown) => { + if (opts.onRefreshAccounts) opts.onRefreshAccounts() + if (delay > 0) await new Promise((r) => setTimeout(r, delay)) + if (result === 'error') throw new Error('quota fetch failed') + if (result === 'ok') { + return accounts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { + groups: { + 'non-gemini': { remainingFraction: 0.5, modelCount: 1 }, + }, + }, + })) + } + return [] + }, + ) + return { + refreshAccounts, + dispose: mock(async () => {}), + } +} + +type TierStubAccount = StubAccount & { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number +} + +function makeTierManager(account: TierStubAccount): PollerAccountView { + return { + getAccounts: () => [account], + getAccountsForQuotaCheck: () => [ + { + refreshToken: account.parts.refreshToken, + enabled: account.enabled, + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock((_index, patch) => Object.assign(account, patch)), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } +} + +type TickablePoller = { + runTick: () => Promise +} + +function makeTwoPollers(options: { + stateFile: string + now: number + firstManager: PollerAccountView + secondManager: PollerAccountView + firstQuotaManager?: ReturnType + secondQuotaManager?: ReturnType + firstLoadAccountTier?: BackgroundQuotaRefreshOptions['loadAccountTier'] + secondLoadAccountTier?: BackgroundQuotaRefreshOptions['loadAccountTier'] +}) { + const first = new BackgroundQuotaRefresh({ + intervalMs: 5 * 60_000, + sidebarStateFile: options.stateFile, + getAccountManager: () => options.firstManager, + quotaManager: (options.firstQuotaManager ?? + makeQuotaManager()) as unknown as import('./quota').QuotaManager, + now: () => options.now, + random: () => 0, + loadAccountTier: options.firstLoadAccountTier, + }) + const second = new BackgroundQuotaRefresh({ + intervalMs: 5 * 60_000, + sidebarStateFile: options.stateFile, + getAccountManager: () => options.secondManager, + quotaManager: (options.secondQuotaManager ?? + makeQuotaManager()) as unknown as import('./quota').QuotaManager, + now: () => options.now, + random: () => 0, + loadAccountTier: options.secondLoadAccountTier, + }) + + return { + first, + second, + tick: (poller: BackgroundQuotaRefresh) => + (poller as unknown as TickablePoller).runTick(), + dispose: () => Promise.all([first.dispose(), second.dispose()]), + } +} + +// ─── Suite ─────────────────────────────────────────────────────────────────── + +describe('BackgroundQuotaRefresh', () => { + let dir: string + let stateFile: string + let savedSidebarEnv: string | undefined + + beforeEach(() => { + savedSidebarEnv = process.env[SIDEBAR_STATE_ENV] + dir = mkdtempSync(join(tmpdir(), 'bg-quota-test-')) + stateFile = join(dir, 'sidebar-state.json') + process.env[SIDEBAR_STATE_ENV] = stateFile + }) + + afterEach(async () => { + // Restore rather than delete — a delete drops resolution to the real state dir. + if (savedSidebarEnv !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebarEnv + else delete process.env[SIDEBAR_STATE_ENV] + rmSync(dir, { recursive: true, force: true }) + }) + + // ── Timer: idempotent start/stop ────────────────────────────────────────── + + it('start is idempotent — a second call before the first tick does not double-schedule', async () => { + const calls: number[] = [] + const manager = makeAccountManager() + const quotaManager = makeQuotaManager({ + onRefreshAccounts: () => calls.push(Date.now()), + }) + const poller = new BackgroundQuotaRefresh({ + intervalMs: 5 * 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + // Delay the first tick past dispose so we can check nothing ran. + random: () => 1, // max startup jitter → no immediate tick + }) + poller.start() + poller.start() + await poller.dispose() + // No tick ran within the jitter window — both starts share one timer. + expect(calls).toHaveLength(0) + }) + + it('dispose after dispose is idempotent', async () => { + const manager = makeAccountManager() + const quotaManager = makeQuotaManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 1, + }) + poller.start() + await poller.dispose() + await poller.dispose() // must not throw + }) + + // ── V4: unexpected throw before internal reschedule still reschedules ────── + + it('V4 liveness: a throw escaping runTick before any scheduleNext still reschedules', async () => { + // runTick has three internal reschedule paths (freshness skip, lock held, + // lock throws). A throw that ESCAPES the function before any of those + // paths previously meant no next tick was queued — the poller died. + // The fix adds a reschedule in the outer .catch() so the poller + // degrades to a delayed retry rather than silently stopping. + // + // We verify this by observing that scheduleNext IS called after the throw, + // rather than waiting for the full ERROR_JITTER_MS delay (15 s). + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const manager = makeAccountManager() + const quotaManager = makeQuotaManager() + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + + // Intercept scheduleNext to observe calls after the catch. + type PollerInternals = { + runTick: () => Promise + scheduleNext: (ms: number) => void + } + const inner = poller as unknown as PollerInternals + const scheduleNextCalls: number[] = [] + const origScheduleNext = inner.scheduleNext.bind(poller) + inner.scheduleNext = (ms: number) => { + scheduleNextCalls.push(ms) + origScheduleNext(ms) + } + + // Replace runTick to throw before any internal scheduleNext can run. + const origRunTick = inner.runTick.bind(poller) + let firstCall = true + inner.runTick = async () => { + if (firstCall) { + firstCall = false + throw new Error('simulated unexpected escape from runTick') + } + return origRunTick() + } + + poller.start() + // The startup scheduleNext fires synchronously with jitteredStartDelay + // (random=0 → 0 ms) — tick runs almost immediately, throws, outer catch + // queues the recovery scheduleNext. Allow a brief event-loop drain. + await new Promise((r) => setTimeout(r, 50)) + await poller.dispose() + + // The outer catch must have called scheduleNext (the recovery reschedule). + // scheduleNextCalls[0] is the startup call; [1] is from the outer catch. + expect(scheduleNextCalls.length).toBeGreaterThanOrEqual(2) + // The recovery delay includes ERROR_JITTER_MS (15_000). + const recoveryDelay = scheduleNextCalls[1] + expect(recoveryDelay).toBeGreaterThanOrEqual(15_000) + }) + + // ── Error in tick does not kill the timer ───────────────────────────────── + + it('an error inside a tick does not prevent subsequent ticks', async () => { + let callCount = 0 + let secondCallResolve!: () => void + const secondCallSeen = new Promise((r) => { + secondCallResolve = r + }) + + const manager = makeAccountManager() + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const quotaManager = { + refreshAccounts: mock(async () => { + callCount++ + if (callCount === 1) throw new Error('first tick error') + secondCallResolve() + return [] + }), + dispose: mock(async () => {}), + } + + const poller = new BackgroundQuotaRefresh({ + // Very short interval so the second tick fires within the timeout. + intervalMs: 100, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + await Promise.race([ + secondCallSeen, + new Promise((_, reject) => + setTimeout( + () => reject(new Error('timeout waiting for second tick')), + 5_000, + ), + ), + ]) + await poller.dispose() + expect(callCount).toBeGreaterThanOrEqual(2) + }) + + // ── In-flight tick not double-entered ──────────────────────────────────── + + it('a slow tick prevents re-entry while it is in flight', async () => { + let active = 0 + let maxActive = 0 + let tickDone!: () => void + const tickBlocked = new Promise((r) => { + tickDone = r + }) + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const quotaManager = { + refreshAccounts: mock(async () => { + active++ + maxActive = Math.max(maxActive, active) + await tickBlocked + active-- + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + // Give the first tick a moment to start. + await new Promise((r) => setTimeout(r, 20)) + tickDone() + await poller.dispose() + expect(maxActive).toBe(1) + }) + + // ── Async dispose awaits in-flight tick ─────────────────────────────────── + + it('dispose awaits an in-flight tick before resolving', async () => { + const events: string[] = [] + + let tickStartedResolve!: () => void + const tickStartedP = new Promise((r) => { + tickStartedResolve = r + }) + let tickRelease!: () => void + const tickGate = new Promise((r) => { + tickRelease = r + }) + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const quotaManager = { + refreshAccounts: mock(async () => { + events.push('tick:start') + tickStartedResolve() + await tickGate + events.push('tick:finish') + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + await tickStartedP + // Tick is now in flight. Dispose must wait for it. + const disposeP = poller.dispose().then(() => { + events.push('dispose:done') + }) + tickRelease() + await disposeP + // dispose:done must follow tick:finish + expect(events).toEqual(['tick:start', 'tick:finish', 'dispose:done']) + }) + + it('S3 boundary: re-entry guard prevents inFlight overwrite when scheduleNext fires mid-tick', async () => { + // The natural test scenario (intervalMs < tick duration) never reaches the + // guard because runTick only calls scheduleNext AFTER it completes, so the + // second timer always fires after inFlight is already null. + // This test triggers the guard directly via reflection: call scheduleNext(0) + // while tick 1 is blocked, then assert the inFlight reference is unchanged. + // Without the guard the 0ms timer would overwrite inFlight; with it, the + // timer callback exits early and the original promise is preserved. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + let tickStartedResolve!: () => void + const tickStartedP = new Promise((r) => { + tickStartedResolve = r + }) + let releaseFirstTick!: () => void + const firstTickGate = new Promise((r) => { + releaseFirstTick = r + }) + + const quotaManager = { + refreshAccounts: mock(async () => { + tickStartedResolve() + await firstTickGate + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + + // Reach private members via cast to verify the guard's effect. + type PollerInternals = { + inFlight: Promise | null + scheduleNext: (ms: number) => void + } + const inner = poller as unknown as PollerInternals + + poller.start() + await tickStartedP // tick 1 is now blocked inside refreshAccounts + + const originalInFlight = inner.inFlight + expect(originalInFlight).not.toBeNull() + + // Force a 0ms timer while the tick is in flight. Without the guard, this + // overwrites inFlight; with it, the timer callback returns immediately. + inner.scheduleNext(0) + await new Promise((r) => setTimeout(r, 20)) // let the 0ms timer fire + + // Guard must have prevented the overwrite. + expect(inner.inFlight).toBe(originalInFlight) + + releaseFirstTick() + await poller.dispose() + }) + + // ── Freshness gate ──────────────────────────────────────────────────────── + + it('freshness gate: skips when checkedAt is fresher than threshold (5-minute interval)', async () => { + const intervalMs = 5 * 60_000 // 300 000 ms + // Threshold = max(300000-60000, floor(300000/2)) = max(240000, 150000) = 240 000 ms + // A checkedAt 239 s ago is still within the threshold → skip. + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 239_000, accounts: [] }, + { stateFile }, + ) + + const quotaManager = makeQuotaManager() + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => now, + }) + poller.start() + // Allow one event loop tick for the timer. + await new Promise((r) => setTimeout(r, 50)) + await poller.dispose() + expect(quotaManager.refreshAccounts).not.toHaveBeenCalled() + }) + + it('freshness gate: refreshes when checkedAt is staler than threshold (5-minute interval)', async () => { + const intervalMs = 5 * 60_000 + // checkedAt 241 s ago → stale → should refresh. + let refreshCalled = false + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 241_000, accounts: [] }, + { stateFile }, + ) + + const quotaManager = { + refreshAccounts: mock(async () => { + refreshCalled = true + refreshResolve() + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => now, + }) + poller.start() + await Promise.race([ + refreshedP, + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 3_000), + ), + ]) + await poller.dispose() + expect(refreshCalled).toBe(true) + }) + + it('freshness gate at 1-minute minimum (floor(intervalMs/2) = 30 s)', async () => { + const intervalMs = 60_000 // 1 minute + // Threshold = max(60000-60000, floor(60000/2)) = max(0, 30000) = 30 000 ms + // checkedAt 29 s ago → skip. + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 29_000, accounts: [] }, + { stateFile }, + ) + + const quotaManager = makeQuotaManager() + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => now, + }) + poller.start() + await new Promise((r) => setTimeout(r, 50)) + await poller.dispose() + expect(quotaManager.refreshAccounts).not.toHaveBeenCalled() + }) + + it('freshness gate at 1-minute minimum refreshes when stale (31 s age)', async () => { + const intervalMs = 60_000 + // checkedAt 31 s ago → stale → refresh. + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 31_000, accounts: [] }, + { stateFile }, + ) + + const quotaManager = { + refreshAccounts: mock(async () => { + refreshResolve() + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => now, + }) + poller.start() + await Promise.race([ + refreshedP, + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 3_000), + ), + ]) + await poller.dispose() + expect(quotaManager.refreshAccounts).toHaveBeenCalled() + }) + + it('freshness gate at 60-minute interval uses intervalMs - 60 000 threshold', async () => { + const intervalMs = 60 * 60_000 // 3 600 000 ms + // Threshold = max(3600000-60000, floor(3600000/2)) = max(3540000, 1800000) = 3 540 000 ms + // checkedAt 3 539 s ago → fresh → skip. + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 3_539_000, accounts: [] }, + { stateFile }, + ) + + const quotaManager = makeQuotaManager() + const manager = makeAccountManager() + const poller = new BackgroundQuotaRefresh({ + intervalMs, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => now, + }) + poller.start() + await new Promise((r) => setTimeout(r, 50)) + await poller.dispose() + expect(quotaManager.refreshAccounts).not.toHaveBeenCalled() + }) + + it('two pollers: a fresh snapshot from one writer does not starve the other tier slot', async () => { + const now = Date.now() + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const firstQuotaManager = makeQuotaManager() + const secondQuotaManager = makeQuotaManager() + const secondLoadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: now, + })) + const pair = makeTwoPollers({ + stateFile, + now, + firstManager: makeTierManager({ + ...makeStubAccount('tok-first-fresh-tier'), + capturedTierId: 'free-tier', + capturedPaidTierId: 'g1-pro-tier', + capturedTierAt: now, + capturedTierSchemaVersion: 1, + }), + secondManager: makeTierManager(makeStubAccount('tok-second-stale-tier')), + firstQuotaManager, + secondQuotaManager, + secondLoadAccountTier, + }) + + await pair.tick(pair.first) + await pair.tick(pair.second) + await pair.dispose() + + expect(firstQuotaManager.refreshAccounts).toHaveBeenCalledTimes(1) + expect(secondQuotaManager.refreshAccounts).not.toHaveBeenCalled() + expect(secondLoadAccountTier).toHaveBeenCalledTimes(1) + }) + + it('two pollers: a shared stale account resolves tier once in one lock window', async () => { + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 60_000, accounts: [] }, + { stateFile }, + ) + + const account = makeStubAccount('tok-shared-stale-tier') + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: now, + })) + const pair = makeTwoPollers({ + stateFile, + now, + firstManager: makeTierManager(account), + secondManager: makeTierManager(account), + firstLoadAccountTier: loadAccountTier, + secondLoadAccountTier: loadAccountTier, + }) + + await Promise.all([pair.tick(pair.first), pair.tick(pair.second)]) + await pair.dispose() + + expect(loadAccountTier).toHaveBeenCalledTimes(1) + }) + + it('two pollers: both tier slots eventually spend while their shared snapshot stays fresh', async () => { + const now = Date.now() + await setSidebarMachineState( + { checkedAt: now - 60_000, accounts: [] }, + { stateFile }, + ) + + const firstLoadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: now, + })) + const secondLoadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: now, + })) + const pair = makeTwoPollers({ + stateFile, + now, + firstManager: makeTierManager(makeStubAccount('tok-first-stale-tier')), + secondManager: makeTierManager(makeStubAccount('tok-second-stale-tier')), + firstLoadAccountTier, + secondLoadAccountTier, + }) + + await pair.tick(pair.first) + await pair.tick(pair.second) + await pair.dispose() + + expect(firstLoadAccountTier).toHaveBeenCalledTimes(1) + expect(secondLoadAccountTier).toHaveBeenCalledTimes(1) + }) + + // ── Lock: held → skip ──────────────────────────────────────────────────── + + it('skips when the fenced lock is held by another process', async () => { + // Acquire the lock ourselves so the poller sees "held". + const { acquireFencedFileLock } = await import( + '@cortexkit/antigravity-auth-core/file-lock' + ) + const lock = await acquireFencedFileLock({ + path: stateFile, + name: 'bg-quota-poll', + ttlMs: 60_000, + renew: false, + }) + expect(lock).not.toBeNull() + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const quotaManager = makeQuotaManager() + const manager = makeAccountManager() + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: Date.now(), + })) + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + loadAccountTier, + }) + poller.start() + await new Promise((r) => setTimeout(r, 100)) + await poller.dispose() + await lock?.release() + + // Lock was held → refresh must not have been called. + expect(quotaManager.refreshAccounts).not.toHaveBeenCalled() + expect(loadAccountTier).not.toHaveBeenCalled() + }) + + // ── Lock throws → FAIL CLOSED ──────────────────────────────────────────── + + it('fail-closed: acquireFencedFileLock throw skips the tick without a fallback mutex', async () => { + // Use a NUL-byte in the path: the kernel rejects mkdir("...\0...") with + // ENOENT/EINVAL immediately, which causes acquireFencedFileLock to throw + // before touching any lock file. Two pollers share the same valid + // sidebarStateFile (for the freshness read), but tick 1 uses the bad lock + // path to trigger the throw, tick 2 uses the real path to confirm the + // timer rescheduled and a real refresh is possible. + // + // Crucially: no "claim marker" or secondary mutex files appear anywhere. + // This is the load-bearing invariant from the module-level comment. + const nullPath = join(dir, 'bad\x00path', 'state.json') + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const quotaManager = { + refreshAccounts: mock(async () => { + refreshResolve() + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: Date.now(), + })) + + // Poller 1: bad lock path → tick throws → fail closed → refresh NOT called. + const pollerBad = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + // The lock is attempted against sidebarStateFile (the valid path) by + // default. We need the lock to use the nullPath instead. Since the lock + // path is derived from sidebarStateFile (by appending `.bg-quota-poll.lock`) + // we point sidebarStateFile at nullPath so the lock write uses it. + // The freshness read will fail (file doesn't exist) and return checkedAt=0, + // which passes the gate — and then the lock attempt throws on NUL-byte. + sidebarStateFile: nullPath, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + loadAccountTier, + }) + pollerBad.start() + // Give it a moment to attempt and fail the lock. + await new Promise((r) => setTimeout(r, 100)) + await pollerBad.dispose() + + // Lock threw → refresh must NOT have been called on the bad poller. + expect(quotaManager.refreshAccounts).not.toHaveBeenCalled() + expect(loadAccountTier).not.toHaveBeenCalled() + + // Confirm NO claim/marker files were created — the fail-closed path + // must not fall back to a secondary mutex. + const { readdirSync } = await import('node:fs') + const anyLockFiles = readdirSync(dir).filter( + (f) => + f.includes('.lock') || f.includes('.claim') || f.includes('.marker'), + ) + expect(anyLockFiles).toHaveLength(0) + + // Poller 2: real path → refresh succeeds → confirms the timer is still + // functional after a fail-closed skip (not a property of pollerBad's timer + // but of the code path: the catch reschedules, so another poller can run). + const pollerGood = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + pollerGood.start() + await Promise.race([ + refreshedP, + new Promise((_, reject) => + setTimeout( + () => reject(new Error('timeout waiting for good refresh')), + 5_000, + ), + ), + ]) + await pollerGood.dispose() + expect(quotaManager.refreshAccounts).toHaveBeenCalledTimes(1) + }) + + // ── N-contender bound ───────────────────────────────────────────────────── + + it('N-contender bound: concurrent pollers against a slow refresh never all refresh', async () => { + // A slow refresh is simulated by inserting a delay. We run N pollers with + // zero startup jitter, all stale, and count how many refreshes execute + // BEFORE the first write lands. The freshness gate + lock means at most + // one poller should execute the refresh. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + let refreshes = 0 + let firstRefreshDone!: () => void + const firstRefreshP = new Promise((r) => { + firstRefreshDone = r + }) + + const quotaManager = { + refreshAccounts: mock(async () => { + refreshes++ + // Simulate a slow fetch — enough for all other pollers to attempt. + await new Promise((r) => setTimeout(r, 150)) + firstRefreshDone() + return [] + }), + dispose: mock(async () => {}), + } + const manager = makeAccountManager() + + const N = 4 + const pollers = Array.from({ length: N }, () => { + const p = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + p.start() + return p + }) + + // Wait for the first refresh to start and finish. + await Promise.race([ + firstRefreshP, + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + // At this point the first refresh just completed. Count before cleanup. + const refreshesSoFar = refreshes + await Promise.all(pollers.map((p) => p.dispose())) + + // Only 1 refresh should have run before the first write landed — + // the freshness gate blocks all others once checkedAt is updated. + expect(refreshesSoFar).toBe(1) + }) + + // ── Windows survive a polled refresh ───────────────────────────────────── + + it('windows-survive: a polled refresh writes per-window data into the sidebar state', async () => { + // This test exercises the real path: + // BackgroundQuotaRefresh.refresh() + // → quotaManager.refreshAccounts (returns per-window data) + // → manager.updateQuotaCache + // → pushSidebarQuotaSnapshot (reads the live view) + // → sidebar state file (verified via readSidebarState) + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const windowData = { + 'non-gemini': { + remainingFraction: 0.72, + modelCount: 1, + windows: [ + { + window: 'weekly' as const, + remainingFraction: 0.72, + resetTime: new Date( + Date.now() + 7 * 24 * 60 * 60_000, + ).toISOString(), + }, + { + window: '5h' as const, + remainingFraction: 0.91, + resetTime: new Date(Date.now() + 5 * 60 * 60_000).toISOString(), + }, + ], + }, + } + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const accounts = [ + { + index: 0, + label: 'Account', + enabled: true, + parts: { refreshToken: 'tok-win' }, + cachedQuota: undefined as typeof windowData | undefined, + cachedQuotaAccountId: undefined as string | undefined, + coolingDownUntil: undefined as number | undefined, + }, + ] + + const quotaManager = { + refreshAccounts: mock(async (accts: Array<{ refreshToken: string }>) => { + return accts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { groups: windowData }, + })) + }), + dispose: mock(async () => {}), + } + + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: (): AccountMetadataV3[] => [ + { + refreshToken: 'tok-win', + enabled: true, + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache: ( + _index: number, + groups: Partial< + Record + >, + _token: string, + ) => { + accounts[0]!.cachedQuota = groups as typeof windowData + // Use the real identity derived from the refresh token so the + // currentQuotaAccountId computed by the poller (same hash) matches. + accounts[0]!.cachedQuotaAccountId = quotaAccountIdentity('tok-win') + refreshResolve() + }, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await Promise.race([ + // Wait a bit extra to let pushSidebarQuotaSnapshot complete after updateQuotaCache. + refreshedP.then(() => new Promise((r) => setTimeout(r, 200))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + await poller.dispose() + + const state = readSidebarState(stateFile) + expect(state.accounts).toHaveLength(1) + const quota = state.accounts[0]?.quota['non-gemini'] + expect(quota?.remainingPercent).toBe(72) + // Per-window entries must survive through the sidebar read path. + expect(quota?.windows).toBeDefined() + expect(quota?.windows?.length).toBeGreaterThanOrEqual(1) + const weekly = quota?.windows?.find((w) => w.window === 'weekly') + expect(weekly).toBeDefined() + expect(weekly?.remainingPercent).toBe(72) + }) + + // ── V1: stale cachedQuotaAccountId is dropped after a polled refresh ──── + + it('V1 identity: poller writes currentQuotaAccountId from live token, not from cached stamp', async () => { + // An account whose cachedQuotaAccountId does NOT match its live + // refresh-token identity — simulating an account that was reordered + // or replaced. After a poll-written snapshot the sidebar must DROP the + // stale cached quota rather than displaying the wrong account's bars. + // + // Without the fix both fields were the same value so the + // staleness check never fired. + + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const refreshToken = 'tok-fresh-account' + // A stale stamp that will never match the hash of refreshToken. + const staleStamp = 'stale-id-000000000000' + + const accounts = [ + { + index: 0, + label: 'Account', + enabled: true, + parts: { refreshToken }, + // Stale quota snapshot from a previous account at this index. + cachedQuota: { + 'non-gemini': { remainingFraction: 0.5, modelCount: 1 }, + } as ReturnType[0]['cachedQuota'], + cachedQuotaAccountId: staleStamp, + coolingDownUntil: undefined as number | undefined, + }, + ] + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const quotaManager = { + // The quota fetch itself succeeds and returns new data. + refreshAccounts: mock(async (accts: Array<{ refreshToken: string }>) => { + return accts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { + groups: { + 'non-gemini': { remainingFraction: 0.8, modelCount: 1 }, + }, + }, + })) + }), + dispose: mock(async () => {}), + } + + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: (_index: number, groups: unknown, _token: string) => { + // Update the in-memory cache but keep the STALE stamp — simulating + // an account where the identity hasn't been updated yet. + accounts[0]!.cachedQuota = groups as (typeof accounts)[0]['cachedQuota'] + // Intentionally DO NOT update cachedQuotaAccountId here: the stamp + // stays stale so the sidebar projection must detect the mismatch. + refreshResolve() + }, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await Promise.race([ + refreshedP.then(() => new Promise((r) => setTimeout(r, 200))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout waiting for poll')), 5_000), + ), + ]) + await poller.dispose() + + const state = readSidebarState(stateFile) + expect(state.accounts).toHaveLength(1) + // The stale stamp must cause the quota to be dropped from the sidebar. + const sidebarQuota = state.accounts[0]?.quota + expect(sidebarQuota?.['non-gemini']).toBeUndefined() + }) + + // ── Fix 1: poller keeps `current` flag ───────────────────────────────────────── + + it('Fix1-current: poller tick keeps current:true on the active account (index 1)', async () => { + // Two accounts; claude-active is index 1. After a poll tick the sidebar + // must carry current:true on account at index 1 and false on index 0. + // Before this fix the poller omitted getActiveIndexByFamily, defaulting + // isAccountCurrent to false for every account. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const accounts: StubAccount[] = [ + { ...makeStubAccount('tok-0'), index: 0 }, + { ...makeStubAccount('tok-1'), index: 1 }, + ] + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const quotaManager = { + refreshAccounts: mock(async (accts: Array<{ refreshToken: string }>) => { + return accts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { + groups: { 'non-gemini': { remainingFraction: 0.5, modelCount: 1 } }, + }, + })) + }), + dispose: mock(async () => {}), + } + + // claude-active = 1, so account at index 1 should be current. + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => + accounts.map((a) => ({ + refreshToken: a.parts.refreshToken, + enabled: true, + addedAt: 0, + lastUsed: 0, + })), + updateQuotaCache: (_i, _g, _t) => { + if (_i === accounts.length - 1) refreshResolve() + }, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 1, gemini: 1 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await Promise.race([ + refreshedP.then(() => new Promise((r) => setTimeout(r, 300))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + await poller.dispose() + + const state = readSidebarState(stateFile) + expect(state.accounts).toHaveLength(2) + const acct0 = state.accounts.find((a) => a.id === 'acct-0') + const acct1 = state.accounts.find((a) => a.id === 'acct-1') + // Index 1 is the active account — must be current. + expect(acct1?.current).toBe(true) + // Index 0 is not active. + expect(acct0?.current).toBe(false) + }) + + // ── Fix 2: tier round-trips through sidebar state file ──────────────────── + + it('Fix2-tier-poller: captured tier appears in sidebar after a poll tick', async () => { + // When the in-memory account carries capturedTierId/capturedTierAt the + // poller must include tier: { id, capturedAt } in the sidebar snapshot. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const capturedAt = Date.now() + const accounts: StubAccount[] = [ + { + ...makeStubAccount('tok-tier'), + index: 0, + capturedTierId: 'enterprise-tier', + capturedPaidTierId: 'g1-pro-tier', + capturedTierAt: capturedAt, + } as StubAccount & { capturedPaidTierId: string }, + ] + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const quotaManager = { + refreshAccounts: mock(async (accts: Array<{ refreshToken: string }>) => { + return accts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { + groups: { 'non-gemini': { remainingFraction: 0.9, modelCount: 1 } }, + }, + })) + }), + dispose: mock(async () => {}), + } + + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { + refreshToken: 'tok-tier', + enabled: true, + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache: (_i, _g, _t) => { + accounts[0]!.cachedQuotaAccountId = quotaAccountIdentity('tok-tier') + refreshResolve() + }, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await Promise.race([ + refreshedP.then(() => new Promise((r) => setTimeout(r, 300))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + await poller.dispose() + + const state = readSidebarState(stateFile) + expect(state.accounts).toHaveLength(1) + const acct = state.accounts[0] + expect(acct?.tier).toBeDefined() + expect(acct?.tier?.id).toBe('enterprise-tier') + expect(acct?.tier?.paidId).toBe('g1-pro-tier') + expect(acct?.tier?.capturedAt).toBe(capturedAt) + }) + + it('keeps a successful quota refresh when the updated account token differs', async () => { + const accounts: StubAccount[] = [makeStubAccount('bare-token')] + const quotaGroups = { + gemini: { remainingFraction: 0.5, modelCount: 1 }, + } + const updateQuotaCache = mock(() => {}) + const quotaManager = { + refreshAccounts: mock(async () => { + accounts[0]!.parts.refreshToken = 'updated-token' + return [ + { + index: 0, + status: 'ok' as const, + quota: { groups: quotaGroups }, + updatedAccount: { + refreshToken: 'updated-token', + addedAt: 0, + lastUsed: 0, + }, + }, + ] + }), + dispose: mock(async () => {}), + } + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { + refreshToken: 'bare-token', + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await new Promise((resolve) => setTimeout(resolve, 300)) + await poller.dispose() + + expect(updateQuotaCache).toHaveBeenCalledWith( + 0, + quotaGroups, + 'updated-token', + ) + }) + + it('Fix2-tier-absent: account without tier yields no tier key in sidebar', async () => { + // An account with no capturedTierId must not emit tier in the snapshot. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const accounts: StubAccount[] = [ + { ...makeStubAccount('tok-notier'), index: 0 }, + ] + + let refreshResolve!: () => void + const refreshedP = new Promise((r) => { + refreshResolve = r + }) + + const quotaManager = { + refreshAccounts: mock(async (accts: Array<{ refreshToken: string }>) => + accts.map((_, i) => ({ + index: i, + status: 'ok' as const, + quota: { groups: {} }, + })), + ), + dispose: mock(async () => {}), + } + + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { + refreshToken: 'tok-notier', + enabled: true, + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache: () => { + refreshResolve() + }, + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + }) + poller.start() + + await Promise.race([ + refreshedP.then(() => new Promise((r) => setTimeout(r, 300))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + await poller.dispose() + + const state = readSidebarState(stateFile) + expect(state.accounts).toHaveLength(1) + const acct = state.accounts[0] + // Must be absent (undefined), not {} or { id: undefined }. + expect(acct?.tier).toBeUndefined() + }) + + // ── P2-A: tier resolved via loadAccountTier for existing accounts ────────────── + + it('P2-A: account with managedProjectId (existing account) gets tier after a tick via loadAccountTier', async () => { + // P2-A: existing accounts have managedProjectId packed in their refresh token. + // ensureProjectContext fast-paths on that field and never calls loadCodeAssist, + // so capturedTierId is never populated by the quota fetch path alone. + // The tier-staleness check must call loadAccountTier for the stalest account + // and persist the result via applyUpdatedAccount. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const refreshToken = 'tok-existing-account' + // Simulate an account with managedProjectId but NO captured tier yet. + const accounts: (StubAccount & { + capturedTierId?: string + capturedTierAt?: number + })[] = [ + { + ...makeStubAccount(refreshToken), + index: 0, + // capturedTierId absent -- the bug: existing account never got tier. + }, + ] + + let tierLoadCalled = false + const applyUpdatedAccountCalls: Array<{ + index: number + patch: { capturedTierId?: string; capturedTierAt?: number } + }> = [] + + const quotaManager = makeQuotaManager({ result: 'ok' }) + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { + refreshToken, + managedProjectId: 'proj-abc123', // existing account has managedProjectId + enabled: true, + addedAt: 0, + lastUsed: 0, + }, + ], + updateQuotaCache: mock((_i, groups, _t) => { + accounts[0]!.cachedQuotaAccountId = quotaAccountIdentity(refreshToken) + }), + applyUpdatedAccount: mock((index, patch) => { + applyUpdatedAccountCalls.push({ index, patch }) + if (patch.capturedTierId !== undefined) { + accounts[0]!.capturedTierId = patch.capturedTierId + accounts[0]!.capturedTierAt = patch.capturedTierAt + } + }), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + let tierResolve!: () => void + const tierResolved = new Promise((r) => { + tierResolve = r + }) + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + loadAccountTier: async (account) => { + tierLoadCalled = true + // Must be called with the AccountMetadataV3 that has managedProjectId. + expect(account.managedProjectId).toBe('proj-abc123') + const result = { id: 'free-tier', capturedAt: Date.now() } + // Signal after returning so the assertions run AFTER applyUpdatedAccount. + setTimeout(tierResolve, 0) + return result + }, + }) + poller.start() + + await Promise.race([ + tierResolved.then(() => new Promise((r) => setTimeout(r, 300))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout waiting for tier')), 5_000), + ), + ]) + await poller.dispose() + + // The tier loader must have been called. + expect(tierLoadCalled).toBe(true) + // applyUpdatedAccount must have been called with the resolved tier. + const tierCall = applyUpdatedAccountCalls.find( + (c) => c.patch.capturedTierId !== undefined, + ) + expect(tierCall).toBeDefined() + expect(tierCall?.patch.capturedTierId).toBe('free-tier') + expect(typeof tierCall?.patch.capturedTierAt).toBe('number') + }) + + it('P2-A-stale: account with outdated capturedTierAt triggers tier refresh', async () => { + // An account that already has a tier but its capturedTierAt is older than + // TIER_STALENESS_TTL_MS must be re-resolved. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const refreshToken = 'tok-stale-tier' + const oldTierAt = Date.now() - 25 * 60 * 60 * 1_000 // 25h ago (> 24h TTL) + const accounts: (StubAccount & { + capturedTierId?: string + capturedTierAt?: number + })[] = [ + { + ...makeStubAccount(refreshToken), + index: 0, + capturedTierId: 'old-tier', + capturedTierAt: oldTierAt, + }, + ] + + let tierLoadCalled = false + const quotaManager = makeQuotaManager({ result: 'ok' }) + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + let tierResolve!: () => void + const tierResolved = new Promise((r) => { + tierResolve = r + }) + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), // real time -- oldTierAt IS stale + loadAccountTier: async (_account) => { + tierLoadCalled = true + setTimeout(tierResolve, 0) + return { id: 'new-tier', capturedAt: Date.now() } + }, + }) + poller.start() + + await Promise.race([ + tierResolved.then(() => new Promise((r) => setTimeout(r, 100))), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + await poller.dispose() + + expect(tierLoadCalled).toBe(true) + }) + + it('backfills a missing paid tier field despite a fresh tier timestamp', async () => { + const now = 1_700_000_000_000 + const refreshToken = 'tok-paid-tier-backfill' + const accounts: Array< + StubAccount & { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number + } + > = [ + { + ...makeStubAccount(refreshToken), + capturedTierId: 'free-tier', + capturedTierAt: now - 60 * 60 * 1_000, + }, + ] + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + paidId: 'g1-pro-tier', + capturedAt: now, + })) + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock((_index, patch) => { + Object.assign(accounts[0]!, patch) + }), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: makeQuotaManager({ + result: 'ok', + }) as unknown as import('./quota').QuotaManager, + loadAccountTier, + now: () => now, + }) + + await (poller as unknown as { refresh: () => Promise }).refresh() + + expect(loadAccountTier).toHaveBeenCalledTimes(1) + expect(accounts[0]?.capturedPaidTierId).toBe('g1-pro-tier') + }) + + it('records a paid-tier-schema attempt when the upstream paid tier is absent', async () => { + const now = 1_700_000_000_000 + const refreshToken = 'tok-empty-paid-tier' + const accounts: Array< + StubAccount & { + capturedTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number + } + > = [ + { + ...makeStubAccount(refreshToken), + capturedTierId: 'free-tier', + capturedTierAt: now - 60 * 60 * 1_000, + }, + ] + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + capturedAt: now, + })) + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock((_index, patch) => { + Object.assign(accounts[0]!, patch) + }), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: makeQuotaManager({ + result: 'ok', + }) as unknown as import('./quota').QuotaManager, + loadAccountTier, + now: () => now, + }) + const refresh = ( + poller as unknown as { refresh: () => Promise } + ).refresh.bind(poller) + + await refresh() + await refresh() + + expect(loadAccountTier).toHaveBeenCalledTimes(1) + expect(accounts[0]?.capturedTierSchemaVersion).toBe(1) + }) + + it('does not re-resolve a fresh tier captured under the paid-tier schema', async () => { + const now = 1_700_000_000_000 + const refreshToken = 'tok-current-tier-schema' + const accounts: Array< + StubAccount & { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number + } + > = [ + { + ...makeStubAccount(refreshToken), + capturedTierId: 'free-tier', + capturedPaidTierId: 'g1-pro-tier', + capturedTierAt: now - 60 * 60 * 1_000, + }, + ] + const loadAccountTier = mock(async () => ({ + id: 'free-tier', + paidId: 'g1-pro-tier', + capturedAt: now, + })) + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: makeQuotaManager({ + result: 'ok', + }) as unknown as import('./quota').QuotaManager, + loadAccountTier, + now: () => now, + }) + + await (poller as unknown as { refresh: () => Promise }).refresh() + + expect(loadAccountTier).not.toHaveBeenCalled() + }) + + it('P2-A-fresh: account with fresh capturedTierAt does NOT trigger tier refresh', async () => { + // An account refreshed 1h ago (< 24h TTL) must not trigger loadAccountTier. + await setSidebarMachineState({ checkedAt: 0, accounts: [] }, { stateFile }) + + const refreshToken = 'tok-fresh-tier' + const freshTierAt = Date.now() - 1 * 60 * 60 * 1_000 // 1h ago (well within 24h) + const accounts: (StubAccount & { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number + })[] = [ + { + ...makeStubAccount(refreshToken), + index: 0, + capturedTierId: 'free-tier', + capturedPaidTierId: 'g1-pro-tier', + capturedTierAt: freshTierAt, + capturedTierSchemaVersion: 1, + }, + ] + + let tierLoadCalled = false + let refreshDone = false + const quotaManager = { + refreshAccounts: mock(async () => { + refreshDone = true + return [] + }), + dispose: mock(async () => {}), + } + const manager: PollerAccountView = { + getAccounts: () => [...accounts], + getAccountsForQuotaCheck: () => [ + { refreshToken, enabled: true, addedAt: 0, lastUsed: 0 }, + ], + updateQuotaCache: mock(() => {}), + applyUpdatedAccount: mock(() => {}), + requestSaveToDisk: mock(() => {}), + getActiveIndexByFamily: () => ({ claude: 0, gemini: 0 }), + } + + const poller = new BackgroundQuotaRefresh({ + intervalMs: 60_000, + sidebarStateFile: stateFile, + getAccountManager: () => manager, + quotaManager: quotaManager as unknown as import('./quota').QuotaManager, + random: () => 0, + now: () => Date.now(), + loadAccountTier: async () => { + tierLoadCalled = true + return null + }, + }) + poller.start() + + // Wait for the quota refresh to complete (a tick ran). + await Promise.race([ + new Promise((r) => { + const interval = setInterval(() => { + if (refreshDone) { + clearInterval(interval) + r() + } + }, 10) + }), + new Promise((_, reject) => + setTimeout(() => reject(new Error('timeout')), 5_000), + ), + ]) + // Small extra drain so any async tier check would have fired. + await new Promise((r) => setTimeout(r, 200)) + await poller.dispose() + + // Fresh tier -- tier loader must NOT have been invoked. + expect(tierLoadCalled).toBe(false) + }) +}) + +// ─── Wiring tests (through the real factory) ────────────────────────────────── + +describe('BackgroundQuotaRefresh wiring through createAntigravityPlugin', () => { + let tempDir: string + let stateFile: string + let fetchSpy: ReturnType + let savedSidebarEnvWiring: string | undefined + + beforeEach(() => { + savedSidebarEnvWiring = process.env[SIDEBAR_STATE_ENV] + tempDir = mkdtempSync(join(tmpdir(), 'bg-quota-wiring-')) + stateFile = join(tempDir, 'sidebar-state.json') + process.env[SIDEBAR_STATE_ENV] = stateFile + // Prevent the version check from hitting the network. + const { spyOn } = require('bun:test') as typeof import('bun:test') + fetchSpy = spyOn(globalThis, 'fetch').mockImplementation( + (async () => + new Response('not-found', { status: 404 })) as unknown as typeof fetch, + ) + }) + + afterEach(async () => { + fetchSpy.mockRestore() + mock.restore() + // Restore rather than delete — a delete drops resolution to the real state dir. + if (savedSidebarEnvWiring !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebarEnvWiring + else delete process.env[SIDEBAR_STATE_ENV] + rmSync(tempDir, { recursive: true, force: true }) + }) + + async function makePlugin( + directory: string, + config: Record, + pollerCreatedCb?: (p: BackgroundQuotaRefresh) => void, + ) { + mkdirSync(join(directory, '.opencode'), { recursive: true }) + await Bun.write( + join(directory, '.opencode', 'antigravity.json'), + JSON.stringify(config), + ) + const client = makeMinimalClient() + const input = makePluginInput(client, directory) + return createAntigravityPlugin('google', { + _onPollerCreated: pollerCreatedCb, + })(input) + } + + it('background_quota_refresh: false → NO poller instance is constructed', async () => { + const pluginDir = join(tempDir, 'plugin-disabled') + let pollerCreated = false + const plugin = await makePlugin( + pluginDir, + { background_quota_refresh: false }, + () => { + pollerCreated = true + }, + ) + await plugin.dispose() + // _onPollerCreated must not have been called — the block is guarded by + // `if (config.background_quota_refresh)` and the seam is wired INSIDE it. + expect(pollerCreated).toBe(false) + }) + + it('background_quota_refresh: true → exactly one poller is constructed and its dispose is registered', async () => { + const pluginDir = join(tempDir, 'plugin-enabled') + const capturedPollers: BackgroundQuotaRefresh[] = [] + const plugin = await makePlugin( + pluginDir, + { + background_quota_refresh: true, + background_quota_refresh_interval_minutes: 60, + }, + (p) => capturedPollers.push(p), + ) + // Exactly one instance was constructed. + expect(capturedPollers).toHaveLength(1) + // Replace dispose with a spy BEFORE plugin.dispose() is called so we can + // confirm the plugin's producer plumbing reaches through to the instance. + let pollerDisposeCalled = false + const original = capturedPollers[0]!.dispose.bind(capturedPollers[0]!) + capturedPollers[0]!.dispose = async () => { + pollerDisposeCalled = true + return original() + } + await plugin.dispose() + expect(pollerDisposeCalled).toBe(true) + }) + + it('plugin dispose awaits the poller in the PRODUCER phase (poller stops before the sidebar drain)', async () => { + // Evidence that the poller is a producer (not consumer): the lifecycle + // disposes producers BEFORE calling drainSidebarWrites. We instrument + // the poller's dispose to record its completion timestamp, and confirm + // plugin.dispose() resolves — which requires the producer chain to have + // completed (poller stopped, any in-flight tick awaited). + // + // A more surgical ordering probe would require intercepting lifecycle + // internals; this test focuses on the externally observable guarantee: + // plugin.dispose() fully awaits the poller, and the poller's dispose + // resolves within a tight deadline. + const pluginDir = join(tempDir, 'plugin-producer') + const capturedPollers: BackgroundQuotaRefresh[] = [] + const plugin = await makePlugin( + pluginDir, + { + background_quota_refresh: true, + background_quota_refresh_interval_minutes: 60, + }, + (p) => capturedPollers.push(p), + ) + expect(capturedPollers).toHaveLength(1) + + // Wrap dispose to measure how long the poller took to stop. + let pollerDisposeResolved = false + const orig = capturedPollers[0]!.dispose.bind(capturedPollers[0]!) + capturedPollers[0]!.dispose = async () => { + await orig() + pollerDisposeResolved = true + } + const disposed = await Promise.race([ + plugin.dispose().then(() => 'done'), + new Promise((r) => setTimeout(() => r('timeout'), 5_000)), + ]) + expect(disposed).toBe('done') + // The poller's dispose must have resolved before plugin.dispose() resolved. + expect(pollerDisposeResolved).toBe(true) + }) +}) diff --git a/packages/opencode/src/plugin/background-quota-refresh.ts b/packages/opencode/src/plugin/background-quota-refresh.ts new file mode 100644 index 0000000..4b13624 --- /dev/null +++ b/packages/opencode/src/plugin/background-quota-refresh.ts @@ -0,0 +1,475 @@ +/** + * Background quota poller. + * + * Runs a jittered, `unref()`'d timer that periodically refreshes quota for + * all accounts and pushes ONE sidebar snapshot per tick. Three-layer + * cross-process dedup: + * + * 1. Freshness gate — the only hard correctness bound for quota snapshots. + * Reads `checkedAt` from the on-disk sidebar state; skips quota work when + * it is fresher than + * `max(intervalMs − 60_000, floor(intervalMs / 2))`. The /2 floor + * ensures the gate still fires at the 1-minute minimum interval. + * + * 2. Advisory fenced lock — an optimization that lets a single process do + * the network work when several wake up simultaneously. Held → skip. + * + * 3. Lock throws → FAIL CLOSED (skip this tick, add jitter to next). A + * fallback mutex is NOT used here. The correct reasoning is: the + * freshness gate (layer 1) already bounds how many refreshes can land; + * a secondary "claim marker" file is just a worse re-implementation of + * the fenced lock we just failed to acquire — it replicates the same + * ownership problem without the protocol guarantees, and any transient + * I/O error that made the lock throw will likely affect a claim file + * too. The previous poller attempt fell into this trap across three + * review rounds; this comment exists so a future reviewer stops + * re-litigating it. + */ + +import { createHash } from 'node:crypto' +import type { + AccountMetadataV3, + QuotaGroup, + QuotaGroupSummary, +} from '@cortexkit/antigravity-auth-core' +import { + acquireFencedFileLock, + type FencedFileLock, +} from '@cortexkit/antigravity-auth-core/file-lock' + +import { readSidebarState, toCapturedTier } from '../sidebar-state' +import type { QuotaManager } from './quota' +import { pushSidebarQuotaSnapshot } from './quota' + +/** + * The subset of AccountManager that the poller needs. Using a narrow + * interface keeps the module decoupled from the large concrete class and + * lets tests pass minimal stubs without satisfying 90+ methods. + */ +export interface PollerAccountView { + getAccounts(): Array<{ + index: number + label?: string + enabled?: boolean + parts: { refreshToken: string } + cachedQuota?: Partial> + cachedQuotaAccountId?: string + coolingDownUntil?: number + /** Captured plan tier ID (raw upstream string). */ + capturedTierId?: string + /** Captured paid-tier ID (raw upstream string). */ + capturedPaidTierId?: string + /** Epoch ms when capturedTierId was last recorded. */ + capturedTierAt?: number + /** Schema version of the most recent tier capture. */ + capturedTierSchemaVersion?: number + }> + /** Must return the full AccountMetadataV3 shape so the quota manager + * can use token-refresh and project-context fields inside each fetch. */ + getAccountsForQuotaCheck(): AccountMetadataV3[] + updateQuotaCache( + accountIndex: number, + quotaGroups: Partial>, + expectedRefreshToken?: string, + ): void + /** + * Apply a subset of quota-result fields (e.g. captured tier) onto the + * in-memory account record without touching the quota cache itself. + */ + applyUpdatedAccount( + accountIndex: number, + patch: { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number + capturedTierSchemaVersion?: number + }, + expectedRefreshToken?: string, + ): void + requestSaveToDisk(): void + /** + * Returns the active account index per model family so the sidebar + * snapshot can stamp the `current` flag on the right account. + */ + getActiveIndexByFamily(): { claude: number; gemini: number } +} + +/** + * Opaque identity derived from a refresh token. Three callers keep + * independent copies: this module, `command-data.ts` (ships into the TUI + * compiled tree and cannot reach the core barrel), and `index.ts`. All + * three must produce the same 16-char hex prefix. + */ +function quotaAccountIdentity(refreshToken: string): string { + return createHash('sha256').update(refreshToken).digest('hex').slice(0, 16) +} + +/** Minimum poll interval (1 minute, matching the config schema min). */ +const MIN_INTERVAL_MS = 60_000 + +/** Spread startup jitter so N processes don't all fire at once on launch. */ +const STARTUP_JITTER_MS = 30_000 + +/** Additional jitter added to the next tick when the lock mechanism throws. */ +const ERROR_JITTER_MS = 15_000 + +/** TTL for the advisory poll lock. Long enough to cover a slow network fetch. */ +const POLL_LOCK_TTL_MS = 60_000 + +/** Maximum age of a captured tier before the poller refreshes it. 24 h. */ +const TIER_STALENESS_TTL_MS = 24 * 60 * 60 * 1_000 + +/** A capture at this version distinguishes an absent paid tier from old data. */ +const CAPTURED_TIER_SCHEMA_VERSION = 1 + +export interface BackgroundQuotaRefreshOptions { + intervalMs: number + sidebarStateFile: string + /** + * Supplier for the live account pool. Called inside each locked phase, so it + * observes any concurrent add/remove that completed while waiting for the + * lock. + */ + getAccountManager: () => PollerAccountView | null + quotaManager: QuotaManager + /** + * Optional tier-refresh callback. When provided, the poller calls this + * once per tick for the account whose capturedTierAt is missing or older + * than TIER_STALENESS_TTL_MS. The callback handles its own token-refresh + * and must resolve to a tier object (or null on failure / nothing to + * report) -- failure must never reject, since tier is best-effort. + * + * Wired in index.ts through the same agyTransport seam the quota manager + * uses, so e2e tests can inject a mock alongside the mock quota fetch. + */ + loadAccountTier?: ( + account: AccountMetadataV3, + ) => Promise<{ id: string; paidId?: string; capturedAt: number } | null> + /** + * Clock / randomness seam. Defaults to the system clock. + */ + now?: () => number + random?: () => number +} + +/** + * Per-loader-instance background quota poller. + * + * Must be created after the plugin has initialised (so `getAccountManager()` + * is populated) and disposed via the producer lifecycle phase. + */ +export class BackgroundQuotaRefresh { + private readonly intervalMs: number + private readonly sidebarStateFile: string + private readonly getAccountManager: () => PollerAccountView | null + private readonly quotaManager: QuotaManager + private readonly loadAccountTier?: ( + account: AccountMetadataV3, + ) => Promise<{ id: string; paidId?: string; capturedAt: number } | null> + private readonly now: () => number + private readonly random: () => number + + private timer: ReturnType | null = null + /** Resolves when the currently-running tick completes (or immediately if none). */ + private inFlight: Promise | null = null + private disposed = false + + constructor(options: BackgroundQuotaRefreshOptions) { + this.intervalMs = options.intervalMs + this.sidebarStateFile = options.sidebarStateFile + this.getAccountManager = options.getAccountManager + this.quotaManager = options.quotaManager + this.loadAccountTier = options.loadAccountTier + this.now = options.now ?? (() => Date.now()) + this.random = options.random ?? Math.random + } + + /** Start the background timer. Idempotent: a second call is a no-op. */ + start(): void { + if (this.disposed || this.timer !== null) return + this.scheduleNext(this.jitteredStartDelay()) + } + + /** Stop the timer and await any in-flight tick. */ + async dispose(): Promise { + if (this.disposed) return + this.disposed = true + if (this.timer !== null) { + clearTimeout(this.timer) + this.timer = null + } + await this.inFlight + } + + // ── internals ───────────────────────────────────────────────────────────── + + private jitteredStartDelay(): number { + return Math.floor(this.random() * STARTUP_JITTER_MS) + } + + private jitteredInterval(extraJitterMs = 0): number { + // ±10% on top of the base interval prevents phase-locking across restarts. + const jitter = Math.floor(this.random() * this.intervalMs * 0.1) + return this.intervalMs + jitter + extraJitterMs + } + + private scheduleNext(delayMs: number): void { + if (this.disposed) return + this.timer = setTimeout(() => { + this.timer = null + // Re-entry guard: if the previous tick is still in flight (refresh + // duration > interval), skip this timer callback. The in-flight tick + // calls scheduleNext itself when it completes, so the poller continues. + // Without this guard, assigning this.inFlight here would overwrite the + // first tick's promise and dispose() would await the wrong one. + if (this.inFlight !== null) return + this.inFlight = this.runTick() + .catch(() => { + // Errors inside runTick are already caught per-phase. An unexpected + // throw here (before any internal reschedule) means no next tick + // was queued — reschedule with error jitter so the poller degrades + // to a delayed tick rather than dying permanently. + this.scheduleNext(this.jitteredInterval(ERROR_JITTER_MS)) + }) + .finally(() => { + this.inFlight = null + }) + }, delayMs) + // Unreffing lets the process exit while the timer is pending — the poller + // must never be the reason an idle session stays alive. + if (this.timer.unref) this.timer.unref() + } + + private async runTick(): Promise { + // ── 1. Quota freshness gate ────────────────────────────────────────── + // + // Bound: `max(intervalMs − 60_000, floor(intervalMs / 2))`. + // + // - intervalMs − 60_000: at 5-minute intervals this is 240 s — tight + // enough to prevent double refreshes while allowing drift. + // - floor(intervalMs / 2): the fallback that keeps the gate non-trivial + // at the 1-minute minimum (floor(60000/2) = 30 000 ms, i.e. 30 s). + const freshnessThresholdMs = Math.max( + this.intervalMs - MIN_INTERVAL_MS, + Math.floor(this.intervalMs / 2), + ) + const state = readSidebarState(this.sidebarStateFile) + const ageMs = this.now() - state.checkedAt + + // ── 2. Advisory fenced lock ────────────────────────────────────────── + let lock: FencedFileLock | null = null + try { + lock = await acquireFencedFileLock({ + path: this.sidebarStateFile, + name: 'bg-quota-poll', + ttlMs: POLL_LOCK_TTL_MS, + // No renewal: the tick completes well within the TTL, and a renewal + // timer would need its own dispose chain. + renew: false, + }) + } catch { + // Lock mechanism threw (I/O error, filesystem issue). FAIL CLOSED: + // skip this tick. See module-level invariant comment for why we do not + // invent a fallback mutex here. + this.scheduleNext(this.jitteredInterval(ERROR_JITTER_MS)) + return + } + + if (lock === null) { + // Another process holds the lock — its tier phase owns this tick's + // metadata work as well as any quota snapshot write. + this.scheduleNext(this.jitteredInterval()) + return + } + + try { + if (ageMs >= freshnessThresholdMs) await this.refresh() + else await this.refreshTierSlot() + } finally { + await lock.release().catch(() => { + // Release is best-effort; the TTL expires it if we can't reach it. + }) + } + + if (!this.disposed) { + this.scheduleNext(this.jitteredInterval()) + } + } + + private async refresh(): Promise { + const refreshedQuota = await this.refreshQuota() + await this.refreshTierSlot() + if (refreshedQuota) await this.pushSnapshot() + } + + private async refreshQuota(): Promise { + if (this.disposed) return false + + const manager = this.getAccountManager() + if (!manager) return false + + const accounts = manager.getAccountsForQuotaCheck() + if (accounts.length === 0) return false + + // Refresh all accounts, then stamp the live manager once per result. + // Using the shared quotaManager deduplicates in-flight refreshes with + // any concurrent modal or fetch-interceptor refresh. + let results: Awaited> + try { + results = await this.quotaManager.refreshAccounts(accounts, { + indexFor: (account) => accounts.indexOf(account), + // Do not force: the quota manager's per-account dedup and backoff + // apply — an account that already refreshed recently is skipped. + force: false, + }) + } catch { + return false + } + + if (this.disposed) return false + + // Re-resolve live indices after the await: a concurrent add/remove may + // have shifted positions. + const liveAccounts = manager.getAccounts() + const liveIndexByToken = new Map() + for (const entry of liveAccounts) { + liveIndexByToken.set(entry.parts.refreshToken, entry.index) + } + + let anyUpdated = false + for (let i = 0; i < results.length; i++) { + const result = results[i] + if (result?.status !== 'ok' || !result.quota?.groups) continue + const candidateTokens = [ + result.updatedAccount?.refreshToken, + accounts[i]?.refreshToken, + ].filter((token): token is string => Boolean(token)) + const liveIndex = candidateTokens + .map((token) => liveIndexByToken.get(token)) + .find((index): index is number => index !== undefined) + if (liveIndex === undefined) continue + const liveToken = liveAccounts[liveIndex]?.parts.refreshToken + if (!liveToken) continue + manager.updateQuotaCache(liveIndex, result.quota.groups, liveToken) + // Persist tier captured during ensureProjectContext (same tick, no extra + // network call) so the sidebar snapshot carries it after every poll. + const uaTier = toCapturedTier(result.updatedAccount ?? {}) + if (uaTier !== undefined) { + manager.applyUpdatedAccount( + liveIndex, + { + capturedTierId: uaTier.id, + ...(uaTier.paidId !== undefined + ? { capturedPaidTierId: uaTier.paidId } + : {}), + capturedTierAt: uaTier.capturedAt, + capturedTierSchemaVersion: CAPTURED_TIER_SCHEMA_VERSION, + }, + liveToken, + ) + } + anyUpdated = true + } + if (anyUpdated) manager.requestSaveToDisk() + + return true + } + + private async refreshTierSlot(): Promise { + if (this.disposed || !this.loadAccountTier) return + + const manager = this.getAccountManager() + if (!manager) return + + const liveAccounts = manager.getAccounts() + const accounts = manager.getAccountsForQuotaCheck() + const accountByToken = new Map( + accounts.map((account) => [account.refreshToken, account]), + ) + const now = this.now() + let stalestEntry: (typeof liveAccounts)[number] | undefined + for (const entry of liveAccounts) { + const needsPaidTierBackfill = + entry.capturedPaidTierId === undefined && + entry.capturedTierSchemaVersion !== CAPTURED_TIER_SCHEMA_VERSION + const isTierStale = + entry.capturedTierAt === undefined || + now - entry.capturedTierAt >= TIER_STALENESS_TTL_MS || + needsPaidTierBackfill + if ( + isTierStale && + (stalestEntry === undefined || + (entry.capturedTierAt ?? 0) < (stalestEntry.capturedTierAt ?? 0)) + ) { + stalestEntry = entry + } + } + + if (stalestEntry === undefined) return + + const refreshToken = stalestEntry.parts.refreshToken + const accountForTier = accountByToken.get(refreshToken) + if (!accountForTier) return + + try { + const tier = await this.loadAccountTier(accountForTier) + if (tier === null || this.disposed) return + const liveEntry = manager + .getAccounts() + .find((entry) => entry.parts.refreshToken === refreshToken) + if (!liveEntry) return + manager.applyUpdatedAccount( + liveEntry.index, + { + capturedTierId: tier.id, + ...(tier.paidId !== undefined + ? { capturedPaidTierId: tier.paidId } + : {}), + capturedTierAt: tier.capturedAt, + capturedTierSchemaVersion: CAPTURED_TIER_SCHEMA_VERSION, + }, + refreshToken, + ) + manager.requestSaveToDisk() + } catch { + // Tier is best-effort; a failure must never break the tick. + } + } + + private async pushSnapshot(): Promise { + // ONE sidebar write per poll. `pushSidebarQuotaSnapshot` reads the live + // account view (with the just-updated cached values) and writes atomically. + const getAccounts = () => { + const m = this.getAccountManager() + return m ? m.getAccounts() : null + } + + await pushSidebarQuotaSnapshot( + () => { + const accts = getAccounts() + if (!accts) return null + return accts.map((entry) => ({ + index: entry.index, + label: entry.label, + enabled: entry.enabled, + coolingDownUntil: entry.coolingDownUntil, + cachedQuota: entry.cachedQuota, + cachedQuotaAccountId: entry.cachedQuotaAccountId, + // Derive current identity from the LIVE refresh token, independent + // of whatever stamp the cached snapshot carries. Mirrors index.ts's + // two correct call sites so the sidebar can detect a stale snapshot + // after an account reorder. + currentQuotaAccountId: quotaAccountIdentity(entry.parts.refreshToken), + tier: toCapturedTier(entry), + })) + }, + 0, + () => { + const m = this.getAccountManager() + return m ? m.getActiveIndexByFamily() : null + }, + ).catch(() => { + // Sidebar write failure is not fatal; the next tick will retry. + }) + } +} diff --git a/packages/opencode/src/plugin/command-data.test.ts b/packages/opencode/src/plugin/command-data.test.ts index e3e480a..6d955fa 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' @@ -55,10 +56,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 { @@ -72,7 +71,12 @@ interface AccountFixture { label?: string cachedQuota?: QuotaGroupFixture cachedQuotaUpdatedAt?: number + cachedQuotaAccountId?: string accountIneligible?: boolean + coolingDownUntil?: number + healthScore?: number + capturedTierId?: string + capturedTierAt?: number } function makeAccountFixture( @@ -102,9 +106,14 @@ 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 + geminiActiveIndex?: number refreshResults?: Map now?: () => number /** When true, also wire a storage adapter (defaults to true). */ @@ -150,6 +159,7 @@ function makeHarness(options: { > | undefined, cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, + cachedQuotaAccountId: entry.cachedQuotaAccountId, accountIneligible: entry.accountIneligible, })), } @@ -202,6 +212,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() { @@ -215,7 +226,12 @@ function makeHarness(options: { | Partial> | undefined, cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, + cachedQuotaAccountId: entry.cachedQuotaAccountId, accountIneligible: entry.accountIneligible, + coolingDownUntil: entry.coolingDownUntil, + healthScore: entry.healthScore, + capturedTierId: entry.capturedTierId, + capturedTierAt: entry.capturedTierAt, })) }, getAccountsForQuotaCheck() { @@ -232,10 +248,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() { @@ -246,8 +269,8 @@ 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 } }, setAccountEnabled(index: number, enabled: boolean): boolean { const account = liveView[index] @@ -260,14 +283,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 { @@ -343,14 +372,19 @@ async function readSidebar(stateFile: string): Promise { } let dir: string +let savedSidebarEnv: string | undefined describe('createCommandDataService', () => { beforeEach(() => { + savedSidebarEnv = process.env[SIDEBAR_STATE_ENV] dir = mkdtempSync(join(tmpdir(), 'agy-command-data-')) }) afterEach(() => { - delete process.env[SIDEBAR_STATE_ENV] + // Restore rather than delete — a delete drops resolution to the real state dir. + if (savedSidebarEnv !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebarEnv + else delete process.env[SIDEBAR_STATE_ENV] rmSync(dir, { recursive: true, force: true }) }) @@ -361,11 +395,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 +407,7 @@ describe('createCommandDataService', () => { label: 'Backup', enabled: false, cachedQuota: { - 'gemini-flash': { remainingFraction: 0.15 }, + gemini: { remainingFraction: 0.15 }, }, }), ], @@ -403,17 +437,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 +465,7 @@ describe('createCommandDataService', () => { refreshToken: 'refresh-a', label: 'A', cachedQuota: { - claude: { remainingFraction: 0.5 }, + 'non-gemini': { remainingFraction: 0.5 }, }, }), ], @@ -446,7 +480,90 @@ describe('createCommandDataService', () => { expect(harness.quotaCallLog).toEqual([]) }) - it('refreshQuota() refreshes through the shared quota manager and persists by refresh token', async () => { + 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('normalizes legacy quota keys and silently skips truly unknown keys', async () => { + // S1 + N3: the legacy keys (claude → non-gemini, gpt-oss → non-gemini, + // gemini-pro/gemini-flash → gemini) are now NORMALIZED at read time, + // not silently dropped. Truly unknown keys (e.g. 'gpt-4') are still + // dropped. The canonical 'gemini' key carries through unchanged. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-a', + label: 'Account A', + cachedQuota: { + // Legacy key: maps to non-gemini with MIN fraction. + claude: { remainingFraction: 0.9 }, + // Truly unknown key: must be dropped (no mapping). + 'gpt-4': { remainingFraction: 0.8 }, + // Canonical key: must render as-is. + gemini: { remainingFraction: 0.5 }, + } as unknown as QuotaGroupFixture, + }), + ], + }) + + const rows = await harness.service.listAccounts() + + expect(rows).toHaveLength(1) + // Both canonical and normalized-from-legacy keys must render. + const quota = rows[0]?.quota + expect(quota?.find((q) => q.key === 'gemini')?.remainingPercent).toBe(50) + expect(quota?.find((q) => q.key === 'non-gemini')?.remainingPercent).toBe( + 90, + ) + // The truly unknown 'gpt-4' key must not appear in the output. + expect(quota?.find((q) => q.key === ('gpt-4' as never))).toBeUndefined() + }) + + it('refreshQuota() fetches every account, including an uncached new account, through the shared quota manager', async () => { const baseAccount = ( refreshToken: string, groups: Record, @@ -468,11 +585,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 +602,7 @@ describe('createCommandDataService', () => { baseAccount( 'refresh-b', { - 'gemini-flash': { + gemini: { remainingFraction: 0.25, resetTime: new Date(0).toISOString(), }, @@ -500,14 +617,12 @@ 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 } }, - cachedQuotaUpdatedAt: 1, }), ], refreshResults, @@ -533,15 +648,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. @@ -559,7 +674,7 @@ describe('createCommandDataService', () => { status: 'ok', quota: { groups: { - claude: { + 'non-gemini': { remainingFraction: 0.7, resetTime: new Date(0).toISOString(), modelCount: 1, @@ -581,13 +696,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, }), ], @@ -603,13 +718,19 @@ 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 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-flash'] - ?.remainingFraction, + harness.storage.accounts[1]?.cachedQuota?.gemini?.remainingFraction, ).toBe(0.05) }) @@ -622,7 +743,7 @@ describe('createCommandDataService', () => { status: 'ok', quota: { groups: { - claude: { + 'non-gemini': { remainingFraction: 0.9, resetTime: new Date(0).toISOString(), modelCount: 1, @@ -644,7 +765,7 @@ describe('createCommandDataService', () => { makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha', - cachedQuota: { claude: { remainingFraction: 0.1 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0.1 } }, cachedQuotaUpdatedAt: 1, }), ], @@ -657,7 +778,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') @@ -695,7 +816,7 @@ describe('createCommandDataService', () => { makeAccountFixture({ refreshToken: 'refresh-a', label: 'Alpha', - cachedQuota: { claude: { remainingFraction: 0.4 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0.4 } }, cachedQuotaUpdatedAt: 100, }), ], @@ -707,7 +828,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) }) @@ -739,6 +860,24 @@ describe('createCommandDataService', () => { expect(harness.storage.accounts[1]?.enabled).toBe(true) }) + it('toggleAccountEnabled() carries healthScore through the command row seam', async () => { + // An account whose tracker score is 42 must keep that score after + // a command mutate — the row seam must forward healthScore so the + // sidebar does not silently reset to the 100 default. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-a', + label: 'Alpha', + healthScore: 42, + }), + ], + }) + const rows = await harness.service.toggleAccountEnabled(0) + expect(rows).not.toBeNull() + expect(rows?.[0]?.healthScore).toBe(42) + }) + it('toggleAccountEnabled() rejects out-of-range indices without touching storage', async () => { const harness = makeHarness({ accounts: [ @@ -824,7 +963,11 @@ 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, @@ -832,6 +975,186 @@ describe('createCommandDataService', () => { }) }) + 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, + gemini: 0, + }) + }) + + 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() 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 + // 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: [ @@ -1047,7 +1370,7 @@ describe('createCommandDataService', () => { status: 'ok', quota: { groups: { - claude: { remainingFraction: 0.1, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.1, modelCount: 1 }, }, modelCount: 1, }, @@ -1064,7 +1387,7 @@ describe('createCommandDataService', () => { makeAccountFixture({ refreshToken: 'refresh-a' }), makeAccountFixture({ refreshToken: 'refresh-b', - cachedQuota: { claude: { remainingFraction: 0.8 } }, + cachedQuota: { 'non-gemini': { remainingFraction: 0.8 } }, }), ], refreshResults, @@ -1076,30 +1399,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 () => { @@ -1139,6 +1462,199 @@ describe('createCommandDataService', () => { expect(harness.storage.activeIndexByFamily?.claude).toBe(1) expect(harness.storage.activeIndexByFamily?.gemini).toBe(1) }) + + it('refreshQuotaRespectingBackoff() folds successful results into the live cache', async () => { + // A stale account whose cached quota is 10% non-gemini. After a + // successful non-forced refresh returning 80%, the live cache must + // carry the fresh 80% — NOT the stale 10%. + const refreshResults = new Map([ + [ + 'refresh-a', + { + index: 0, + status: 'ok', + quota: { + groups: { + 'non-gemini': { + remainingFraction: 0.8, + resetTime: new Date(0).toISOString(), + modelCount: 1, + }, + }, + modelCount: 1, + }, + updatedAccount: { + refreshToken: 'refresh-a', + addedAt: 0, + lastUsed: 0, + }, + }, + ], + ]) + + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-a', + label: 'Alpha', + cachedQuota: { 'non-gemini': { remainingFraction: 0.1 } }, + cachedQuotaUpdatedAt: 1, + }), + ], + refreshResults, + }) + + // Live cache carries stale 10% before the call. + const before = harness.liveView[0]?.cachedQuota as + | QuotaGroupFixture + | undefined + expect(before?.['non-gemini']?.remainingFraction).toBe(0.1) + + await harness.service.refreshQuotaRespectingBackoff() + + // Live cache updated to fresh 80%. + const after = harness.liveView[0]?.cachedQuota as + | QuotaGroupFixture + | undefined + expect(after?.['non-gemini']?.remainingFraction).toBe(0.8) + // Identity stamp must be set. + expect(harness.liveView[0]?.cachedQuotaAccountId).toBe( + quotaAccountIdentity('refresh-a'), + ) + // Must have been a non-forced refresh. + expect(harness.quotaCallLog.every((call) => call.force === false)).toBe( + true, + ) + }) + + it('refreshQuotaRespectingBackoff() preserves existing backoff behaviour when all accounts are fresh', async () => { + // When refreshResults is empty (no stubs match — simulating the quota + // manager skipping every account due to backoff/freshness), the + // method must NOT crash, NOT update the cache, and NOT write a + // different sidebar. + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'fresh-account', + label: 'Fresh', + cachedQuota: { 'non-gemini': { remainingFraction: 0.5 } }, + }), + ], + refreshResults: new Map(), // Empty → every result is 'error' + }) + + const beforeFraction = ( + harness.liveView[0]?.cachedQuota as QuotaGroupFixture | undefined + )?.['non-gemini']?.remainingFraction + + await harness.service.refreshQuotaRespectingBackoff() + + // Stale cache unchanged — no successful refresh occurred. + const afterFraction = ( + harness.liveView[0]?.cachedQuota as QuotaGroupFixture | undefined + )?.['non-gemini']?.remainingFraction + expect(afterFraction).toBe(beforeFraction) + // Identity stamp must NOT have been set (no update happened). + expect(harness.liveView[0]?.cachedQuotaAccountId).toBeUndefined() + }) +}) + +describe('P2-B: writeSidebar lands tier in the state FILE', () => { + let dir: string + let savedSidebarEnvTier: string | undefined + + beforeEach(() => { + savedSidebarEnvTier = process.env[SIDEBAR_STATE_ENV] + dir = mkdtempSync(join(tmpdir(), 'agy-tier-sidebar-')) + }) + + afterEach(() => { + if (savedSidebarEnvTier !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebarEnvTier + else delete process.env[SIDEBAR_STATE_ENV] + rmSync(dir, { recursive: true, force: true }) + }) + + it('P2-B-file: refreshQuota writes tier to sidebar state file (not just the returned row)', async () => { + // Previously writeSidebar built SidebarAccountRedactionInput WITHOUT tier, + // so the field landed in the CommandAccountRow but was thrown away before + // reaching the disk. This test asserts the state FILE carries it. + const capturedAt = 1_700_000_001_000 + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-tier-test', + capturedTierId: 'free-tier', + capturedTierAt: capturedAt, + cachedQuota: { 'non-gemini': { remainingFraction: 0.5 } }, + cachedQuotaAccountId: quotaAccountIdentity('refresh-tier-test'), + }), + ], + refreshResults: new Map([ + [ + 'refresh-tier-test', + { + index: 0, + status: 'ok' as const, + quota: { + groups: { + 'non-gemini': { remainingFraction: 0.6, modelCount: 1 }, + }, + modelCount: 1, + }, + updatedAccount: { + refreshToken: 'refresh-tier-test', + addedAt: 0, + lastUsed: 0, + }, + }, + ], + ]), + }) + + await harness.service.refreshQuota() + const state = await readSidebar(harness.stateFile) + + expect(state.accounts).toHaveLength(1) + const acct = state.accounts[0] + // Tier must reach the FILE, not just the returned row. + expect(acct?.tier).toBeDefined() + expect(acct?.tier?.id).toBe('free-tier') + expect(acct?.tier?.capturedAt).toBe(capturedAt) + }) + + it('P2-B-absent: account without tier produces no tier key in the sidebar file', async () => { + const harness = makeHarness({ + accounts: [ + makeAccountFixture({ + refreshToken: 'refresh-no-tier', + cachedQuota: { 'non-gemini': { remainingFraction: 0.5 } }, + cachedQuotaAccountId: quotaAccountIdentity('refresh-no-tier'), + }), + ], + refreshResults: new Map([ + [ + 'refresh-no-tier', + { + index: 0, + status: 'ok' as const, + quota: { groups: {}, modelCount: 0 }, + updatedAccount: { + refreshToken: 'refresh-no-tier', + addedAt: 0, + lastUsed: 0, + }, + }, + ], + ]), + }) + + await harness.service.refreshQuota() + const state = await readSidebar(harness.stateFile) + + expect(state.accounts).toHaveLength(1) + expect(state.accounts[0]?.tier).toBeUndefined() + }) }) // Keep `mock` import alive for symmetry with sibling suites. diff --git a/packages/opencode/src/plugin/command-data.ts b/packages/opencode/src/plugin/command-data.ts index 0c85a5d..eb72733 100644 --- a/packages/opencode/src/plugin/command-data.ts +++ b/packages/opencode/src/plugin/command-data.ts @@ -40,10 +40,14 @@ * wrong account. */ +import { createHash } from 'node:crypto' + import { buildSidebarMachineStateFromAccounts, + normalizeLegacyCachedQuota, type SidebarAccountRedactionInput, setSidebarMachineState, + toCapturedTier, } from '../sidebar-state' /** @@ -70,19 +74,21 @@ 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 resetTime?: string modelCount: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> } type CommandDataAccountQuotaResult = { @@ -105,7 +111,7 @@ type CommandDataAccountQuotaResult = { updatedAccount?: CommandDataAccountMetadata } -type CommandDataAccountStorage = { +export interface CommandDataAccountStorage { version: 4 activeIndex: number activeIndexByFamily?: { claude?: number; gemini?: number } @@ -133,12 +139,27 @@ 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 + /** Health score in [0, 100], forwarded from the live tracker. */ + healthScore?: number quota: Array<{ - key: 'claude' | 'gemini-pro' | 'gemini-flash' | 'gpt-oss' + key: 'gemini' | 'non-gemini' label: string remainingPercent: number | null resetAt?: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingPercent: number + resetAt?: number + }> }> + /** Captured plan tier. Absent when unknown; never defaulted. */ + tier?: { id: string; capturedAt: number } } /** @@ -149,17 +170,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 +189,36 @@ interface LiveAccountSnapshot { Record > cachedQuotaUpdatedAt?: number + cachedQuotaAccountId?: string accountIneligible?: boolean + coolingDownUntil?: number + healthScore?: number + /** Captured plan tier from the most recent loadCodeAssist response. */ + capturedTierId?: string + /** Captured paid-tier ID from the most recent loadCodeAssist response. */ + capturedPaidTierId?: string + /** Epoch ms when capturedTierId was last recorded. */ + capturedTierAt?: number } function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { - 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 rawCached = + entry.cachedQuotaAccountId && + entry.cachedQuotaAccountId !== quotaAccountIdentity(entry.refreshToken) + ? undefined + : entry.cachedQuota + // Normalize legacy pool keys (gemini-pro, gemini-flash, claude, gpt-oss) + // to canonical keys before the SUPPORTED_QUOTA_KEYS loop. Without this + // a legacy snapshot produces empty quota rows, and writeSidebar would + // permanently re-emit legacy keys on every dialog open until the next + // real quota refresh rewrites the on-disk snapshot. + const cached = normalizeLegacyCachedQuota( + rawCached as Parameters[0], + ) const quota: CommandAccountRow['quota'] = [] for (const key of SUPPORTED_QUOTA_KEYS) { const cachedEntry = cached?.[key] @@ -194,11 +236,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}` @@ -208,10 +266,41 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { label, enabled: entry.enabled, current: entry.active, + coolingDownUntil: entry.coolingDownUntil, + healthScore: entry.healthScore, quota, + tier: toCapturedTier(entry), } } +/** + * 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[] { + 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, + cachedQuotaAccountId: entry.cachedQuotaAccountId, + }), + ) +} + /** * Live AccountManager view the data service needs. * @@ -234,8 +323,6 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { * 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[] @@ -245,10 +332,17 @@ export interface CommandDataAccountManagerView { groups: Partial< Record >, + expectedRefreshToken?: string, ): 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 + * 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. */ @@ -328,8 +422,20 @@ export interface CommandDataService { * refresh token, bump `cachedQuotaUpdatedAt`, and push a label-only * sidebar snapshot. Returns the freshly persisted rows so the * dialog can re-render in place. + * + * Bypass flag is intentional: explicit user-triggered refreshes + * ("Refresh" button, direct slash-command with `refresh` argument) + * must always fetch, even if the quota manager has backed off. */ refreshQuota(): Promise + /** + * Non-forced quota check that RESPECTS the quota manager's per-account + * backoff. Used by the dialog OPEN path — which is a VIEW, not a user + * request to fetch. Accounts already fresh (or in backoff) are skipped; + * the result is discarded (the sidebar poller carries the freshness + * guarantee for idle accounts). + */ + refreshQuotaRespectingBackoff(): Promise /** * Pin `index` as the active account for every family the dialog * tracks (claude + gemini). Mutates the live AccountManager AND the @@ -381,39 +487,70 @@ export function createCommandDataService( 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 gptOss = row.quota.find((q) => q.key === 'gpt-oss') - const toFraction = ( - q: { remainingPercent: number | null } | undefined, - ): { remainingFraction?: number; resetTime?: string } | undefined => { - if (!q || q.remainingPercent == null) return undefined - return { remainingFraction: q.remainingPercent / 100 } - } + // Map a CommandAccountRow back to the SidebarAccountRedactionInput shape. + // Extracted to prevent the "sixth dropped field" pattern: every field in + // CommandAccountRow that must reach the sidebar lives here once, not as an + // inline literal that a future author has to remember to update in two places. + const toRedactionInput = ( + row: CommandAccountRow, + ): SidebarAccountRedactionInput => { + const gemini = row.quota.find((q) => q.key === 'gemini') + const nonGemini = row.quota.find((q) => q.key === 'non-gemini') + // Map a CommandAccountRow quota entry to cachedQuota pool shape + // so projectQuotaPoolForSidebar (the canonical projection seam) + // carries the per-window breakdown into the sidebar state. + const toPool = ( + q: + | { + remainingPercent: number | null + resetAt?: number + windows?: CommandAccountRow['quota'][number]['windows'] + } + | undefined, + ) => { + if (!q || q.remainingPercent == null) return undefined return { - index: row.index, - label: row.label, - enabled: row.enabled, - current: row.current, - cachedQuota: { - claude: toFraction(claude), - 'gemini-pro': toFraction(geminiPro), - 'gemini-flash': toFraction(geminiFlash), - 'gpt-oss': toFraction(gptOss), - }, + remainingFraction: q.remainingPercent / 100, + resetTime: + typeof q.resetAt === 'number' && Number.isFinite(q.resetAt) + ? new Date(q.resetAt).toISOString() + : undefined, + windows: q.windows?.map((w) => ({ + window: w.window, + remainingFraction: w.remainingPercent / 100, + resetTime: + typeof w.resetAt === 'number' && Number.isFinite(w.resetAt) + ? new Date(w.resetAt).toISOString() + : '', + })), } - }) - // Fire-and-forget — the sidebar writer is fenced by its own queue, + } + return { + index: row.index, + label: row.label, + enabled: row.enabled, + current: row.current, + healthScore: row.healthScore, + // Tier passes the redaction boundary unchanged (plan metadata, not PII). + // The stamp check was already done by toCommandAccountRow. + tier: row.tier, + cachedQuota: { + gemini: toPool(gemini), + 'non-gemini': toPool(nonGemini), + }, + } + } + + const writeSidebar = (rows: CommandAccountRow[]): void => { + const accounts = rows.map(toRedactionInput) + // Fire-and-forget -- the sidebar writer is fenced by its own queue, // so a transient lock contention cannot block the dialog response. void setSidebarMachineState( buildSidebarMachineStateFromAccounts(accounts, { checkedAt: now() }), { 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. }) } @@ -466,7 +603,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() @@ -486,6 +627,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, } } @@ -510,6 +655,77 @@ export function createCommandDataService( return rows }, + async refreshQuotaRespectingBackoff() { + const accountsForQuota = accountManagerView.getAccountsForQuotaCheck() + if (accountsForQuota.length === 0) return + + // Non-forced: the quota manager's per-account backoff and dedup + // apply. Accounts that are fresh or in backoff are skipped. Unlike + // the earlier fire-and-forget path, successful results ARE folded + // into the live cache so the sidebar snapshot reflects fresh data + // immediately — the dialog renders ONLY from the sidebar file, and + // the quota manager's push path does not update AccountManager + // entries for this caller. + let results: Awaited> + try { + results = await quotaManager.refreshAccounts(accountsForQuota, { + indexFor: (account) => accountsForQuota.indexOf(account), + force: false, + }) + } catch { + // Non-critical: dialog open must never fail because a background + // refresh check encountered backoff or a transient error. + return + } + + // Build updates akin to the forced refresh path — resolve live + // indexes after the network call (a concurrent add/remove may have + // shifted positions), fold successful results into the live cache, + // then push a sidebar snapshot. + const updates: Array<{ + refreshToken: string + groups?: Partial< + Record + > + }> = [] + for (const result of results) { + const refreshToken = + result.updatedAccount?.refreshToken ?? + accountsForQuota[result.index]?.refreshToken + if (!refreshToken) continue + const groups = + result.status === 'ok' && result.quota?.groups + ? result.quota.groups + : undefined + updates.push({ refreshToken, groups }) + } + + if (updates.length === 0) return + + const liveIndexByRefreshToken = new Map() + for (const entry of accountManagerView.getAccounts()) { + liveIndexByRefreshToken.set(entry.refreshToken, entry.index) + } + let liveQuotaChanged = false + for (const update of updates) { + const liveIndex = liveIndexByRefreshToken.get(update.refreshToken) + if (liveIndex === undefined || !update.groups) continue + accountManagerView.updateQuotaCache( + liveIndex, + update.groups, + update.refreshToken, + ) + liveQuotaChanged = true + } + if (liveQuotaChanged) accountManagerView.requestSaveToDisk() + + // Push a fresh sidebar snapshot — the live view now carries the + // just-refreshed quota plus any skipped (backoff/fresh) accounts + // whose cached percentages are unchanged. + const rows = projectRows() + writeSidebar(rows) + }, + async setCurrentAccount(index) { return mutateLiveAndStorage({ action: 'setCurrent', index }) }, @@ -549,9 +765,32 @@ 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. 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 liveIndexes = accountManagerView.getActiveIndexByFamily() + liveCurrentTokens.claude = + liveAccounts[liveIndexes.claude]?.refreshToken ?? undefined + liveCurrentTokens.gemini = + liveAccounts[liveIndexes.gemini]?.refreshToken ?? undefined + } + 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, @@ -587,13 +826,54 @@ 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 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, + ) + 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 found + // 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). + // + // Use tokenIdx (the storage-snapshot position) rather than the + // outer-scope `index` (live-read position): a concurrent mutation + // between the live read and the lock can diverge the two, and the + // storage-level cursor must match the storage-level array. + return tokenIdx < nextAccounts.length ? tokenIdx : 0 + } + const legacyClaude = + current.activeIndexByFamily?.claude ?? current.activeIndex + 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/plugin/commands.test.ts b/packages/opencode/src/plugin/commands.test.ts index d4d2c47..af0fa60 100644 --- a/packages/opencode/src/plugin/commands.test.ts +++ b/packages/opencode/src/plugin/commands.test.ts @@ -3,12 +3,16 @@ 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 type { CommandAccountRow } from './command-data' import { applyCommand, buildDialogPayload, createCommandExecuteBefore, + createSidebarRefresher, MODAL_COMMANDS, } from './commands' import { GEMINI_DUMP_COMMAND_NAME } from './gemini-dump' @@ -240,6 +244,8 @@ describe('createCommandExecuteBefore', () => { quota: {}, }, ], + refreshQuota: async () => [], + refreshQuotaRespectingBackoff: async () => {}, } as never, { isTuiConnected: () => true }, ) @@ -256,6 +262,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 +626,570 @@ 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('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: 'non-gemini' as const, + label: 'Non-Gemini', + 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['non-gemini']?.remainingPercent).toBe(50) + expect(state.accounts[1]?.quota).toEqual({}) + } finally { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + + it('V2 windows: per-window quota arrays survive the post-apply sidebar projection', async () => { + // After an account or quota apply, createSidebarRefresher rebuilds + // cachedQuota from CommandAccountRow.quota. Without the fix the + // `windows` array was not carried and the sidebar collapsed to a + // single aggregate bar after every apply. + const sidebarFile = join(dir, 'sidebar-windows-v2.json') + const previousSidebarFile = process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = sidebarFile + try { + const resetAt = Date.now() + 7 * 24 * 60 * 60_000 + const dialogRows: CommandAccountRow[] = [ + { + id: 'acct-0', + index: 0, + label: 'Account 1', + enabled: true, + current: true, + quota: [ + { + key: 'non-gemini' as const, + label: 'Non-Gemini', + remainingPercent: 60, + resetAt, + windows: [ + { + window: '5h' as const, + remainingPercent: 85, + resetAt: resetAt - 1000, + }, + { window: 'weekly' as const, remainingPercent: 60, resetAt }, + ], + }, + ], + }, + ] + + const refresher = createSidebarRefresher(() => []) + await refresher(dialogRows) + + const state = readSidebarState(sidebarFile) + expect(state.accounts).toHaveLength(1) + const quota = state.accounts[0]?.quota['non-gemini'] + expect(quota?.remainingPercent).toBe(60) + // Per-window entries must survive the post-apply sidebar write. + expect(quota?.windows).toBeDefined() + expect(quota?.windows?.length).toBeGreaterThanOrEqual(2) + const weekly = quota?.windows?.find((w) => w.window === 'weekly') + expect(weekly).toBeDefined() + expect(weekly?.remainingPercent).toBe(60) + } finally { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + + it('carries healthScore from dialog rows into the sidebar state file', async () => { + // Covers the M1 seam AND the M3 gap: the final projection inside + // createSidebarRefresher must not re-map fields (dropping healthScore), + // and the sidebar state file (not just the returned row) must carry + // the tracker score. A score of 42 should appear as health: 42 in the + // persisted state - not reset to the default 100. + const sidebarFile = join(dir, 'sidebar-health-seam.json') + const previousSidebarFile = process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = sidebarFile + try { + const refresher = createSidebarRefresher(() => []) + const dialogRows: CommandAccountRow[] = [ + { + id: 'acct-0', + index: 0, + label: 'Alpha', + enabled: true, + current: true, + healthScore: 42, + quota: [ + { + key: 'gemini' as const, + label: 'Gemini', + remainingPercent: 80, + windows: [ + { + window: '5h' as const, + remainingPercent: 80, + resetAt: Date.now() + 5 * 60 * 1000, + }, + { + window: 'weekly' as const, + remainingPercent: 75, + resetAt: Date.now() + 7 * 24 * 60 * 60_000, + }, + ], + }, + ], + }, + { + id: 'acct-1', + index: 1, + label: 'Beta', + enabled: false, + current: false, + healthScore: 17, + quota: [], + }, + ] + + await refresher(dialogRows) + + const state = readSidebarState(sidebarFile) + expect(state.accounts).toHaveLength(2) + // Core seam assertion: tracker scores must reach the state file. + expect(state.accounts[0]?.health).toBe(42) + expect(state.accounts[1]?.health).toBe(17) + // Per-window data must also survive the same projection. + const gemini = state.accounts[0]?.quota.gemini + expect(gemini?.windows).toBeDefined() + expect(gemini?.windows?.length).toBeGreaterThanOrEqual(2) + } finally { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + + 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 { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + + it('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 { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + + it('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 { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + + it('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 { + // Restore rather than delete when previousSidebarFile was unset — + // a delete drops resolution to the operator's real state dir. + if (previousSidebarFile !== undefined) { + process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = previousSidebarFile + } + } + }) + + it('starts a backoff-respecting quota check when the quota dialog opens without delaying its cached payload', async () => { + // N6: the dialog OPEN path now calls refreshQuotaRespectingBackoff + // (non-forced) instead of refreshQuota (forced). The fire-and-forget + // semantics are preserved: if buildDialogPayload accidentally awaits + // the call this test hangs and surfaces as a timeout. + const listAccounts = mock(async () => [ + { + id: 'acct-0', + index: 0, + label: 'Primary account', + enabled: true, + current: true, + quota: [], + }, + ]) + let refreshStarted = false + const refreshQuotaRespectingBackoff = mock(() => { + refreshStarted = true + return new Promise(() => {}) + }) + + const payload = await buildDialogPayload('antigravity-quota', '', { + client: {} as never, + sessionID: 'session-1', + settings: ctx.settings, + commandData: { listAccounts, refreshQuotaRespectingBackoff } as never, + }) + + expect(payload.knobs.accounts).toHaveLength(1) + expect(refreshStarted).toBe(true) + expect(refreshQuotaRespectingBackoff).toHaveBeenCalledTimes(1) + }) + + 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..6dfc584 100644 --- a/packages/opencode/src/plugin/commands.ts +++ b/packages/opencode/src/plugin/commands.ts @@ -25,8 +25,10 @@ import { isTuiConnected as defaultIsTuiConnected } from '../rpc/notifications' import type { CommandModalName } from '../rpc/protocol' import { buildSidebarMachineStateFromAccounts, + type SidebarAccountRedactionInput, setSidebarMachineState, } from '../sidebar-state' +import type { AccountCommandOAuthService } from './account-command-oauth' import type { CommandAccountRow, CommandDataService } from './command-data' import { executeGeminiDumpCommand, @@ -117,7 +119,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 +128,7 @@ interface CommandContext { * undefined and the dialog falls back to the legacy placeholder. */ commandData?: CommandDataService + accountOAuth?: AccountCommandOAuthService } /** @@ -149,15 +152,16 @@ 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. Opening the dialog is a + // VIEW — it must respect the quota manager's per-account backoff so + // a rate-limited account isn't hammered every time the dialog opens. + // Explicit user-triggered refreshes (inside the dialog) keep force:true. + if (context.commandData) { + void context.commandData.refreshQuotaRespectingBackoff().catch(() => {}) + } return { command, text: 'Antigravity quota', @@ -169,6 +173,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 +187,7 @@ export async function buildDialogPayload( action === 'list' ? action : 'list', + accounts, }, } } @@ -323,6 +331,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 +361,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 +422,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 +520,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 +529,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 @@ -696,30 +751,100 @@ 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 } } + /** Captured plan tier. Absent when unknown. */ + tier?: { id: string; capturedAt: number } }> | null, -): () => Promise { - return async () => { - const accounts = getAccounts() - if (!accounts || accounts.length === 0) return +): (accounts?: CommandAccountRow[]) => Promise { + return async (dialogAccounts) => { + // 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 { - await setSidebarMachineState( - buildSidebarMachineStateFromAccounts( - accounts.map((entry) => ({ + 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) => { + // 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 + // only for legacy rows that lack the field entirely (?? can't + // distinguish absent from present-but-undefined). + const liveCooldown = + 'coolingDownUntil' in entry + ? entry.coolingDownUntil + : liveByIndex.get(entry.index) + return { index: entry.index, label: entry.label, enabled: entry.enabled, - coolingDownUntil: entry.coolingDownUntil, - cachedQuota: entry.cachedQuota, - })), + current: entry.current, + coolingDownUntil: liveCooldown, + healthScore: entry.healthScore, + tier: entry.tier, + // Rebuild cachedQuota carrying the windows array so the sidebar + // doesn't collapse from per-window rows to a single aggregate bar + // after every apply. Mirrors the QuotaGroupSummary shape expected + // by projectQuotaPoolForSidebar inside buildSidebarMachineStateFromAccounts. + cachedQuota: Object.fromEntries( + entry.quota.flatMap((group) => { + if (group.remainingPercent == null) return [] + return [ + [ + group.key, + { + remainingFraction: group.remainingPercent / 100, + ...(group.resetAt === undefined + ? {} + : { resetTime: new Date(group.resetAt).toISOString() }), + ...(group.windows && group.windows.length > 0 + ? { + windows: group.windows.map((w) => ({ + window: w.window, + remainingFraction: w.remainingPercent / 100, + resetTime: + typeof w.resetAt === 'number' && + Number.isFinite(w.resetAt) + ? new Date(w.resetAt).toISOString() + : '', + })), + } + : {}), + }, + ], + ] + }), + ), + } + }) + : getAccounts() + if (!accounts || accounts.length === 0) return + try { + // Pass the already-projected accounts directly so every field the + // intermediate map built (healthScore, cachedQuota with windows, ...) + // reaches the redactor unchanged. A second .map here would silently + // drop any field not explicitly forwarded -- that is the class of + // defect this cast prevents. + await setSidebarMachineState( + buildSidebarMachineStateFromAccounts( + accounts as SidebarAccountRedactionInput[], ), ) } 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/plugin/config/schema.ts b/packages/opencode/src/plugin/config/schema.ts index 7f4ecad..d5aa8ca 100644 --- a/packages/opencode/src/plugin/config/schema.ts +++ b/packages/opencode/src/plugin/config/schema.ts @@ -513,6 +513,31 @@ export const AntigravityConfigSchema = z.object({ // Auto-Update // ========================================================================= + // ========================================================================= + // Background Quota Refresh + // ========================================================================= + + /** + * Enable a background timer that periodically refreshes quota for all + * accounts. Without this, idle sessions show stale sidebar bars + * because Antigravity is a poll-only API (no quota headers in responses). + * + * @default true + */ + background_quota_refresh: z.boolean().default(true), + + /** + * How often the background poller refreshes quota (in minutes). + * Valid range: 1–60. + * + * @default 5 + */ + background_quota_refresh_interval_minutes: z + .number() + .min(1) + .max(60) + .default(5), + /** * Enable automatic plugin updates. * @default true @@ -603,6 +628,8 @@ export const DEFAULT_CONFIG: AntigravityConfig = { quota_refresh_interval_minutes: 30, soft_quota_cache_ttl_minutes: 'auto', proactive_rotation_threshold_percent: 20, + background_quota_refresh: true, + background_quota_refresh_interval_minutes: 5, auto_update: true, signature_cache: { enabled: true, diff --git a/packages/opencode/src/plugin/debug.test.ts b/packages/opencode/src/plugin/debug.test.ts index 8249385..0288384 100644 --- a/packages/opencode/src/plugin/debug.test.ts +++ b/packages/opencode/src/plugin/debug.test.ts @@ -126,9 +126,8 @@ describe('debug sink redaction', () => { }) it('masks full project IDs in debug logs', async () => { - const { initializeDebug, startAntigravityDebugRequest } = await import( - './debug' - ) + const { initializeDebug, startAntigravityDebugRequest, closeDebugLog } = + await import('./debug') initializeDebug({ ...DEFAULT_CONFIG, @@ -154,8 +153,9 @@ describe('debug sink redaction', () => { projectId: knownProjectId, }) - // Allow the writeStream a tick to flush the line. - await new Promise((r) => setTimeout(r, 50)) + // Close the write stream so all buffered data is flushed to disk + // before reading the log file. + await closeDebugLog() // The debug log file is the most recent `antigravity-debug-*.log` in // the log dir. @@ -178,9 +178,8 @@ describe('debug sink redaction', () => { }) it('masks fingerprint User-Agent headers in the recorded headers dump', async () => { - const { initializeDebug, startAntigravityDebugRequest } = await import( - './debug' - ) + const { initializeDebug, startAntigravityDebugRequest, closeDebugLog } = + await import('./debug') initializeDebug({ ...DEFAULT_CONFIG, @@ -204,7 +203,7 @@ describe('debug sink redaction', () => { streaming: false, }) - await new Promise((r) => setTimeout(r, 50)) + await closeDebugLog() const files = readdirSync(logDir) const logFile = files diff --git a/packages/opencode/src/plugin/debug.ts b/packages/opencode/src/plugin/debug.ts index a1aaf1e..7fe9f7d 100644 --- a/packages/opencode/src/plugin/debug.ts +++ b/packages/opencode/src/plugin/debug.ts @@ -60,6 +60,30 @@ function closeDebugStream(): void { } } +/** + * Close the debug log stream and resolve once all buffered writes + * have been flushed to disk. Tests use this to wait deterministically + * for log lines before reading the log file; the internal reinitialize + * path calls {@link closeDebugStream} directly (fire-and-forget). + */ +export function closeDebugLog(): Promise { + const current = debugState?.logStream + if (!current) return Promise.resolve() + // Detach the reference so reinitialization won't double-close. + if (debugState) { + debugState.logStream = null + } + return new Promise((resolve) => { + current.once('close', () => resolve()) + current.once('error', () => resolve()) + try { + current.end() + } catch { + resolve() + } + }) +} + /** * Get the OS-specific config directory. */ diff --git a/packages/opencode/src/plugin/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/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 bac9fb9..179316a 100644 --- a/packages/opencode/src/plugin/index.ts +++ b/packages/opencode/src/plugin/index.ts @@ -1,16 +1,24 @@ +import { createHash } from 'node:crypto' import { join } from 'node:path' import { ANTIGRAVITY_PROVIDER_ID } from '../constants' import { createAutoUpdateCheckerHook } from '../hooks/auto-update-checker' import { drainNotifications, pushNotification } from '../rpc/notifications' import { getRpcDir } from '../rpc/rpc-dir' import { startRpcServer } from '../rpc/rpc-server' -import { drainSidebarWrites, getSidebarStateFile } from '../sidebar-state' +import { + drainSidebarWrites, + getSidebarStateFile, + isAccountCurrent, + toCapturedTier, +} from '../sidebar-state' import { createAccountAccessService, promptAccountIndexForVerification, promptOpenVerificationUrl, } from './account-access' +import { createAccountCommandOAuthService } from './account-command-oauth' import { createAuthLoader } from './auth-loader' +import { BackgroundQuotaRefresh } from './background-quota-refresh' import { initDiskSignatureCache, shutdownDiskSignatureCache } from './cache' import { applyAntigravityProviderCatalog, @@ -19,6 +27,7 @@ import { import { type CommandDataService, createCommandDataService, + projectCommandAccountRows, } from './command-data' import { applyCommand, @@ -26,8 +35,8 @@ import { createSidebarRefresher, } from './commands' import { initRuntimeConfig, loadConfig } from './config' -import { getUserConfigPath as getUserConfigDir } from './config/loader' -import { initializeDebug } from './debug' +import { getUserConfigPath } from './config/loader' +import { closeDebugLog, initializeDebug } from './debug' import { type PluginDependencies, type PluginDependencyOverrides, @@ -44,9 +53,17 @@ import { type OperatorSettingsController, } from './operator-settings' import { persistAccountPool } from './persist-account-pool' -import { createOpenCodeQuotaManager, type QuotaManager } from './quota' +import { + createOpenCodeQuotaManager, + makeTierLoader, + type QuotaManager, +} from './quota' import { createSessionRecoveryHook } from './recovery' -import { initHealthTracker, initTokenTracker } from './rotation' +import { + getHealthTracker, + initHealthTracker, + initTokenTracker, +} from './rotation' import { AgySessionRegistry } from './session-context' import { clearAccounts, @@ -62,6 +79,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 @@ -74,6 +104,12 @@ export interface CreateAntigravityPluginOptions { * randomness. Defaults to production implementations. */ dependencies?: PluginDependencyOverrides + /** + * Test-only seam: called synchronously after the background poller is + * constructed, before `start()`. Lets wiring tests capture the instance + * without reaching into lifecycle internals. + */ + _onPollerCreated?: (poller: BackgroundQuotaRefresh) => void } export function registerQuotaManagerProducer( @@ -135,6 +171,12 @@ export const createAntigravityPlugin = drainSidebarWrites, }) const quotaManager = createOpenCodeQuotaManager(client, providerId, { + // Route quota fetches through the injected transport so e2e and + // custom-host deployments stay on the loopback/mock server for both + // the request path and the background poller. Without this the + // poller always hits the production Antigravity endpoint even when + // the caller injected a mock transport. + fetchVia: dependencies.agyTransport, // Bind to the live AccountManager so every refresh (manual or // background) pushes the freshly-updated quota percentages into the // sidebar without an extra RPC. The wrapper reads lazily so the @@ -142,14 +184,27 @@ export const createAntigravityPlugin = getAccountsForSidebar: () => { const manager = lifecycle.getAccountManager() if (!manager) return null + const activeByFamily = manager.getActiveIndexByFamily() return manager.getAccounts().map((entry) => ({ index: entry.index, label: entry.label, enabled: entry.enabled, + current: isAccountCurrent(entry.index, activeByFamily), coolingDownUntil: entry.coolingDownUntil, cachedQuota: entry.cachedQuota, + // Carry the identity stamp so the sidebar projection can + // detect a stale snapshot that landed on the wrong account + // after an index shift (see `redactAccountForSidebar`). + cachedQuotaAccountId: entry.cachedQuotaAccountId, + currentQuotaAccountId: quotaAccountIdentity(entry.parts.refreshToken), + tier: toCapturedTier(entry), })) }, + getActiveIndexByFamily: () => { + const manager = lifecycle.getAccountManager() + if (!manager) return null + return manager.getActiveIndexByFamily() + }, }) // Producer phase: the quota manager emits fire-and-forget sidebar // writes after every refresh. Its dispose() awaits any in-flight @@ -159,6 +214,30 @@ export const createAntigravityPlugin = // after drainSidebarWrites() already asserted the queue was empty. registerQuotaManagerProducer(lifecycle, quotaManager) + // Background quota poller: per-loader-instance timer that keeps idle + // account sidebar bars fresh between real requests. Registered as a + // producer so it is stopped and awaited BEFORE the sidebar drain, + // guaranteeing any final poll write lands before the drain flushes. + if (config.background_quota_refresh) { + const poller = new BackgroundQuotaRefresh({ + intervalMs: config.background_quota_refresh_interval_minutes * 60_000, + sidebarStateFile: getSidebarStateFile(), + getAccountManager: () => lifecycle.getAccountManager(), + quotaManager, + // Resolve plan tier for accounts with stale capturedTierAt. Uses + // the same token-refresh path as the quota manager without a separate + // network seam -- tier is best-effort metadata, not quota. + loadAccountTier: makeTierLoader(client, providerId), + // Thread the injected clock so e2e tests can control startup + // jitter and timing without needing a real 30-second wait. + now: dependencies.clock.now, + random: dependencies.clock.random, + }) + options._onPollerCreated?.(poller) + poller.start() + lifecycle.register({ dispose: () => poller.dispose() }, 'producer') + } + // Operator settings controller backs the /antigravity-* slash commands. // The controller loads existing persisted settings at first read, mutates // runtime config immediately, and serializes through the fenced-lock @@ -166,7 +245,10 @@ export const createAntigravityPlugin = const operatorSettings: OperatorSettingsController = createOperatorSettingsController({ projectConfigPath: join(directory, '.opencode', 'antigravity.json'), - userConfigPath: join(getUserConfigDir(), 'antigravity.json'), + // getUserConfigPath() already returns the full file path including + // 'antigravity.json' — do NOT join again or the path double-nests + // into /antigravity.json/antigravity.json. + userConfigPath: getUserConfigPath(), }) setRuntimeLogLevel(operatorSettings.get().log_level) lifecycle.register({ dispose: () => operatorSettings.dispose() }) @@ -196,25 +278,32 @@ export const createAntigravityPlugin = getAccounts: () => { const manager = lifecycle.getAccountManager() if (!manager) return [] - const current = - manager.getCurrentAccountForFamily('claude')?.index ?? 0 + const activeByFamily = manager.getActiveIndexByFamily() return manager.getAccounts().map((entry, index) => ({ index, refreshToken: entry.parts.refreshToken, label: entry.label, enabled: entry.enabled !== false, - active: index === current, + active: isAccountCurrent(index, activeByFamily), cachedQuota: entry.cachedQuota, cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt, + cachedQuotaAccountId: entry.cachedQuotaAccountId, accountIneligible: entry.accountIneligible, + coolingDownUntil: entry.coolingDownUntil, + healthScore: getHealthTracker().getScore(entry.index), + capturedTierId: entry.capturedTierId, + capturedPaidTierId: entry.capturedPaidTierId, + capturedTierAt: entry.capturedTierAt, })) }, getAccountsForQuotaCheck: () => { 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() @@ -222,11 +311,10 @@ export const createAntigravityPlugin = flushSaveToDisk: async () => { await lifecycle.getAccountManager()?.flushSaveToDisk() }, - activeIndex: () => { + getActiveIndexByFamily: () => { const manager = lifecycle.getAccountManager() - return manager - ? (manager.getCurrentAccountForFamily('claude')?.index ?? 0) - : 0 + if (!manager) return { claude: 0, gemini: 0 } + return manager.getActiveIndexByFamily() }, setAccountEnabled: (index, enabled) => { const manager = lifecycle.getAccountManager() @@ -285,15 +373,6 @@ export const createAntigravityPlugin = confirmOpenVerificationUrl: promptOpenVerificationUrl, }, }) - const oauthMethods = createOAuthMethods({ - client, - providerId, - config, - lifecycle, - accountAccess, - quotaManager, - getAuth: async () => (cachedGetAuth ? cachedGetAuth() : undefined), - }) const authLoader = createAuthLoader({ client, providerId, @@ -317,6 +396,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), @@ -327,12 +440,22 @@ export const createAntigravityPlugin = const refreshSidebar = createSidebarRefresher(() => { const manager = lifecycle.getAccountManager() if (!manager) return null + const activeByFamily = manager.getActiveIndexByFamily() return manager.getAccounts().map((entry) => ({ index: entry.index, label: entry.label, enabled: entry.enabled, + current: isAccountCurrent(entry.index, activeByFamily), coolingDownUntil: entry.coolingDownUntil, cachedQuota: entry.cachedQuota, + // Carry the identity stamp so the sidebar projection can + // detect a stale snapshot that landed on the wrong account + // after an index shift (see `redactAccountForSidebar`). + cachedQuotaAccountId: entry.cachedQuotaAccountId, + currentQuotaAccountId: quotaAccountIdentity( + entry.parts.refreshToken, + ), + tier: toCapturedTier(entry), })) }) const result = await applyCommand(request, { @@ -341,6 +464,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. @@ -350,6 +474,9 @@ export const createAntigravityPlugin = drain: drainNotifications, }) lifecycle.register({ dispose: () => rpcServer.stop() }) + // Flush the debug log stream after the sidebar drain so the last + // buffered lines land on disk before the process exits. Best-effort. + lifecycle.register({ dispose: () => closeDebugLog().catch(() => {}) }) return { dispose: async () => { diff --git a/packages/opencode/src/plugin/killswitch.test.ts b/packages/opencode/src/plugin/killswitch.test.ts index 53c4348..16b4bb4 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,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 — the - // candidate is allowed because gemini-flash is at 90%. + // 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-pro': { remainingFraction: 0.02, modelCount: 1 }, - 'gemini-flash': { remainingFraction: 0.9, modelCount: 1 }, + gemini: { remainingFraction: 0.02, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.9, modelCount: 1 }, }, cachedQuotaUpdatedAt: NOW, }, @@ -143,16 +152,17 @@ describe('evaluateKillswitchForAccount', () => { expect(decision.remainingPercent).toBe(90) }) - 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. + // 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-pro': { remainingFraction: 0.02, modelCount: 1 }, - 'gemini-flash': { remainingFraction: 0.9, modelCount: 1 }, + gemini: { remainingFraction: 0.02, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.9, modelCount: 1 }, }, cachedQuotaUpdatedAt: NOW, }, @@ -165,16 +175,15 @@ 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 }, + 'non-gemini': { remainingFraction: 0.9, modelCount: 1 }, }, cachedQuotaUpdatedAt: NOW, }, @@ -182,9 +191,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 +202,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 +210,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 +223,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 +249,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 +279,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 +309,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..0d805ee 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,16 @@ 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' + // 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' } return null } diff --git a/packages/opencode/src/plugin/lifecycle.test.ts b/packages/opencode/src/plugin/lifecycle.test.ts index 5d73d25..93bd41a 100644 --- a/packages/opencode/src/plugin/lifecycle.test.ts +++ b/packages/opencode/src/plugin/lifecycle.test.ts @@ -1,6 +1,11 @@ -import { describe, expect, it, mock } from 'bun:test' +import { afterEach, beforeEach, describe, expect, it, mock } from 'bun:test' +import { mkdirSync, readdirSync, readFileSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' import type { AccountManager } from './accounts' +import { DEFAULT_CONFIG } from './config' +import { initializeDebug } from './debug' import { createPluginLifecycle } from './lifecycle' import type { ProactiveRefreshQueue } from './refresh-queue' @@ -344,3 +349,107 @@ describe('PluginLifecycle RPC ownership', () => { expect(stop).toHaveBeenCalledTimes(1) }) }) + +describe('PluginLifecycle debug log disposal', () => { + let logDir: string + + beforeEach(() => { + logDir = join( + tmpdir(), + `agy-lifecycle-debug-${Math.random().toString(36).slice(2)}`, + ) + mkdirSync(logDir, { recursive: true }) + }) + + afterEach(() => { + // Reset module-global debug state before removing the temp dir so any + // stream targeting logDir is closed and the global points nowhere live. + // Without this, a stale WriteStream in debugState outlives the directory; + // under --isolate this is benign today (each file gets a fresh registry) + // but correct teardown guards against future tests added to this suite. + initializeDebug({ ...DEFAULT_CONFIG, debug: false }) + rmSync(logDir, { recursive: true, force: true }) + }) + + it('calls closeDebugLog as a consumer-phase disposable during lifecycle disposal', async () => { + const events: string[] = [] + const lifecycle = createPluginLifecycle({ + sessionRegistry: { clear: () => {} }, + shutdownDiskSignatureCache: async () => {}, + clearFetchState: () => {}, + drainSidebarWrites: async () => { + events.push('sidebar:drain') + }, + }) + // Simulate the plugin registration: closeDebugLog is registered as a + // consumer so it runs AFTER the sidebar drain. + lifecycle.register( + { + dispose: async () => { + events.push('debug-log:close') + }, + }, + 'consumer', + ) + + await lifecycle.dispose() + + // The debug log close must happen after the sidebar drain — consumers + // are disposed in phase 3, after the drain in phase 2. + expect(events.indexOf('sidebar:drain')).toBeLessThan( + events.indexOf('debug-log:close'), + ) + }) + + it('flushes buffered debug lines to disk when closeDebugLog runs during dispose', async () => { + // Initialize debug with a known log dir so closeDebugLog targets a + // predictable file path. + const { + initializeDebug, + getLogFilePath, + closeDebugLog, + startAntigravityDebugRequest, + } = await import('./debug') + + initializeDebug({ + ...DEFAULT_CONFIG, + debug: true, + debug_tui: false, + log_dir: logDir, + }) + + // Write a log line that would be buffered in the WriteStream. + startAntigravityDebugRequest({ + originalUrl: 'https://example.com/v1', + resolvedUrl: 'https://example.com/v1', + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: '"dispose-test-body"', + streaming: false, + }) + + const lifecycle = createPluginLifecycle({ + sessionRegistry: { clear: () => {} }, + shutdownDiskSignatureCache: async () => {}, + clearFetchState: () => {}, + }) + lifecycle.register({ + dispose: () => closeDebugLog().catch(() => {}), + }) + + await lifecycle.dispose() + + // After dispose, the debug log file must exist and contain the + // test marker written before shutdown. + const logPath = getLogFilePath() + expect(logPath).toBeTruthy() + const files = readdirSync(logDir) + const logFile = files + .filter((f) => f.startsWith('antigravity-debug-') && f.endsWith('.log')) + .sort() + .pop() + expect(logFile).toBeTruthy() + const contents = readFileSync(join(logDir, logFile!), 'utf8') + expect(contents).toContain('dispose-test-body') + }) +}) diff --git a/packages/opencode/src/plugin/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..5aea92e 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', ) }) @@ -76,15 +78,21 @@ describe('classifyQuotaGroup', () => { describe('pushSidebarQuotaSnapshot', () => { let dir: string let stateFile: string + let savedSidebarEnv: string | undefined beforeEach(() => { + // Save preload-pinned value so afterEach can restore it instead of + // deleting — a delete drops resolution to the operator's real state dir. + savedSidebarEnv = process.env[SIDEBAR_STATE_ENV] dir = mkdtempSync(join(tmpdir(), 'agy-quota-sidebar-')) stateFile = join(dir, 'sidebar-state.json') process.env[SIDEBAR_STATE_ENV] = stateFile }) afterEach(() => { - delete process.env[SIDEBAR_STATE_ENV] + if (savedSidebarEnv !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebarEnv + else delete process.env[SIDEBAR_STATE_ENV] rmSync(dir, { recursive: true, force: true }) }) @@ -100,12 +108,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 +122,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 +136,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 +150,7 @@ describe('pushSidebarQuotaSnapshot', () => { label: 'Primary Account', enabled: true, cachedQuota: { - claude: { remainingFraction: 0.6, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.6, modelCount: 1 }, }, }, ] @@ -154,7 +162,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 () => { @@ -174,6 +182,186 @@ 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('marks the active claude account as current when getActiveIndexByFamily is passed', async () => { + const getAccounts = (): QuotaSnapshotAccount[] => [ + { + index: 0, + label: 'Active', + enabled: true, + cachedQuota: { + 'non-gemini': { remainingFraction: 0.8, modelCount: 1 }, + }, + }, + { + index: 1, + label: 'Idle', + enabled: true, + }, + ] + + await pushSidebarQuotaSnapshot(getAccounts, 0, () => ({ + claude: 0, + gemini: 0, + })) + + const state = read() + expect(state.accounts).toHaveLength(2) + expect(state.accounts[0]?.current).toBe(true) + expect(state.accounts[1]?.current).toBe(false) + }) + + it('marks both accounts current when each family points to a different index', async () => { + const getAccounts = (): QuotaSnapshotAccount[] => [ + { index: 0, label: 'Claude', enabled: true }, + { index: 1, label: 'Middle', enabled: true }, + { index: 2, label: 'Gemini', enabled: true }, + ] + + await pushSidebarQuotaSnapshot(getAccounts, 0, () => ({ + claude: 0, + gemini: 2, + })) + + const state = read() + expect(state.accounts).toHaveLength(3) + expect(state.accounts[0]?.current).toBe(true) + expect(state.accounts[1]?.current).toBe(false) + expect(state.accounts[2]?.current).toBe(true) + }) + + it('defaults current to false when getActiveIndexByFamily is omitted (backward compat)', async () => { + const getAccounts = (): QuotaSnapshotAccount[] => [ + { index: 0, label: 'Acc', enabled: true }, + ] + + await pushSidebarQuotaSnapshot(getAccounts, 0) + + const state = read() + expect(state.accounts[0]?.current).toBe(false) + }) + + it('returns null from getActiveIndexByFamily → all accounts false', async () => { + const getAccounts = (): QuotaSnapshotAccount[] => [ + { index: 0, label: 'A', enabled: true }, + ] + + await pushSidebarQuotaSnapshot(getAccounts, 0, () => null) + + const state = read() + expect(state.accounts[0]?.current).toBe(false) + }) + it('fences the real quota wrapper sidebar enqueue before the lifecycle drain', async () => { const events: string[] = [] let releaseFetch!: () => void @@ -206,7 +394,7 @@ describe('pushSidebarQuotaSnapshot', () => { index: 0, email: 'primary@example.test', cachedQuota: { - claude: { remainingFraction: 0.42, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0.42, modelCount: 1 }, }, }, ], @@ -254,6 +442,12 @@ describe('pushSidebarQuotaSnapshot', () => { }) await drainSidebarWrites() expect(events).toEqual([ + 'fetch:start', + 'fetch:start', + 'fetch:start', + 'fetch:start', + // fetchGeminiCliQuota now also uses fetchVia (2 endpoints via the + // FetchGeminiCliQuotaOptions.fetchVia seam added for N2 testability). 'fetch:start', 'fetch:start', 'sidebar:write-start', @@ -267,4 +461,169 @@ describe('pushSidebarQuotaSnapshot', () => { fetchSpy.mockRestore() } }) + + it('N2: CLI rejection carries the real error message instead of the generic no-CLI-configured string', async () => { + // fetchGeminiCliQuota now propagates transport-level errors (network abort, + // DNS, socket hang, timeout) by collecting them per-endpoint and throwing + // when all endpoints fail. That throw reaches the outer .catch() in + // quota.ts, which sets geminiCliFetchError. Before N2 the catch was + // present but never fired — fetchGeminiCliQuota silently swallowed errors + // and returned { buckets: [] }, landing on the generic message. + // + // The new FetchGeminiCliQuotaOptions.fetchVia seam lets the test inject + // a transport that throws, driving the outer .catch() directly. + const CLI_THROW_MSG = 'socket hang up' + + const fetchSpy = spyOn(globalThis, 'fetch').mockImplementation((async ( + input: unknown, + ) => { + const url = + typeof input === 'string' + ? input + : input instanceof URL + ? input.toString() + : (input as Request).url + + if (url.includes('retrieveUserQuotaSummary')) { + return new Response( + JSON.stringify({ + groups: [ + { + displayName: 'Gemini Models', + buckets: [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly', + window: 'weekly', + resetTime: '2026-01-08T00:00:00Z', + remainingFraction: 0.6, + }, + ], + }, + ], + }), + { status: 200 }, + ) + } + if (url.includes('retrieveUserQuota')) { + // Simulate a transport-level rejection (socket hang up). + // fetchGeminiCliQuota collects this as an error and re-throws + // at end-of-loop, triggering the outer .catch() in quota.ts. + throw new Error(CLI_THROW_MSG) + } + return new Response( + JSON.stringify({ access_token: 'access-token', expires_in: 3600 }), + { status: 200 }, + ) + }) as unknown as typeof fetch) + + const client = { + auth: { set: mock(async () => {}) }, + } as unknown as PluginClient + const account: AccountMetadataV3 = { + refreshToken: `n2-rejection-${Date.now()}-${Math.random()}`, + managedProjectId: 'managed-n2', + projectId: 'project-n2', + addedAt: 0, + lastUsed: 0, + } + const manager = createOpenCodeQuotaManager(client, 'google') + let result: import('./quota.ts').AccountQuotaResult | undefined + try { + const results = await manager.refreshAccounts([account], { + indexFor: () => 0, + force: true, + }) + result = results[0] + } finally { + fetchSpy.mockRestore() + } + + // Summary succeeded — overall status must stay 'ok'. + expect(result?.status).toBe('ok') + if (result?.status !== 'ok') return + + // Summary groups must survive the CLI rejection. + expect(result.quota?.groups).toBeDefined() + + // The annotation must carry the REAL thrown message, not the generic + // 'No Gemini CLI quota available' that indicates a permanent absence. + // Multiple endpoints may each contribute the message (joined by '; '). + expect(result.geminiCliQuota?.error).toContain(CLI_THROW_MSG) + expect(result.geminiCliQuota?.error).not.toBe( + 'No Gemini CLI quota available', + ) + }) + + it('N2: HTTP-500 on CLI endpoint does NOT kill the summary (parallel-fetch isolation)', async () => { + // Complement: an HTTP 500 is treated as a transport error by the updated + // fetchGeminiCliQuota (errors[] + re-throw), so the same .catch() path + // fires and the summary still flows through. + const fetchSpy = spyOn(globalThis, 'fetch').mockImplementation((async ( + input: unknown, + ) => { + const url = + typeof input === 'string' + ? input + : input instanceof URL + ? input.toString() + : (input as Request).url + + if (url.includes('retrieveUserQuotaSummary')) { + return new Response( + JSON.stringify({ + groups: [ + { + displayName: 'Gemini Models', + buckets: [ + { + bucketId: 'gemini-weekly', + displayName: 'Weekly', + window: 'weekly', + resetTime: '2026-01-08T00:00:00Z', + remainingFraction: 0.6, + }, + ], + }, + ], + }), + { status: 200 }, + ) + } + if (url.includes('retrieveUserQuota')) { + return new Response('internal error', { status: 500 }) + } + return new Response( + JSON.stringify({ access_token: 'access-token', expires_in: 3600 }), + { status: 200 }, + ) + }) as unknown as typeof fetch) + + const client = { + auth: { set: mock(async () => {}) }, + } as unknown as PluginClient + const account: AccountMetadataV3 = { + refreshToken: `n2-500-${Date.now()}-${Math.random()}`, + managedProjectId: 'managed-n2b', + projectId: 'project-n2b', + addedAt: 0, + lastUsed: 0, + } + const manager = createOpenCodeQuotaManager(client, 'google') + let result: import('./quota.ts').AccountQuotaResult | undefined + try { + const results = await manager.refreshAccounts([account], { + indexFor: () => 0, + force: true, + }) + result = results[0] + } finally { + fetchSpy.mockRestore() + } + + expect(result?.status).toBe('ok') + if (result?.status !== 'ok') return + expect(result.quota?.groups).toBeDefined() + expect(result.geminiCliQuota?.error).toBeTruthy() + }) }) diff --git a/packages/opencode/src/plugin/quota.ts b/packages/opencode/src/plugin/quota.ts index 5fbe0c9..2f9a1d3 100644 --- a/packages/opencode/src/plugin/quota.ts +++ b/packages/opencode/src/plugin/quota.ts @@ -20,14 +20,16 @@ import { type AccountQuotaResult, aggregateGeminiCliQuota, aggregateQuota, + aggregateQuotaSummary, createQuotaManager, defaultKeyOf, type FetchAccountQuota, type FetchAvailableModelsOptions, - type FetchAvailableModelsResponse, fetchAvailableModels, fetchGeminiCliQuota, + fetchQuotaSummary, type GeminiCliQuotaSummary, + getHealthTracker, type QuotaManager, type QuotaSummary, } from '@cortexkit/antigravity-auth-core' @@ -39,6 +41,7 @@ import { } from '../constants' import { buildSidebarMachineStateFromAccounts, + isAccountCurrent, setSidebarMachineState, } from '../sidebar-state' import { @@ -49,7 +52,7 @@ import { import { logQuotaFetch, logQuotaStatus } from './debug' import { buildAntigravityHarnessUserAgent } from './fingerprint' import { createLogger } from './logger' -import { ensureProjectContext } from './project' +import { ensureProjectContext, loadManagedProject } from './project' import { refreshAccessToken } from './token' import type { OAuthAuthDetails, PluginClient } from './types' @@ -114,7 +117,18 @@ export function createOpenCodeQuotaManager( enabled?: boolean coolingDownUntil?: number cachedQuota?: AccountMetadataV3['cachedQuota'] + cachedQuotaAccountId?: string + currentQuotaAccountId?: string }> | null + /** + * Optional provider for the active-account indexes per model family. + * Wired by the plugin entry so every quota-refresh sidebar snapshot + * carries the real `current` flag — not a hardcoded `false`. + */ + getActiveIndexByFamily?: () => { + claude: number + gemini: number + } | null /** * Optional transport adapter used for both `fetchAvailableModels` * and the project-context lookup. When omitted, the production @@ -140,6 +154,7 @@ export function createOpenCodeQuotaManager( const originalRefreshAccount = manager.refreshAccount const originalRefreshAccounts = manager.refreshAccounts const getAccountsForSidebar = options.getAccountsForSidebar + const getActiveIndexByFamily = options.getActiveIndexByFamily let disposed = false const inFlight = new Set>() @@ -150,6 +165,7 @@ export function createOpenCodeQuotaManager( await pushSidebarQuotaSnapshot( getAccountsForSidebar, manager.getBackoffUntil(account), + getActiveIndexByFamily, ).catch(() => { // Sidebar persistence remains best-effort when lock contention // outlives its retry budget. @@ -285,11 +301,20 @@ export async function pushSidebarQuotaSnapshot( enabled?: boolean coolingDownUntil?: number cachedQuota?: AccountMetadataV3['cachedQuota'] + cachedQuotaAccountId?: string + currentQuotaAccountId?: string + /** Captured plan tier to surface in the sidebar state file. */ + tier?: { id: string; paidId?: string; capturedAt: number } }> | null, backoffUntil: number = 0, + getActiveIndexByFamily?: () => { + claude: number + gemini: number + } | null, ): Promise { const accounts = getAccounts() if (!accounts || accounts.length === 0) return + const activeByFamily = getActiveIndexByFamily?.() ?? null try { await setSidebarMachineState( buildSidebarMachineStateFromAccounts( @@ -297,9 +322,15 @@ export async function pushSidebarQuotaSnapshot( index: entry.index, label: entry.label, enabled: entry.enabled, - current: false, + current: activeByFamily + ? isAccountCurrent(entry.index, activeByFamily) + : false, coolingDownUntil: entry.coolingDownUntil, cachedQuota: entry.cachedQuota, + cachedQuotaAccountId: entry.cachedQuotaAccountId, + currentQuotaAccountId: entry.currentQuotaAccountId, + healthScore: getHealthTracker().getScore(entry.index), + tier: entry.tier, })), { checkedAt: Date.now(), @@ -312,6 +343,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, @@ -360,41 +428,90 @@ function makeFetchAccountQuota( const projectContext = await ensureProjectContext(auth) auth = projectContext.auth - const updatedAccount = applyAccountUpdates(account, auth) + const updatedAccount = applyAccountUpdates( + account, + auth, + projectContext.capturedTier, + ) if (rotatedRefresh && client) { await persistRotatedRefresh(client, providerId, auth).catch(() => {}) } - const [antigravityResponse, geminiCliResponse] = await Promise.all([ - fetchAvailableModels({ - accessToken: auth.access ?? '', - 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', + let fellBackToLegacy = false + + const authParts = parseRefreshParts(auth.refresh) + // 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 + + // 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 summaryResult = await fetchQuotaSummary({ + accessToken: auth.access ?? '', + managedProjectId, + projectId: projectContext.effectiveProjectId, + endpoints: ANTIGRAVITY_ENDPOINT_FALLBACKS, + userAgent: buildAntigravityHarnessUserAgent(), + timeoutMs: 10_000, + ...(fetchVia ? { fetchVia } : {}), + }) + return { + result: aggregateQuotaSummary(summaryResult.summary), + fellBackToLegacy: summaryResult.fellBackToLegacy ?? false, + } + } catch { + return { + result: await fetchLegacyModelsFallback({ + accessToken: auth.access ?? '', + projectId: projectContext.effectiveProjectId, + fetchVia, + }), + fellBackToLegacy: true, + } } - } else { - quotaResult = aggregateQuota(antigravityResponse.models) - } + })() + + // CLI fetch is independent of the summary fetch. A CLI failure must + // NOT kill the summary result, but it also must not be laundered into + // "No Gemini CLI quota available" (a permanent-looking status) when + // the real cause is a transient network error. Capture the error + // message separately so the annotated result can carry it. + let geminiCliFetchError: string | undefined + const fetchGeminiCliPayload = fetchGeminiCliQuota({ + accessToken: auth.access ?? '', + projectId: projectContext.effectiveProjectId, + endpoints: ANTIGRAVITY_ENDPOINT_FALLBACKS, + userAgent: buildGeminiCliUserAgent(), + timeoutMs: 10_000, + ...(fetchVia ? { fetchVia } : {}), + }).catch((error: unknown) => { + geminiCliFetchError = + error instanceof Error ? error.message : String(error) + log.debug('fetchGeminiCliQuota failed', { error: geminiCliFetchError }) + 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 = @@ -403,9 +520,12 @@ function makeFetchAccountQuota( ? { ...geminiCliQuotaResult, error: - geminiCliQuotaResult.models.length === 0 + // A real fetch exception is a transient failure, not a + // "no CLI configured" scenario — propagate the actual message. + geminiCliFetchError ?? + (geminiCliQuotaResult.models.length === 0 ? 'No Gemini CLI quota available' - : undefined, + : undefined), } : geminiCliQuotaResult @@ -414,7 +534,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, @@ -458,6 +579,7 @@ function buildAuthFromAccount(account: AccountMetadataV3): OAuthAuthDetails { function applyAccountUpdates( account: AccountMetadataV3, auth: OAuthAuthDetails, + capturedTier?: { id: string; paidId?: string; capturedAt: number }, ): AccountMetadataV3 | undefined { const parts = parseRefreshParts(auth.refresh) if (!parts.refreshToken) { @@ -469,12 +591,30 @@ function applyAccountUpdates( refreshToken: parts.refreshToken, projectId: parts.projectId ?? account.projectId, managedProjectId: parts.managedProjectId ?? account.managedProjectId, + // Persist the captured tier alongside the project-context write. Only + // present when the loadCodeAssist payload returned a non-empty currentTier.id. + ...(capturedTier + ? { + capturedTierId: capturedTier.id, + ...(capturedTier.paidId !== undefined + ? { capturedPaidTierId: capturedTier.paidId } + : {}), + capturedTierAt: capturedTier.capturedAt, + } + : {}), } const changed = updated.refreshToken !== account.refreshToken || updated.projectId !== account.projectId || - updated.managedProjectId !== account.managedProjectId + updated.managedProjectId !== account.managedProjectId || + updated.capturedTierId !== account.capturedTierId || + updated.capturedPaidTierId !== account.capturedPaidTierId || + // capturedAt represents when the tier was LAST CONFIRMED, not when it + // changed -- always update it on a successful observation so consumers + // can gate staleness on that timestamp even when the id stays the same. + (capturedTier !== undefined && + updated.capturedTierAt !== account.capturedTierAt) return changed ? updated : undefined } @@ -494,3 +634,57 @@ async function persistRotatedRefresh( }, }) } + +/** + * Build a per-account tier-loader callback for the background poller. + * + * Calls `loadManagedProject` (loadCodeAssist) directly, bypassing the + * `ensureProjectContext` cache that fast-paths on `managedProjectId` and + * never returns a tier for existing accounts. One call per account per 24 h. + * + * Uses the same token-refresh infrastructure as `makeFetchAccountQuota` so + * an expired access token does not silently fail the tier lookup. + * + * `loadManagedProject` uses the production TLS transport (`fetchWithAgyCliTransport`) + * and is not interceptable via `fetchVia` -- the same design constraint applies to + * `ensureProjectContext`. Tier lookup is best-effort; any failure resolves `null`. + */ +export function makeTierLoader( + client: PluginClient | undefined, + providerId: string, +): ( + account: AccountMetadataV3, +) => Promise<{ id: string; paidId?: string; capturedAt: number } | null> { + return async (account) => { + try { + let auth = buildAuthFromAccount(account) + if (accessTokenExpired(auth)) { + const refreshed = await refreshAccessToken( + auth, + client as PluginClient, + providerId, + ) + if (!refreshed) return null + auth = refreshed + } + + const accessToken = auth.access + if (!accessToken) return null + + const payload = await loadManagedProject(accessToken) + if (!payload?.currentTier?.id) return null + const paidTierId = + typeof payload.paidTier === 'string' + ? payload.paidTier + : payload.paidTier?.id + + return { + id: payload.currentTier.id, + ...(paidTierId ? { paidId: paidTierId } : {}), + capturedAt: Date.now(), + } + } catch { + return null + } + } +} diff --git a/packages/opencode/src/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..0df7c1f 100644 --- a/packages/opencode/src/plugin/ui/quota-status.test.ts +++ b/packages/opencode/src/plugin/ui/quota-status.test.ts @@ -263,24 +263,38 @@ 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({ - 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 +303,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..41141b2 100644 --- a/packages/opencode/src/sidebar-state.test.ts +++ b/packages/opencode/src/sidebar-state.test.ts @@ -24,7 +24,10 @@ import { join } from 'node:path' import { acquireFencedFileLock } from '@cortexkit/antigravity-auth-core' import { + buildSidebarMachineStateFromAccounts, DEFAULT_SIDEBAR_STATE, + isAccountCurrent, + normalizeLegacyCachedQuota, pruneActiveRouting, readSidebarState, redactAccountForSidebar, @@ -105,6 +108,27 @@ async function waitFor( } } +describe('isAccountCurrent', () => { + it('returns true when index matches the claude active index', () => { + expect(isAccountCurrent(0, { claude: 0, gemini: 0 })).toBe(true) + }) + + it('returns true when index matches the gemini active index but not claude', () => { + expect(isAccountCurrent(2, { claude: 0, gemini: 2 })).toBe(true) + }) + + it('returns false when index matches neither family', () => { + expect(isAccountCurrent(1, { claude: 0, gemini: 0 })).toBe(false) + }) + + it('returns true for both indices when families point to different accounts', () => { + const active = { claude: 0, gemini: 2 } + expect(isAccountCurrent(0, active)).toBe(true) + expect(isAccountCurrent(2, active)).toBe(true) + expect(isAccountCurrent(1, active)).toBe(false) + }) +}) + describe('setSidebarMachineState — freshness merge', () => { let fixture: Fixture @@ -444,7 +468,7 @@ describe('redaction', () => { health: 100, current: false, quota: { - claude: { remainingPercent: 80 }, + gemini: { remainingPercent: 80 }, }, }, ], @@ -492,6 +516,19 @@ describe('redaction', () => { } }) + it('defaults healthScore to 100 when missing', () => { + // The doc comment says "Defaults to 100 when missing" — the + // implementation must NOT render 0 just because no producer + // passed the field. + const redacted = redactAccountForSidebar({ index: 0 }) + expect(redacted.health).toBe(100) + }) + + it('passes through an explicit healthScore', () => { + const redacted = redactAccountForSidebar({ index: 0, healthScore: 42 }) + expect(redacted.health).toBe(42) + }) + it('replaces profile labels with privacy-safe ordinal labels', () => { const redacted = redactAccountForSidebar({ index: 0, @@ -502,13 +539,413 @@ 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: { + 'non-gemini': { remainingFraction: 0.42 }, + }, + }) + 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: { - 'gpt-oss': { remainingFraction: 0.42 }, + 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) + }) + + 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', () => { + 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: '5h', + remainingFraction: 0.99, + resetTime: '2026-07-24T20:43:21Z', + }, + { + window: 'weekly', + remainingFraction: 0.92, + resetTime: '2026-07-28T18:24:21Z', + }, + ], + }, + 'non-gemini': { + remainingFraction: 0.96, + resetTime: '2026-07-24T18:41:52Z', + windows: [ + { + window: '5h', + remainingFraction: 0.96, + resetTime: '2026-07-24T18:41:52Z', + }, + { + window: 'weekly', + remainingFraction: 0.99, + resetTime: '2026-07-31T13: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('5h') + expect(gemini!.windows![0]!.remainingPercent).toBe(99) + expect(gemini!.windows![1]!.window).toBe('weekly') + expect(gemini!.windows![1]!.remainingPercent).toBe(92) + + const nonGemini = redacted.quota['non-gemini'] + expect(nonGemini).toBeDefined() + expect(nonGemini!.windows).toHaveLength(2) + expect(nonGemini!.windows![0]!.remainingPercent).toBe(96) + expect(nonGemini!.windows![1]!.remainingPercent).toBe(99) + }) + + it('redactAccountForSidebar normalizes legacy cachedQuota keys and drops windows', () => { + // Uses real legacy keys (gemini-pro) to exercise normalizeLegacyCachedQuota, + // not the canonical 'gemini' key that bypasses the normalization path. + const redacted = redactAccountForSidebar({ + index: 0, + cachedQuota: { + 'gemini-pro': { remainingFraction: 0.5 }, + }, + }) + // Legacy key must be normalized to 'gemini'. + expect(redacted.quota.gemini?.remainingPercent).toBe(50) + // No windows were in the legacy snapshot — none must be invented. + expect(redacted.quota.gemini?.windows).toBeUndefined() + }) + + 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: '5h', + remainingFraction: 0.85, + resetTime: '2026-07-25T00:00:00Z', + }, + { + window: 'weekly', + remainingFraction: 0.7, + resetTime: '2026-08-01T00: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('5h') + expect(gemini!.windows![0]!.remainingPercent).toBe(85) + expect(gemini!.windows![1]!.window).toBe('weekly') + expect(gemini!.windows![1]!.remainingPercent).toBe(70) + } finally { + rmSync(root, { recursive: true, force: true }) + } + }) + + it('normalizeLegacyCachedQuota: gemini-pro + gemini-flash collapse to gemini with MIN fraction', () => { + // M5: verify the two legacy gemini keys collapse with the most-constrained + // (lowest remainingFraction) winning. Also verifies the M4 resetTime fix: + // the winner (gemini-flash, 0.3) carries resetTime=Aug-01; gemini-pro (0.5) + // has an earlier resetTime=Jul-26 — the merged result must use Jul-26. + const result = normalizeLegacyCachedQuota({ + 'gemini-pro': { + remainingFraction: 0.5, + resetTime: '2026-07-26T00:00:00Z', + }, + 'gemini-flash': { + remainingFraction: 0.3, + resetTime: '2026-08-01T00:00:00Z', + }, + }) + expect(result?.gemini?.remainingFraction).toBe(0.3) // gemini-flash wins fraction + // M4 regression: the winner\'s resetTime (Aug-01) is LATER than the loser\'s + // (Jul-26). The merged result must carry the earlier one (Jul-26). + expect(result?.gemini?.resetTime).toBe('2026-07-26T00:00:00Z') + // No windows must be invented for migrated entries. + expect(result?.gemini?.windows).toBeUndefined() + }) + + it('normalizeLegacyCachedQuota: M4 a-wins branch also merges resetTime', () => { + // Specifically tests the a-wins path (fa <= fb): gemini-pro has lower + // fraction than gemini-flash AND an earlier resetTime — both must survive. + const result = normalizeLegacyCachedQuota({ + 'gemini-pro': { + remainingFraction: 0.2, // wins + resetTime: '2026-07-26T00:00:00Z', // earlier + }, + 'gemini-flash': { + remainingFraction: 0.5, + resetTime: '2026-08-05T00:00:00Z', + }, + }) + expect(result?.gemini?.remainingFraction).toBe(0.2) + // Even though a wins, the merged resetTime must be the earlier of the two. + expect(result?.gemini?.resetTime).toBe('2026-07-26T00:00:00Z') + }) + + it('normalizeLegacyCachedQuota: a-wins branch with loser having earlier resetTime (M4 bug scenario)', () => { + // This is the exact M4 bug: a wins the fraction comparison but b has the + // earlier resetTime. Before the fix, b\'s resetTime was silently dropped. + const result = normalizeLegacyCachedQuota({ + 'gemini-pro': { + remainingFraction: 0.3, // wins (lower) + resetTime: '2026-08-01T00:00:00Z', // later + }, + 'gemini-flash': { + remainingFraction: 0.5, + resetTime: '2026-07-26T00:00:00Z', // earlier — must survive + }, + }) + expect(result?.gemini?.remainingFraction).toBe(0.3) + // M4 fix: loser\'s earlier resetTime must win over the winner\'s later one. + expect(result?.gemini?.resetTime).toBe('2026-07-26T00:00:00Z') + }) + + it('normalizeLegacyCachedQuota: claude + gpt-oss collapse to non-gemini with MIN fraction', () => { + const result = normalizeLegacyCachedQuota({ + claude: { remainingFraction: 0.7, resetTime: '2026-08-01T00:00:00Z' }, + 'gpt-oss': { remainingFraction: 0.4, resetTime: '2026-07-28T00:00:00Z' }, + }) + expect(result?.['non-gemini']?.remainingFraction).toBe(0.4) // gpt-oss wins + expect(result?.['non-gemini']?.resetTime).toBe('2026-07-28T00:00:00Z') // earlier + }) + + it('normalizeLegacyCachedQuota: no windows invented for migrated entries', () => { + // Legacy entries never had windows; the normalizer must not synthesize them. + const result = normalizeLegacyCachedQuota({ + 'gemini-pro': { remainingFraction: 0.6 }, + claude: { remainingFraction: 0.8 }, + }) + expect(result?.gemini?.windows).toBeUndefined() + expect(result?.['non-gemini']?.windows).toBeUndefined() + }) + + it('normalizeLegacyCachedQuota: canonical keys pass through unchanged', () => { + // Non-legacy input must be returned as-is without any transformation. + const canonical = { + gemini: { remainingFraction: 0.6, resetTime: '2026-08-01T00:00:00Z' }, + 'non-gemini': { remainingFraction: 0.4 }, + } + const result = normalizeLegacyCachedQuota(canonical) + expect(result).toBe(canonical) // same reference — fast path returned as-is + }) + + it('identity mismatch drops the entire cachedQuota including windows', () => { + const redacted = redactAccountForSidebar({ + index: 0, + 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) + }) +}) + +describe('render boundary — current flag', () => { + let fixture: Fixture + + beforeEach(() => { + fixture = makeFixture() + }) + + afterEach(() => { + fixture.cleanup() + }) + + it('preserves current: true from producer input through to sidebar state', async () => { + const state = buildSidebarMachineStateFromAccounts([ + { index: 0, current: true }, + { index: 1, current: false }, + ]) + // The render boundary is `redactAccountForSidebar` inside + // `buildSidebarMachineStateFromAccounts`. A `current: true` from any + // producer must survive to the state the TUI reads so that + // `visibleAccounts` (when fallbackAccounts is off) renders the + // active account instead of an empty list. + expect(state.accounts).toHaveLength(2) + expect(state.accounts[0]?.current).toBe(true) + expect(state.accounts[1]?.current).toBe(false) + + // Also exercise the full write→read round-trip the TUI actually uses. + await setSidebarMachineState(state, { stateFile: fixture.stateFile }) + const disk = readSidebarState(fixture.stateFile) + expect(disk.accounts[0]?.current).toBe(true) + expect(disk.accounts[1]?.current).toBe(false) + }) + + it('without fallbackAccounts, current: false for all accounts → no accounts visible', () => { + // The filter the TUI runs is: + // state.accounts.filter(account => account.current) + // When every account has current: false, the result is empty — that + // was the live bug before this fix. + const state = buildSidebarMachineStateFromAccounts([ + { index: 0, current: false }, + { index: 1, current: false }, + ]) + const visible = state.accounts.filter((entry) => entry.current) + expect(visible).toHaveLength(0) + }) + + it('without fallbackAccounts, current: true on one account → that account is visible', () => { + const state = buildSidebarMachineStateFromAccounts([ + { index: 0, current: true }, + { index: 1, current: false }, + ]) + const visible = state.accounts.filter((entry) => entry.current) + expect(visible).toHaveLength(1) + expect(visible[0]?.id).toBe('acct-0') + }) + + // ── Fix 2: tier round-trip ────────────────────────────────────────────────── + + it('Fix2-tier-roundtrip: tier writes to disk and reads back with capturedAt', async () => { + // An account with a captured tier must survive a full write→read cycle. + const capturedAt = 1_785_005_949_000 + const state = buildSidebarMachineStateFromAccounts([ + { index: 0, tier: { id: 'free-tier', capturedAt } }, + ]) + await setSidebarMachineState(state, { stateFile: fixture.stateFile }) + const disk = readSidebarState(fixture.stateFile) + expect(disk.accounts[0]?.tier).toEqual({ id: 'free-tier', capturedAt }) + }) + + it('Fix2-tier-absent: account without tier yields no tier key', () => { + // An account with no tier in the redaction input must not emit tier in the + // rendered state — absent ≠ free, absent ≠ {} . + const state = buildSidebarMachineStateFromAccounts([{ index: 0 }]) + expect(state.accounts[0]?.tier).toBeUndefined() + // Confirm `tier` is not a key at all in serialized form. + const serialized = JSON.stringify(state.accounts[0]) + expect(serialized).not.toContain('tier') + }) + + it('Fix2-tier-redaction: tier survives redactAccountForSidebar (not PII)', () => { + // Plan tier is metadata, not a credential — it must pass the redaction boundary. + const capturedAt = 1_785_000_000_000 + const result = redactAccountForSidebar({ + index: 0, + tier: { id: 'paid-tier', capturedAt }, }) - expect(redacted.quota['gpt-oss']?.remainingPercent).toBe(42) + expect(result.tier).toEqual({ id: 'paid-tier', capturedAt }) }) }) diff --git a/packages/opencode/src/sidebar-state.ts b/packages/opencode/src/sidebar-state.ts index 7cce03d..2dbcdfc 100644 --- a/packages/opencode/src/sidebar-state.ts +++ b/packages/opencode/src/sidebar-state.ts @@ -52,15 +52,20 @@ 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 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 { @@ -71,6 +76,11 @@ export interface SidebarAccountState { current: boolean cooldownUntil?: number quota: Partial> + /** + * Captured plan tier. Absent when unknown — do NOT default to `free-tier`. + * `id` is the raw upstream string (e.g. `"free-tier"`); never normalised. + */ + tier?: { id: string; paidId?: string; capturedAt: number } } export interface SidebarRoutingEntry { @@ -294,17 +304,13 @@ 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 } } + const tier = normalizeTier(record.tier) return { id, label, @@ -313,18 +319,60 @@ function normalizeAccount(input: unknown): SidebarAccountState | null { current, cooldownUntil, quota, + ...(tier !== undefined ? { tier } : {}), } } +function normalizeTier( + input: unknown, +): { id: string; paidId?: string; capturedAt: number } | undefined { + if (!isObject(input)) return undefined + const record = input as Record + const id = + typeof record.id === 'string' && record.id.length > 0 ? record.id : null + const capturedAt = toFiniteNumber(record.capturedAt) + if (!id || capturedAt === null) return undefined + const paidId = + typeof record.paidId === 'string' && record.paidId.length > 0 + ? record.paidId + : undefined + return { id, ...(paidId ? { paidId } : {}), capturedAt } +} + function normalizeQuota(input: unknown): SidebarQuotaEntry | null { if (!isObject(input)) return null const record = input as Record 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, } } @@ -441,13 +489,236 @@ 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 + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } + 'non-gemini'?: { + remainingFraction?: number + resetTime?: string + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } + // Legacy per-pool keys written by older versions of the plugin. + // Accepted here so the normalizer can map them to the current pool + // keys at read time without requiring a disk migration. + [legacyKey: string]: + | { + remainingFraction?: number + resetTime?: string + modelCount?: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } + | undefined + } + /** + * Opaque identity stamp that was attached to the persisted quota snapshot. + * 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 + /** + * Captured plan tier from `loadCodeAssist`. Absent when unknown. The `id` + * is the raw upstream string; `capturedAt` is epoch ms. Not PII — tier + * metadata survives the redaction boundary unchanged. + */ + tier?: { id: string; paidId?: string; capturedAt: number } +} + +/** + * Build a `{ id, capturedAt }` tier object from account fields, or `undefined` + * when either field is absent. Centralised so every producer (index.ts, + * background-quota-refresh.ts, command-data.ts) uses the same guard and the + * same shape -- the same pattern shipped as inline literals across three files + * and produced six missed-field bugs; one helper ends that. + */ +export function toCapturedTier(account: { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number +}): { id: string; paidId?: string; capturedAt: number } | undefined { + return account.capturedTierId !== undefined && + account.capturedTierAt !== undefined + ? { + id: account.capturedTierId, + ...(account.capturedPaidTierId !== undefined + ? { paidId: account.capturedPaidTierId } + : {}), + capturedAt: account.capturedTierAt, + } + : undefined +} + +/** + * Map pre-pool legacy quota keys to the current two-pool schema at read time. + * + * Legacy mappings (empirically settled from burn tests): + * `gemini-pro` + `gemini-flash` → `gemini` + * `claude` + `gpt-oss` → `non-gemini` + * + * Where multiple legacy keys collapse into one pool the MIN remainingFraction + * and earliest resetTime are used (most-constrained-first, matching the + * window-level rule). No `windows` arrays are invented for migrated entries; + * the single aggregate bar is the correct render for pre-window snapshots. + * + * Non-legacy keys that are already canonical (`gemini`, `non-gemini`) are + * carried through unchanged. The next real quota refresh overwrites any + * migrated snapshot with authoritative data. + */ +export function normalizeLegacyCachedQuota( + raw: SidebarAccountRedactionInput['cachedQuota'], +): SidebarAccountRedactionInput['cachedQuota'] { + if (!raw) return raw + + // Pool already carries current keys — fast path when no legacy keys present. + const hasLegacy = + 'gemini-pro' in raw || + 'gemini-flash' in raw || + 'claude' in raw || + 'gpt-oss' in raw + if (!hasLegacy) return raw + + type PoolEntry = NonNullable< + SidebarAccountRedactionInput['cachedQuota'] + >['gemini'] + + // Pick the earlier of two ISO reset timestamps, preferring a defined + // value over undefined. Extracted so both branches of minFraction use + // the same merge rather than the previous asymmetry where the a-wins + // branch dropped b's possibly-earlier resetTime. + const earlierResetTime = ( + a: string | undefined, + b: string | undefined, + ): string | undefined => { + if (!a) return b + if (!b) return a + return a < b ? a : b + } + + const minFraction = (a: PoolEntry, b: PoolEntry): PoolEntry => { + if (!a) return b + if (!b) return a + const fa = a.remainingFraction ?? 1 + const fb = b.remainingFraction ?? 1 + const winner = fa <= fb ? a : b + const loser = fa <= fb ? b : a + // Always merge to the earliest resetTime regardless of which pool wins + // the fraction comparison — an earlier window expiry from the loser + // pool must not be silently dropped. + return { + ...winner, + resetTime: earlierResetTime(winner.resetTime, loser.resetTime), + } + } + + const gemini: PoolEntry = + minFraction( + raw.gemini, + minFraction( + raw['gemini-pro'] as PoolEntry, + raw['gemini-flash'] as PoolEntry, + ), + ) ?? raw.gemini + + const nonGemini: PoolEntry = + minFraction( + raw['non-gemini'], + minFraction(raw.claude as PoolEntry, raw['gpt-oss'] as PoolEntry), + ) ?? raw['non-gemini'] + + return { + ...(gemini !== undefined ? { gemini } : {}), + ...(nonGemini !== undefined ? { 'non-gemini': nonGemini } : {}), } } +/** + * 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 } +} + +/** + * Returns `true` when an account at `index` is the active account for + * at least one model family. Two accounts CAN both be current — one + * serving claude, one serving gemini — so this must not collapse to one. + */ +export function isAccountCurrent( + index: number, + activeIndexByFamily: { claude: number; gemini: number }, +): boolean { + return ( + index === activeIndexByFamily.claude || index === activeIndexByFamily.gemini + ) +} + /** * Convert a live account snapshot into the redacted shape the TUI renders. * The redacted `SidebarAccountState` carries no email, refresh token, access @@ -467,31 +738,32 @@ export function redactAccountForSidebar( ? source.coolingDownUntil : undefined const health = clampNumber( - typeof source.healthScore === 'number' ? source.healthScore : null, + typeof source.healthScore === 'number' ? source.healthScore : 100, 0, 100, ) 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 + // Normalize legacy per-pool keys to the current two-pool schema before + // reading. This is non-destructive: the next real quota refresh will + // overwrite the migrated snapshot with authoritative data. + const cached = staleCachedQuota + ? undefined + : normalizeLegacyCachedQuota(source.cachedQuota) if (cached) { - for (const key of [ - '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 - 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 - } - quota[key] = { remainingPercent, resetAt } + const projected = projectQuotaPoolForSidebar(entry) + if (projected) quota[key] = projected } } @@ -503,6 +775,8 @@ export function redactAccountForSidebar( current, cooldownUntil, quota, + // Tier is plan metadata, not PII — it passes the redaction boundary. + ...(source.tier !== undefined ? { tier: source.tier } : {}), } } diff --git a/packages/opencode/src/tui-compiled/plugin/command-data.ts b/packages/opencode/src/tui-compiled/plugin/command-data.ts new file mode 100644 index 0000000..eb72733 --- /dev/null +++ b/packages/opencode/src/tui-compiled/plugin/command-data.ts @@ -0,0 +1,931 @@ +/** + * 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 { createHash } from 'node:crypto' + +import { + buildSidebarMachineStateFromAccounts, + normalizeLegacyCachedQuota, + type SidebarAccountRedactionInput, + setSidebarMachineState, + toCapturedTier, +} 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 + cachedQuotaAccountId?: string + accountIneligible?: boolean +} + +type CommandDataQuotaGroup = 'gemini' | 'non-gemini' + +type CommandDataQuotaGroupSummary = { + remainingFraction?: number + resetTime?: string + modelCount: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> +} + +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 + /** Privacy-safe ordinal display label (for example, `Account 1`). */ + label: string + 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 + /** Health score in [0, 100], forwarded from the live tracker. */ + healthScore?: number + quota: Array<{ + key: 'gemini' | 'non-gemini' + label: string + remainingPercent: number | null + resetAt?: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingPercent: number + resetAt?: number + }> + }> + /** Captured plan tier. Absent when unknown; never defaulted. */ + tier?: { id: string; capturedAt: 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 +> = { + gemini: 'Gemini', + 'non-gemini': 'Non-Gemini', +} + +const SUPPORTED_QUOTA_KEYS = [ + 'gemini', + 'non-gemini', +] 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 + cachedQuotaAccountId?: string + accountIneligible?: boolean + coolingDownUntil?: number + healthScore?: number + /** Captured plan tier from the most recent loadCodeAssist response. */ + capturedTierId?: string + /** Captured paid-tier ID from the most recent loadCodeAssist response. */ + capturedPaidTierId?: string + /** Epoch ms when capturedTierId was last recorded. */ + capturedTierAt?: number +} + +function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow { + // 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 rawCached = + entry.cachedQuotaAccountId && + entry.cachedQuotaAccountId !== quotaAccountIdentity(entry.refreshToken) + ? undefined + : entry.cachedQuota + // Normalize legacy pool keys (gemini-pro, gemini-flash, claude, gpt-oss) + // to canonical keys before the SUPPORTED_QUOTA_KEYS loop. Without this + // a legacy snapshot produces empty quota rows, and writeSidebar would + // permanently re-emit legacy keys on every dialog open until the next + // real quota refresh rewrites the on-disk snapshot. + const cached = normalizeLegacyCachedQuota( + rawCached as Parameters[0], + ) + const quota: CommandAccountRow['quota'] = [] + for (const key of SUPPORTED_QUOTA_KEYS) { + const cachedEntry = cached?.[key] + 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 + } + // 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}` + return { + id: `acct-${entry.index}`, + index: entry.index, + label, + enabled: entry.enabled, + current: entry.active, + coolingDownUntil: entry.coolingDownUntil, + healthScore: entry.healthScore, + quota, + tier: toCapturedTier(entry), + } +} + +/** + * 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[] { + 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, + cachedQuotaAccountId: entry.cachedQuotaAccountId, + }), + ) +} + +/** + * 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. + */ +export interface CommandDataAccountManagerView { + getAccounts(): LiveAccountSnapshot[] + getAccountsForQuotaCheck(): CommandDataAccountMetadata[] + updateQuotaCache( + index: number, + groups: Partial< + Record + >, + expectedRefreshToken?: string, + ): void + requestSaveToDisk(): void + flushSaveToDisk(): Promise + /** + * 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. */ + 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. + * + * Bypass flag is intentional: explicit user-triggered refreshes + * ("Refresh" button, direct slash-command with `refresh` argument) + * must always fetch, even if the quota manager has backed off. + */ + refreshQuota(): Promise + /** + * Non-forced quota check that RESPECTS the quota manager's per-account + * backoff. Used by the dialog OPEN path — which is a VIEW, not a user + * request to fetch. Accounts already fresh (or in backoff) are skipped; + * the result is discarded (the sidebar poller carries the freshness + * guarantee for idle accounts). + */ + refreshQuotaRespectingBackoff(): Promise + /** + * Pin `index` as the active account for every family the dialog + * tracks (claude + gemini). Mutates the live AccountManager AND the + * 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) + + // Map a CommandAccountRow back to the SidebarAccountRedactionInput shape. + // Extracted to prevent the "sixth dropped field" pattern: every field in + // CommandAccountRow that must reach the sidebar lives here once, not as an + // inline literal that a future author has to remember to update in two places. + const toRedactionInput = ( + row: CommandAccountRow, + ): SidebarAccountRedactionInput => { + const gemini = row.quota.find((q) => q.key === 'gemini') + const nonGemini = row.quota.find((q) => q.key === 'non-gemini') + // Map a CommandAccountRow quota entry to cachedQuota pool shape + // so projectQuotaPoolForSidebar (the canonical projection seam) + // carries the per-window breakdown into the sidebar state. + const toPool = ( + q: + | { + remainingPercent: number | null + resetAt?: number + windows?: CommandAccountRow['quota'][number]['windows'] + } + | undefined, + ) => { + if (!q || q.remainingPercent == null) return undefined + return { + remainingFraction: q.remainingPercent / 100, + resetTime: + typeof q.resetAt === 'number' && Number.isFinite(q.resetAt) + ? new Date(q.resetAt).toISOString() + : undefined, + windows: q.windows?.map((w) => ({ + window: w.window, + remainingFraction: w.remainingPercent / 100, + resetTime: + typeof w.resetAt === 'number' && Number.isFinite(w.resetAt) + ? new Date(w.resetAt).toISOString() + : '', + })), + } + } + return { + index: row.index, + label: row.label, + enabled: row.enabled, + current: row.current, + healthScore: row.healthScore, + // Tier passes the redaction boundary unchanged (plan metadata, not PII). + // The stamp check was already done by toCommandAccountRow. + tier: row.tier, + cachedQuota: { + gemini: toPool(gemini), + 'non-gemini': toPool(nonGemini), + }, + } + } + + const writeSidebar = (rows: CommandAccountRow[]): void => { + const accounts = rows.map(toRedactionInput) + // Fire-and-forget -- the sidebar writer is fenced by its own queue, + // so a transient lock contention cannot block the dialog response. + void setSidebarMachineState( + buildSidebarMachineStateFromAccounts(accounts, { checkedAt: now() }), + { 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, + }) + + const refreshedAt = now() + const updates: Array<{ + refreshToken: string + groups?: Partial< + Record + > + }> = [] + for (const result of results) { + const refreshToken = + result.updatedAccount?.refreshToken ?? + accountsForQuota[result.index]?.refreshToken + if (!refreshToken) continue + const groups = + result.status === 'ok' && result.quota?.groups + ? result.quota.groups + : undefined + updates.push({ refreshToken, groups }) + } + + // 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 + } + if (liveQuotaChanged) accountManagerView.requestSaveToDisk() + + // 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 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, + } + } + // Error result keeps the previous cached percentage and only + // records that a refresh was attempted. + return { + ...entry, + cachedQuotaUpdatedAt: refreshedAt, + } + }), + })) + await Promise.resolve(writeResult).catch(() => { + // The live AccountManager already carries successful refreshes; its + // next save can reconcile a transient storage-lock failure. + }) + } + + // Re-read the live view post-mutate so the rows reflect the + // freshly persisted percentages. + const rows = projectRows() + writeSidebar(rows) + return rows + }, + + async refreshQuotaRespectingBackoff() { + const accountsForQuota = accountManagerView.getAccountsForQuotaCheck() + if (accountsForQuota.length === 0) return + + // Non-forced: the quota manager's per-account backoff and dedup + // apply. Accounts that are fresh or in backoff are skipped. Unlike + // the earlier fire-and-forget path, successful results ARE folded + // into the live cache so the sidebar snapshot reflects fresh data + // immediately — the dialog renders ONLY from the sidebar file, and + // the quota manager's push path does not update AccountManager + // entries for this caller. + let results: Awaited> + try { + results = await quotaManager.refreshAccounts(accountsForQuota, { + indexFor: (account) => accountsForQuota.indexOf(account), + force: false, + }) + } catch { + // Non-critical: dialog open must never fail because a background + // refresh check encountered backoff or a transient error. + return + } + + // Build updates akin to the forced refresh path — resolve live + // indexes after the network call (a concurrent add/remove may have + // shifted positions), fold successful results into the live cache, + // then push a sidebar snapshot. + const updates: Array<{ + refreshToken: string + groups?: Partial< + Record + > + }> = [] + for (const result of results) { + const refreshToken = + result.updatedAccount?.refreshToken ?? + accountsForQuota[result.index]?.refreshToken + if (!refreshToken) continue + const groups = + result.status === 'ok' && result.quota?.groups + ? result.quota.groups + : undefined + updates.push({ refreshToken, groups }) + } + + if (updates.length === 0) return + + const liveIndexByRefreshToken = new Map() + for (const entry of accountManagerView.getAccounts()) { + liveIndexByRefreshToken.set(entry.refreshToken, entry.index) + } + let liveQuotaChanged = false + for (const update of updates) { + const liveIndex = liveIndexByRefreshToken.get(update.refreshToken) + if (liveIndex === undefined || !update.groups) continue + accountManagerView.updateQuotaCache( + liveIndex, + update.groups, + update.refreshToken, + ) + liveQuotaChanged = true + } + if (liveQuotaChanged) accountManagerView.requestSaveToDisk() + + // Push a fresh sidebar snapshot — the live view now carries the + // just-refreshed quota plus any skipped (backoff/fresh) accounts + // whose cached percentages are unchanged. + const rows = projectRows() + writeSidebar(rows) + }, + + async setCurrentAccount(index) { + return mutateLiveAndStorage({ action: 'setCurrent', index }) + }, + + async toggleAccountEnabled(index) { + return mutateLiveAndStorage({ action: 'toggleEnabled', index }) + }, + + async removeAccount(index) { + return mutateLiveAndStorage({ action: 'remove', index }) + }, + } + + async function mutateLiveAndStorage(args: { + action: 'setCurrent' | 'toggleEnabled' | 'remove' + index: number + }): Promise { + 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) { + throw new Error( + 'CommandDataService is missing a locked-storage adapter; account mutations are disabled.', + ) + } + + // 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. 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 liveIndexes = accountManagerView.getActiveIndexByFamily() + liveCurrentTokens.claude = + liveAccounts[liveIndexes.claude]?.refreshToken ?? undefined + liveCurrentTokens.gemini = + liveAccounts[liveIndexes.gemini]?.refreshToken ?? undefined + } + + 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, + ) + if (tokenIdx === -1) return current + foundInStorage = true + + if (action === 'setCurrent') { + return { + ...current, + activeIndex: tokenIdx, + activeIndexByFamily: { claude: tokenIdx, gemini: tokenIdx }, + } + } + + if (action === 'toggleEnabled') { + const entry = current.accounts[tokenIdx] + if (!entry) return current + 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((account) => + account.refreshToken === refreshToken + ? { ...account, enabled: desiredEnabled } + : account, + ), + } + } + + // 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 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, + ) + 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 found + // 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). + // + // Use tokenIdx (the storage-snapshot position) rather than the + // outer-scope `index` (live-read position): a concurrent mutation + // between the live read and the lock can diverge the two, and the + // storage-level cursor must match the storage-level array. + return tokenIdx < nextAccounts.length ? tokenIdx : 0 + } + const legacyClaude = + current.activeIndexByFamily?.claude ?? current.activeIndex + 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: nextAccounts, + activeIndex: nextActiveIndex, + activeIndexByFamily: nextActiveIndexByFamily, + } + }) + + 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(() => { + // The locked storage mutation already committed; a later periodic flush + // can reconcile transient AccountManager lock contention. + }) + + const rows = projectRows() + writeSidebar(rows) + return rows + } +} diff --git a/packages/opencode/src/tui-compiled/sidebar-state.ts b/packages/opencode/src/tui-compiled/sidebar-state.ts index 7cce03d..2dbcdfc 100644 --- a/packages/opencode/src/tui-compiled/sidebar-state.ts +++ b/packages/opencode/src/tui-compiled/sidebar-state.ts @@ -52,15 +52,20 @@ 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 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 { @@ -71,6 +76,11 @@ export interface SidebarAccountState { current: boolean cooldownUntil?: number quota: Partial> + /** + * Captured plan tier. Absent when unknown — do NOT default to `free-tier`. + * `id` is the raw upstream string (e.g. `"free-tier"`); never normalised. + */ + tier?: { id: string; paidId?: string; capturedAt: number } } export interface SidebarRoutingEntry { @@ -294,17 +304,13 @@ 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 } } + const tier = normalizeTier(record.tier) return { id, label, @@ -313,18 +319,60 @@ function normalizeAccount(input: unknown): SidebarAccountState | null { current, cooldownUntil, quota, + ...(tier !== undefined ? { tier } : {}), } } +function normalizeTier( + input: unknown, +): { id: string; paidId?: string; capturedAt: number } | undefined { + if (!isObject(input)) return undefined + const record = input as Record + const id = + typeof record.id === 'string' && record.id.length > 0 ? record.id : null + const capturedAt = toFiniteNumber(record.capturedAt) + if (!id || capturedAt === null) return undefined + const paidId = + typeof record.paidId === 'string' && record.paidId.length > 0 + ? record.paidId + : undefined + return { id, ...(paidId ? { paidId } : {}), capturedAt } +} + function normalizeQuota(input: unknown): SidebarQuotaEntry | null { if (!isObject(input)) return null const record = input as Record 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, } } @@ -441,13 +489,236 @@ 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 + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } + 'non-gemini'?: { + remainingFraction?: number + resetTime?: string + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } + // Legacy per-pool keys written by older versions of the plugin. + // Accepted here so the normalizer can map them to the current pool + // keys at read time without requiring a disk migration. + [legacyKey: string]: + | { + remainingFraction?: number + resetTime?: string + modelCount?: number + windows?: Array<{ + window: 'weekly' | '5h' + remainingFraction: number + resetTime: string + }> + } + | undefined + } + /** + * Opaque identity stamp that was attached to the persisted quota snapshot. + * 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 + /** + * Captured plan tier from `loadCodeAssist`. Absent when unknown. The `id` + * is the raw upstream string; `capturedAt` is epoch ms. Not PII — tier + * metadata survives the redaction boundary unchanged. + */ + tier?: { id: string; paidId?: string; capturedAt: number } +} + +/** + * Build a `{ id, capturedAt }` tier object from account fields, or `undefined` + * when either field is absent. Centralised so every producer (index.ts, + * background-quota-refresh.ts, command-data.ts) uses the same guard and the + * same shape -- the same pattern shipped as inline literals across three files + * and produced six missed-field bugs; one helper ends that. + */ +export function toCapturedTier(account: { + capturedTierId?: string + capturedPaidTierId?: string + capturedTierAt?: number +}): { id: string; paidId?: string; capturedAt: number } | undefined { + return account.capturedTierId !== undefined && + account.capturedTierAt !== undefined + ? { + id: account.capturedTierId, + ...(account.capturedPaidTierId !== undefined + ? { paidId: account.capturedPaidTierId } + : {}), + capturedAt: account.capturedTierAt, + } + : undefined +} + +/** + * Map pre-pool legacy quota keys to the current two-pool schema at read time. + * + * Legacy mappings (empirically settled from burn tests): + * `gemini-pro` + `gemini-flash` → `gemini` + * `claude` + `gpt-oss` → `non-gemini` + * + * Where multiple legacy keys collapse into one pool the MIN remainingFraction + * and earliest resetTime are used (most-constrained-first, matching the + * window-level rule). No `windows` arrays are invented for migrated entries; + * the single aggregate bar is the correct render for pre-window snapshots. + * + * Non-legacy keys that are already canonical (`gemini`, `non-gemini`) are + * carried through unchanged. The next real quota refresh overwrites any + * migrated snapshot with authoritative data. + */ +export function normalizeLegacyCachedQuota( + raw: SidebarAccountRedactionInput['cachedQuota'], +): SidebarAccountRedactionInput['cachedQuota'] { + if (!raw) return raw + + // Pool already carries current keys — fast path when no legacy keys present. + const hasLegacy = + 'gemini-pro' in raw || + 'gemini-flash' in raw || + 'claude' in raw || + 'gpt-oss' in raw + if (!hasLegacy) return raw + + type PoolEntry = NonNullable< + SidebarAccountRedactionInput['cachedQuota'] + >['gemini'] + + // Pick the earlier of two ISO reset timestamps, preferring a defined + // value over undefined. Extracted so both branches of minFraction use + // the same merge rather than the previous asymmetry where the a-wins + // branch dropped b's possibly-earlier resetTime. + const earlierResetTime = ( + a: string | undefined, + b: string | undefined, + ): string | undefined => { + if (!a) return b + if (!b) return a + return a < b ? a : b + } + + const minFraction = (a: PoolEntry, b: PoolEntry): PoolEntry => { + if (!a) return b + if (!b) return a + const fa = a.remainingFraction ?? 1 + const fb = b.remainingFraction ?? 1 + const winner = fa <= fb ? a : b + const loser = fa <= fb ? b : a + // Always merge to the earliest resetTime regardless of which pool wins + // the fraction comparison — an earlier window expiry from the loser + // pool must not be silently dropped. + return { + ...winner, + resetTime: earlierResetTime(winner.resetTime, loser.resetTime), + } + } + + const gemini: PoolEntry = + minFraction( + raw.gemini, + minFraction( + raw['gemini-pro'] as PoolEntry, + raw['gemini-flash'] as PoolEntry, + ), + ) ?? raw.gemini + + const nonGemini: PoolEntry = + minFraction( + raw['non-gemini'], + minFraction(raw.claude as PoolEntry, raw['gpt-oss'] as PoolEntry), + ) ?? raw['non-gemini'] + + return { + ...(gemini !== undefined ? { gemini } : {}), + ...(nonGemini !== undefined ? { 'non-gemini': nonGemini } : {}), } } +/** + * 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 } +} + +/** + * Returns `true` when an account at `index` is the active account for + * at least one model family. Two accounts CAN both be current — one + * serving claude, one serving gemini — so this must not collapse to one. + */ +export function isAccountCurrent( + index: number, + activeIndexByFamily: { claude: number; gemini: number }, +): boolean { + return ( + index === activeIndexByFamily.claude || index === activeIndexByFamily.gemini + ) +} + /** * Convert a live account snapshot into the redacted shape the TUI renders. * The redacted `SidebarAccountState` carries no email, refresh token, access @@ -467,31 +738,32 @@ export function redactAccountForSidebar( ? source.coolingDownUntil : undefined const health = clampNumber( - typeof source.healthScore === 'number' ? source.healthScore : null, + typeof source.healthScore === 'number' ? source.healthScore : 100, 0, 100, ) 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 + // Normalize legacy per-pool keys to the current two-pool schema before + // reading. This is non-destructive: the next real quota refresh will + // overwrite the migrated snapshot with authoritative data. + const cached = staleCachedQuota + ? undefined + : normalizeLegacyCachedQuota(source.cachedQuota) if (cached) { - for (const key of [ - '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 - 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 - } - quota[key] = { remainingPercent, resetAt } + const projected = projectQuotaPoolForSidebar(entry) + if (projected) quota[key] = projected } } @@ -503,6 +775,8 @@ export function redactAccountForSidebar( current, cooldownUntil, quota, + // Tier is plan metadata, not PII — it passes the redaction boundary. + ...(source.tier !== undefined ? { tier: source.tier } : {}), } } diff --git a/packages/opencode/src/tui-compiled/tui.tsx b/packages/opencode/src/tui-compiled/tui.tsx index d3adc16..5e06ed5 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,14 @@ 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 = ['gemini', 'non-gemini']; +const WINDOW_GUTTER = { + weekly: '7d', + '5h': '5h' }; -const QUOTA_ORDER = ['claude', 'gemini-pro', 'gemini-flash', 'gpt-oss']; // --- Theme tokens ---------------------------------------------------------- // @@ -440,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, @@ -469,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); @@ -572,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 => _$createComponent(QuotaRow, { + get theme() { + return props.theme; + }, + get appearance() { + return props.appearance; + }, + get label() { + return `${QUOTA_LABELS[key]} ${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%'); @@ -837,8 +865,23 @@ export function SidebarPanel(props) { return (() => { const account = () => visibleAccounts().find(entry => entry.current) ?? visibleAccounts()[0]; const used = () => { - const entry = account()?.quota.claude; - return entry ? 100 - clamp(entry.remainingPercent, 0, 100) : null; + const q = account()?.quota; + // Derive tone from the WORST (most-constrained) pool so a + // critical non-gemini pool doesn't hide behind a healthy + // gemini tone. Worst = highest used-percent across available pools. + const geminiUsed = q?.gemini != null ? 100 - clamp(q.gemini.remainingPercent, 0, 100) : null; + const nonGeminiUsed = q?.['non-gemini'] != null ? 100 - clamp(q['non-gemini'].remainingPercent, 0, 100) : null; + if (geminiUsed === null && nonGeminiUsed === null) return null; + return Math.max(geminiUsed ?? 0, nonGeminiUsed ?? 0); + }; + const poolText = key => { + const entry = account()?.quota[key]; + if (!entry) return `${QUOTA_LABELS[key]}: —`; + const pct = Math.round(100 - clamp(entry.remainingPercent, 0, 100)); + // 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(); @@ -863,8 +906,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$() ? '—' : `Cl: ${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-compiled/tui/command-dialogs.tsx b/packages/opencode/src/tui-compiled/tui/command-dialogs.tsx index dca8198..0031ca3 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,153 @@ 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) { + // 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); + }, + 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 +423,7 @@ function renderAccountDialog(api, initialPayload, apply) { }); } })); - return _el$; + return _el$5; })()); }; const promptRemove = index => { @@ -335,19 +461,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 +496,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 +510,7 @@ function renderAccountDialog(api, initialPayload, apply) { }); } })); - return _el$5; + return _el$9; })()); }; renderMain(); @@ -452,19 +572,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 +618,7 @@ function renderRoutingDialog(api, initialPayload, apply) { }); } })); - return _el$8; + return _el$10; })()); }; renderMain(); @@ -559,9 +679,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 +722,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 +772,7 @@ function renderKillswitchDialog(api, initialPayload, apply) { }); } })); - return _el$11; + return _el$15; })()); }; renderMain(); 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..5b530a2 --- /dev/null +++ b/packages/opencode/src/tui-windows-frames.test.tsx @@ -0,0 +1,371 @@ +/** @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') + // Save preload-pinned values to restore on cleanup instead of deleting — + // a delete drops resolution to the operator's real state/config dirs. + const savedSidebar = process.env[SIDEBAR_STATE_ENV] + const savedTuiLog = process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE + const savedPrefs = process.env[TUI_PREFS_FILE_ENV] + process.env[SIDEBAR_STATE_ENV] = statePath + process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE = join(root, 'tui.log') + process.env[TUI_PREFS_FILE_ENV] = prefsPath + return { + statePath, + prefsPath, + cleanup: () => { + if (savedSidebar !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebar + else delete process.env[SIDEBAR_STATE_ENV] + if (savedTuiLog !== undefined) + process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE = savedTuiLog + else delete process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE + if (savedPrefs !== undefined) process.env[TUI_PREFS_FILE_ENV] = savedPrefs + else delete process.env[TUI_PREFS_FILE_ENV] + rmSync(root, { recursive: true, force: true }) + }, + } +} + +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: '5h', remainingPercent: 99, resetAt: future }, + { window: 'weekly', remainingPercent: 92, resetAt: future2 }, + ], + }, + 'non-gemini': { + remainingPercent: 96, + resetAt: future2, + windows: [ + { window: '5h', remainingPercent: 96, resetAt: future }, + { window: 'weekly', remainingPercent: 99, resetAt: future2 }, + ], + }, + }, + }, + ], + }) + 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') + // Window order: aggregator now sorts shortest-first so 5h must + // render before 7d. The fixture carries windows in [5h, weekly] + // order (matching aggregateQuotaSummary output) and the renderer + // preserves input order, so the assertion is deterministic. + expect(frame.indexOf('Gm 5h')).toBeLessThan(frame.indexOf('Gm 7d')) + expect(frame.indexOf('NG 5h')).toBeLessThan(frame.indexOf('NG 7d')) + testSetup.renderer.destroy() + }) + + 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: '5h', remainingPercent: 99, resetAt: future }, + { window: 'weekly', remainingPercent: 92, resetAt: future2 }, + ], + }, + 'non-gemini': { + remainingPercent: 96, + resetAt: future2, + windows: [ + { window: '5h', remainingPercent: 96, resetAt: future }, + { window: 'weekly', remainingPercent: 99, resetAt: future2 }, + ], + }, + }, + }, + ], + }) + 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') + // Same order assertion as (a): shortest window first. + expect(frame.indexOf('Gm 5h')).toBeLessThan(frame.indexOf('Gm 7d')) + expect(frame.indexOf('NG 5h')).toBeLessThan(frame.indexOf('NG 7d')) + testSetup.renderer.destroy() + }) +}) diff --git a/packages/opencode/src/tui.test.tsx b/packages/opencode/src/tui.test.tsx index 68b32fb..455a702 100644 --- a/packages/opencode/src/tui.test.tsx +++ b/packages/opencode/src/tui.test.tsx @@ -80,6 +80,11 @@ function makeFixture(): Fixture { const statePath = join(root, 'sidebar-state.json') const logPath = join(root, 'tui.log') const prefsPath = join(root, 'tui-preferences.jsonc') + // Save preload-pinned values to restore on cleanup instead of deleting — + // a delete drops resolution to the operator's real state/config dirs. + const savedSidebar = process.env[SIDEBAR_STATE_ENV] + const savedTuiLog = process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE + const savedPrefs = process.env[TUI_PREFS_FILE_ENV] process.env[SIDEBAR_STATE_ENV] = statePath process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE = logPath process.env[TUI_PREFS_FILE_ENV] = prefsPath @@ -88,9 +93,14 @@ function makeFixture(): Fixture { logPath, prefsPath, cleanup: () => { - delete process.env[SIDEBAR_STATE_ENV] - delete process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE - delete process.env[TUI_PREFS_FILE_ENV] + if (savedSidebar !== undefined) + process.env[SIDEBAR_STATE_ENV] = savedSidebar + else delete process.env[SIDEBAR_STATE_ENV] + if (savedTuiLog !== undefined) + process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE = savedTuiLog + else delete process.env.ANTIGRAVITY_AUTH_TUI_LOG_FILE + if (savedPrefs !== undefined) process.env[TUI_PREFS_FILE_ENV] = savedPrefs + else delete process.env[TUI_PREFS_FILE_ENV] rmSync(root, { recursive: true, force: true }) }, } @@ -195,9 +205,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 +216,7 @@ describe('SidebarPanel', () => { health: 42, current: false, quota: { - claude: { remainingPercent: 60 }, + 'non-gemini': { remainingPercent: 60 }, }, }, ], @@ -229,13 +238,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 +453,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 +482,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 +749,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 }, }, }, ], @@ -777,17 +781,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('Cl: 25%') + 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('Claude') - expect(frame).not.toContain('Gemini Pro') - expect(frame).not.toContain('Gemini Flash') expect(frame).not.toContain('cooldown') expect(frame).not.toContain('Awaiting Antigravity state') testSetup.renderer.destroy() @@ -834,9 +836,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 +864,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 +877,8 @@ describe('SidebarPanel collapse/expand + compact row', () => { await testSetup.flush() const collapsedFrame = testSetup.captureCharFrame() expect(collapsedFrame).toContain('▶') - expect(collapsedFrame).not.toContain('GP') + expect(collapsedFrame).toContain('Gm:') + expect(collapsedFrame).toContain('NG:') testSetup.renderer.destroy() }) @@ -895,9 +897,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 +961,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { enabled: true, health: 80, current: true, - quota: { claude: { remainingPercent: 75 } }, + quota: { 'non-gemini': { remainingPercent: 75 } }, }, ], }) @@ -1102,7 +1103,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 +1111,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { enabled: true, health: 60, current: false, - quota: { claude: { remainingPercent: 50 } }, + quota: { 'non-gemini': { remainingPercent: 50 } }, }, ], }) @@ -1154,7 +1155,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { enabled: true, health: 80, current: true, - quota: { claude: { remainingPercent: 75 } }, + quota: { 'non-gemini': { remainingPercent: 75 } }, }, ], }) @@ -1233,7 +1234,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { health: 80, current: true, quota: { - claude: { remainingPercent: 75 }, + 'non-gemini': { remainingPercent: 75 }, }, }, ], @@ -1283,7 +1284,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 +1292,7 @@ describe('SidebarPanel sections + themed border (T6)', () => { enabled: true, health: 60, current: false, - quota: { claude: { remainingPercent: 50 } }, + quota: { 'non-gemini': { remainingPercent: 50 } }, }, ], }) @@ -1337,7 +1338,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..3fb397b 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,16 @@ 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'] + +const WINDOW_GUTTER: Record = { + weekly: '7d', + '5h': '5h', +} // --- Theme tokens ---------------------------------------------------------- // @@ -517,14 +515,14 @@ function QuotaRow(props: { when={used() != null} fallback={ - {props.label.padEnd(3)} + {props.label.padEnd(5)} {'\u2014'} } > - + {props.label} - {(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) => ( + + ))} + + ) + }} {` ${healthText()}`} @@ -835,8 +857,39 @@ export function SidebarPanel(props: SidebarPanelProps): JSX.Element { visibleAccounts().find((entry) => entry.current) ?? visibleAccounts()[0] const used = () => { - const entry = account()?.quota.claude - return entry ? 100 - clamp(entry.remainingPercent, 0, 100) : null + const q = account()?.quota + // Derive tone from the WORST (most-constrained) pool so a + // critical non-gemini pool doesn't hide behind a healthy + // gemini tone. Worst = highest used-percent across available pools. + const geminiUsed = + q?.gemini != null + ? 100 - clamp(q.gemini.remainingPercent, 0, 100) + : null + const nonGeminiUsed = + q?.['non-gemini'] != null + ? 100 - clamp(q['non-gemini'].remainingPercent, 0, 100) + : null + if (geminiUsed === null && nonGeminiUsed === null) return null + return Math.max(geminiUsed ?? 0, nonGeminiUsed ?? 0) + } + const poolText = (key: 'gemini' | 'non-gemini') => { + const entry = account()?.quota[key] + if (!entry) return `${QUOTA_LABELS[key]}: —` + const pct = Math.round(100 - clamp(entry.remainingPercent, 0, 100)) + // 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() @@ -861,9 +914,10 @@ export function SidebarPanel(props: SidebarPanelProps): JSX.Element { )} > - {used() == null + {account()?.quota.gemini == null && + account()?.quota['non-gemini'] == null ? '—' - : `Cl: ${Math.round(used() ?? 0)}%`} + : `${poolText('gemini')} · ${poolText('non-gemini')}`} { } }) + 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') + // 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', + ) + 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) + // 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 () => { const localFake = makeFakeApi() dispatcher.openCommandDialog( @@ -1230,7 +1318,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 +1350,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..b8bac5c 100644 --- a/packages/opencode/src/tui/command-dialogs.tsx +++ b/packages/opencode/src/tui/command-dialogs.tsx @@ -252,6 +252,141 @@ 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) { + // 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) + }} + 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 +530,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__ ')) { diff --git a/test/environment.test.ts b/test/environment.test.ts index ddd32ec..def61ec 100644 --- a/test/environment.test.ts +++ b/test/environment.test.ts @@ -1,6 +1,14 @@ import { describe, expect, it } from 'bun:test' +import { homedir } from 'node:os' +import { join } from 'node:path' +import { getUserConfigPath } from '../packages/opencode/src/plugin/config/loader.ts' +import { getOpencodeConfigDir } from '../packages/opencode/src/plugin/config/updater.ts' +import { defaultFilesystemRoots } from '../packages/opencode/src/plugin/dependencies.ts' import { getConfigDir as getStorageConfigDir } from '../packages/opencode/src/plugin/storage.ts' +import { getRpcDir } from '../packages/opencode/src/rpc/rpc-dir.ts' +import { getSidebarStateFile } from '../packages/opencode/src/sidebar-state.ts' +import { resolveTuiLogPath } from '../packages/opencode/src/tui/file-logger.ts' import { getPiAntigravityAuthFile, getPiConfigDir, @@ -63,3 +71,227 @@ describe('test environment isolation', () => { expect(recovery.PART_STORAGE.startsWith(root)).toBe(true) }) }) + +/** + * Guard: every writable-path resolver must resolve inside the test root. + * + * Two assertions per resolver: + * (a) under-test: the resolved path IS inside the test root + * (b) production default: the resolved path is NOT inside the test root + * (i.e. it would write to the real user dir in production, proving + * the test-env override is actually doing work — a resolver + * accidentally hardcoded to a temp path would pass (a) but fail (b)) + * + * If any assertion fails the suite is silently writing to the operator's + * live state. The error message names the offending resolver and the + * path it resolved to so the root cause is immediately visible. + */ +describe('state leak guard — every path resolver must stay inside the test root', () => { + const root = process.env.ANTIGRAVITY_TEST_ROOT + if (!root) throw new Error('ANTIGRAVITY_TEST_ROOT not set by preload') + + // The operator's real home — used to assert production defaults differ. + // homedir() reads HOME (pinned by preload), so the production default is + // computed by stripping the test indirection. + const realHome = homedir() + // Production xdg-state default that the real user sees outside of tests. + const productionXdgState = join(realHome, '.local', 'state') + + // ── Sidebar state ───────────────────────────────────────────────────────── + + it('getSidebarStateFile() resolves under the test root', () => { + const resolved = getSidebarStateFile() + expect( + resolved.startsWith(root), + `getSidebarStateFile() resolved OUTSIDE test root: ${resolved}\n` + + "The suite would corrupt the operator's live sidebar state file. " + + 'Check that ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE is pinned in test/setup.ts.', + ).toBe(true) + }) + + it('getSidebarStateFile() production default is NOT under the test root', () => { + // Verifies the test-env override is doing real work: without it, the + // resolver would fall back to this production path. + const productionDefault = join( + productionXdgState, + 'cortexkit', + 'antigravity-auth', + 'sidebar-state.json', + ) + expect(productionDefault.startsWith(root)).toBe(false) + }) + + // ── RPC dir ─────────────────────────────────────────────────────────────── + + it('getRpcDir() resolves under the test root', () => { + // Use an arbitrary project directory; what matters is that the root dir + // is isolated, not the per-project hash suffix. + const resolved = getRpcDir('/some/project') + expect( + resolved.startsWith(root), + `getRpcDir() resolved OUTSIDE test root: ${resolved}\n` + + "The suite would write port files into the operator's real state dir " + + '(2119 leaked project dirs observed). Check that ANTIGRAVITY_AUTH_RPC_DIR ' + + 'is pinned in test/setup.ts.', + ).toBe(true) + }) + + it('getRpcDir() production default is NOT under the test root', () => { + const productionDefault = join( + productionXdgState, + 'cortexkit', + 'antigravity-auth', + 'rpc', + ) + expect(productionDefault.startsWith(root)).toBe(false) + }) + + // ── defaultFilesystemRoots() — same XDG_STATE_HOME gap ─────────────────── + + it('defaultFilesystemRoots() sidebarStateRoot resolves under the test root', () => { + const { sidebarStateRoot } = defaultFilesystemRoots() + expect( + sidebarStateRoot.startsWith(root), + `defaultFilesystemRoots().sidebarStateRoot resolved OUTSIDE test root: ${sidebarStateRoot}`, + ).toBe(true) + }) + + it('defaultFilesystemRoots() rpcRoot resolves under the test root', () => { + const { rpcRoot } = defaultFilesystemRoots() + expect( + rpcRoot.startsWith(root), + `defaultFilesystemRoots().rpcRoot resolved OUTSIDE test root: ${rpcRoot}`, + ).toBe(true) + }) + + it('defaultFilesystemRoots() production xdg-state defaults are NOT under the test root', () => { + const productionSidebar = join( + productionXdgState, + 'cortexkit', + 'antigravity-auth', + ) + const productionRpc = join( + productionXdgState, + 'cortexkit', + 'antigravity-auth', + 'rpc', + ) + expect(productionSidebar.startsWith(root)).toBe(false) + expect(productionRpc.startsWith(root)).toBe(false) + }) + + // ── TUI log ─────────────────────────────────────────────────────────────── + + it('resolveTuiLogPath() resolves under the test root', () => { + const resolved = resolveTuiLogPath() + expect( + resolved.startsWith(root), + `resolveTuiLogPath() resolved OUTSIDE test root: ${resolved}\n` + + 'Check that ANTIGRAVITY_AUTH_TUI_LOG_FILE or XDG_STATE_HOME is pinned.', + ).toBe(true) + }) + + it('resolveTuiLogPath() production default is NOT under the test root', () => { + const productionDefault = join( + productionXdgState, + 'cortexkit', + 'antigravity-auth', + 'tui.log', + ) + expect(productionDefault.startsWith(root)).toBe(false) + }) + + // ── Plugin config (getUserConfigPath, getOpencodeConfigDir) ─────────────── + + it('getUserConfigPath() resolves under the test root', () => { + const resolved = getUserConfigPath() + expect( + resolved.startsWith(root), + `getUserConfigPath() resolved OUTSIDE test root: ${resolved}\n` + + 'Check that OPENCODE_CONFIG_DIR is pinned in test/setup.ts.', + ).toBe(true) + }) + + it('getOpencodeConfigDir() resolves under the test root', () => { + // getOpencodeConfigDir() uses XDG_CONFIG_HOME (pinned), not OPENCODE_CONFIG_DIR. + const resolved = getOpencodeConfigDir() + expect( + resolved.startsWith(root), + `getOpencodeConfigDir() resolved OUTSIDE test root: ${resolved}`, + ).toBe(true) + }) + + it('config resolver production defaults are NOT under the test root', () => { + const productionConfig = join(realHome, '.config', 'opencode') + expect(productionConfig.startsWith(root)).toBe(false) + }) + + // ── Storage config dir ──────────────────────────────────────────────────── + + it('storage getConfigDir() resolves under the test root', () => { + const resolved = getStorageConfigDir() + expect( + resolved.startsWith(root), + `storage getConfigDir() resolved OUTSIDE test root: ${resolved}`, + ).toBe(true) + }) + + // ── auto-update-checker constants (module-level captures, separate class) ──── + // + // auto-update-checker/constants.ts uses os.homedir() directly (not HOME + // env) for CACHE_DIR. Bun's os.homedir() is an OS-level call that does NOT + // read process.env.HOME — it returns the passwd entry. This means CACHE_DIR + // resolves to the real ~/.cache/opencode regardless of env pinning. + // + // This is a named finding, not a failure to fix: the test suite does NOT + // write to CACHE_DIR. The write paths (installPackage, invalidatePackage) + // only fire from the plugin's session.created event, which unit tests never + // trigger. The only test-time read is an existsSync check in checker.ts:149 + // which is a no-op when the directory is absent. CACHE_DIR is therefore + // UNREACHABLE FOR WRITES in the unit test suite. + // + // USER_OPENCODE_CONFIG uses XDG_CONFIG_HOME (pinned) and resolves correctly. + // It is asserted here separately because its failure mode (import-order + // sensitivity) differs from the XDG_STATE_HOME gap above. + + it('auto-update-checker USER_OPENCODE_CONFIG resolves under the test root', async () => { + const { USER_OPENCODE_CONFIG } = await import( + `../packages/opencode/src/hooks/auto-update-checker/constants.ts?bust=${Date.now()}` + ) + expect( + (USER_OPENCODE_CONFIG as string).startsWith(root), + `auto-update-checker USER_OPENCODE_CONFIG resolved OUTSIDE test root: ${USER_OPENCODE_CONFIG as string}`, + ).toBe(true) + }) + + it('auto-update-checker USER_OPENCODE_CONFIG production default is NOT under the test root', () => { + const productionUserConfig = join( + realHome, + '.config', + 'opencode', + 'opencode.json', + ) + expect(productionUserConfig.startsWith(root)).toBe(false) + }) + + it('auto-update-checker CACHE_DIR resolves OUTSIDE the test root (write-unreachable)', async () => { + // CACHE_DIR uses os.homedir() (OS-level, not env-based) and resolves + // to the real ~/.cache/opencode even under test — env pinning cannot + // redirect it. This is the (b)-half guard: CACHE_DIR must be outside + // the test root. If it ever drifts inside, env isolation has broken + // and the situation must be re-evaluated. + const { CACHE_DIR } = await import( + `../packages/opencode/src/hooks/auto-update-checker/constants.ts?bust=${Date.now()}` + ) + expect( + (CACHE_DIR as string).startsWith(root), + `CACHE_DIR resolved INSIDE test root: ${CACHE_DIR as string}\n` + + 'os.homedir() now returns the test home — the isolation rationale no longer holds. ' + + 'Audit write paths before the suite is considered safe.', + ).toBe(false) + // The real protection: none of the writing callsites are reachable from + // bun test. cache.ts::installPackage and invalidatePackage are only called + // from runBackgroundUpdateCheck which fires on session.created events that + // unit tests never emit. Verified by grep: no test imports cache.ts. + }) +}) diff --git a/test/setup.ts b/test/setup.ts index a39fc81..8dbe06c 100644 --- a/test/setup.ts +++ b/test/setup.ts @@ -8,9 +8,10 @@ const home = join(root, 'home') const config = join(root, 'config') const cache = join(root, 'cache') const data = join(root, 'data') +const state = join(root, 'state') const pi = join(root, 'pi-agent') -for (const path of [home, config, cache, data, pi]) +for (const path of [home, config, cache, data, state, pi]) mkdirSync(path, { recursive: true }) process.env.ANTIGRAVITY_TEST_ROOT = root @@ -19,12 +20,32 @@ process.env.USERPROFILE = home process.env.XDG_CONFIG_HOME = config process.env.XDG_CACHE_HOME = cache process.env.XDG_DATA_HOME = data +// XDG_STATE_HOME must be pinned before sidebar-state.ts and rpc-dir.ts import, +// because xdg-basedir reads state from env once at module load, and both +// resolvers fall back to when their env +// override is unset. Without this pin, any test that deletes the per-test +// override drops resolution to the operator's real ~/.local/state dir. +process.env.XDG_STATE_HOME = state process.env.APPDATA = config process.env.LOCALAPPDATA = cache process.env.OPENCODE_CONFIG_DIR = join(config, 'opencode') process.env.PI_AGENT_DIR = pi process.env.PI_ANTIGRAVITY_AUTH_FILE = join(pi, 'antigravity-accounts.json') process.env.OPENCODE_ANTIGRAVITY_GEMINI_DUMP_DIR = join(root, 'gemini-dumps') +// Pin the two state-dir overrides so tests that reset them restore to an +// in-root path rather than dropping to the real default. +process.env.ANTIGRAVITY_AUTH_SIDEBAR_STATE_FILE = join( + state, + 'cortexkit', + 'antigravity-auth', + 'sidebar-state.json', +) +process.env.ANTIGRAVITY_AUTH_RPC_DIR = join( + state, + 'cortexkit', + 'antigravity-auth', + 'rpc', +) const stubbedGlobals = new Map() const dateNowSpyState: { active: boolean } = { active: false }