diff --git a/src/plugin.ts b/src/plugin.ts index d3862f1..1cfe8e3 100644 --- a/src/plugin.ts +++ b/src/plugin.ts @@ -1417,7 +1417,7 @@ export const createAntigravityPlugin = (providerId: string) => async ( // Note: AccountManager now ensures the current auth is always included in accounts - const accountManager = await AccountManager.loadFromDisk(auth); + let accountManager = await AccountManager.loadFromDisk(auth); activeAccountManager = accountManager; if (accountManager.getAccountCount() > 0) { accountManager.requestSaveToDisk(); @@ -1468,9 +1468,8 @@ export const createAntigravityPlugin = (providerId: string) => async ( return fetch(input, init); } - if (accountManager.getAccountCount() === 0) { - throw new Error("No Antigravity accounts configured. Run `opencode auth login`."); - } + accountManager = await AccountManager.loadFromDisk(latestAuth); + activeAccountManager = accountManager; const urlString = toUrlString(input); const family = getModelFamilyFromUrl(urlString); @@ -1481,6 +1480,26 @@ export const createAntigravityPlugin = (providerId: string) => async ( debugLines.push(line); }; pushDebug(`request=${urlString}`); + let reloadedProviderState = false; + const reloadProviderState = async (reason: string): Promise => { + if (reloadedProviderState) { + return false; + } + const reloadedAuth = await getAuth(); + if (!isOAuthAuth(reloadedAuth)) { + return false; + } + reloadedProviderState = true; + accountManager = await AccountManager.loadFromDisk(reloadedAuth); + activeAccountManager = accountManager; + rateLimitStateByAccountQuota.clear(); + emptyResponseAttempts.clear(); + accountFailureState.clear(); + rateLimitToastShown = false; + softQuotaToastShown = false; + pushDebug(`reload-provider-state reason=${reason} accounts=${accountManager.getAccountCount()}`); + return true; + }; type FailureContext = { response: Response; @@ -1561,6 +1580,9 @@ export const createAntigravityPlugin = (providerId: string) => async ( } = routingDecision; if (accountCount === 0) { + if (await reloadProviderState("no-accounts")) { + continue; + } throw new Error("No Antigravity accounts available. Run `opencode auth login`."); } @@ -1653,6 +1675,9 @@ export const createAntigravityPlugin = (providerId: string) => async ( // 0 means disabled (wait indefinitely) const maxWaitMs = (config.max_rate_limit_wait_seconds ?? 300) * 1000; if (maxWaitMs > 0 && waitMs > maxWaitMs) { + if (await reloadProviderState(`all-rate-limited:${family}`)) { + continue; + } const waitTimeFormatted = formatWaitTime(waitMs); await showToast( `Rate limited for ${waitTimeFormatted}. Try again later or add another account.`, @@ -1832,6 +1857,7 @@ export const createAntigravityPlugin = (providerId: string) => async ( const warmupUrl = toWarmupStreamUrl(prepared.request); const warmupHeaders = new Headers(prepared.init.headers ?? {}); warmupHeaders.set("accept", "text/event-stream"); + warmupHeaders.delete("x-goog-api-key"); const warmupInit: RequestInit = { ...prepared.init, @@ -1984,6 +2010,13 @@ export const createAntigravityPlugin = (providerId: string) => async ( }, ); + const requestHeaders = new Headers(prepared.init.headers ?? {}); + requestHeaders.delete("x-goog-api-key"); + const requestInit: RequestInit = { + ...prepared.init, + headers: requestHeaders, + }; + const originalUrl = toUrlString(input); const resolvedUrl = toUrlString(prepared.request); pushDebug(`endpoint=${currentEndpoint}`); @@ -1991,9 +2024,9 @@ export const createAntigravityPlugin = (providerId: string) => async ( const debugContext = startAntigravityDebugRequest({ originalUrl, resolvedUrl, - method: prepared.init.method, - headers: prepared.init.headers, - body: prepared.init.body, + method: requestInit.method, + headers: requestInit.headers, + body: requestInit.body, streaming: prepared.streaming, projectId: projectContext.effectiveProjectId, }); @@ -2027,7 +2060,7 @@ export const createAntigravityPlugin = (providerId: string) => async ( tokenConsumed = getTokenTracker().consume(account.index); } - const response = await fetch(prepared.request, prepared.init); + const response = await fetch(prepared.request, requestInit); pushDebug(`status=${response.status} ${response.statusText}`); diff --git a/src/plugin/accounts.test.ts b/src/plugin/accounts.test.ts index ceda7db..d64fd3c 100644 --- a/src/plugin/accounts.test.ts +++ b/src/plugin/accounts.test.ts @@ -1,7 +1,8 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { AccountManager, type ModelFamily, type HeaderStyle, parseRateLimitReason, calculateBackoffMs, type RateLimitReason, resolveQuotaGroup } from "./accounts"; -import { saveAccounts, type AccountStorageV4 } from "./storage"; +import { saveAccounts } from "./storage"; +import type { AccountStorageV4 } from "./storage"; import type { OAuthAuthDetails } from "./types"; // Mock storage to prevent test data from leaking to real config files @@ -255,6 +256,40 @@ describe("AccountManager", () => { expect(snapshot[1]?.expires).toBe(123); }); + it("saveToDisk preserves empty rateLimitResetTimes so stale limits can be cleared", async () => { + const stored: AccountStorageV4 = { + version: 4, + accounts: [ + { + refreshToken: "r1", + projectId: "p1", + addedAt: 1, + lastUsed: 0, + rateLimitResetTimes: { claude: 123456 }, + }, + ], + activeIndex: 0, + }; + + const manager = new AccountManager(undefined, stored); + const account = manager.getAccounts()[0]; + if (!account) throw new Error("Account not found"); + + account.rateLimitResetTimes = {}; + await manager.saveToDisk(); + + expect(vi.mocked(saveAccounts)).toHaveBeenCalledWith( + expect.objectContaining({ + accounts: [ + expect.objectContaining({ + refreshToken: "r1", + rateLimitResetTimes: {}, + }), + ], + }), + ); + }); + it("debounces toast display for same account", () => { vi.useFakeTimers(); vi.setSystemTime(new Date(0)); diff --git a/src/plugin/accounts.ts b/src/plugin/accounts.ts index 885ae1a..d0cef3a 100644 --- a/src/plugin/accounts.ts +++ b/src/plugin/accounts.ts @@ -1007,7 +1007,8 @@ export class AccountManager { lastUsed: a.lastUsed, enabled: a.enabled, lastSwitchReason: a.lastSwitchReason, - rateLimitResetTimes: Object.keys(a.rateLimitResetTimes).length > 0 ? { ...a.rateLimitResetTimes } : undefined, + // Persist an empty object when limits were explicitly cleared so disk state stays authoritative. + rateLimitResetTimes: { ...a.rateLimitResetTimes }, coolingDownUntil: a.coolingDownUntil, cooldownReason: a.cooldownReason, fingerprint: a.fingerprint, diff --git a/src/plugin/storage.test.ts b/src/plugin/storage.test.ts index 6120303..dfee945 100644 --- a/src/plugin/storage.test.ts +++ b/src/plugin/storage.test.ts @@ -1,11 +1,13 @@ import { describe, expect, it, vi, beforeEach } from "vitest"; import { deduplicateAccountsByEmail, + mergeAccountStorage, migrateV2ToV3, loadAccounts, saveAccounts, type AccountMetadata, type AccountStorage, + type AccountStorageV4, } from "./storage"; import { promises as fs } from "node:fs"; import { @@ -214,6 +216,71 @@ describe("deduplicateAccountsByEmail", () => { }); }); +describe("mergeAccountStorage", () => { + it("lets an empty incoming rateLimitResetTimes clear stale stored limits", () => { + const existing: AccountStorageV4 = { + version: 4, + accounts: [ + { + refreshToken: "r1", + projectId: "p1", + addedAt: 1, + lastUsed: 10, + rateLimitResetTimes: { claude: 123456 }, + }, + ], + activeIndex: 0, + }; + const incoming: AccountStorageV4 = { + version: 4, + accounts: [ + { + refreshToken: "r1", + projectId: "p1", + addedAt: 1, + lastUsed: 20, + rateLimitResetTimes: {}, + }, + ], + activeIndex: 0, + }; + + const merged = mergeAccountStorage(existing, incoming); + expect(merged.accounts[0]?.rateLimitResetTimes).toBeUndefined(); + }); + + it("preserves existing limits when incoming storage does not specify them", () => { + const existing: AccountStorageV4 = { + version: 4, + accounts: [ + { + refreshToken: "r1", + projectId: "p1", + addedAt: 1, + lastUsed: 10, + rateLimitResetTimes: { claude: 123456 }, + }, + ], + activeIndex: 0, + }; + const incoming: AccountStorageV4 = { + version: 4, + accounts: [ + { + refreshToken: "r1", + projectId: "p1", + addedAt: 1, + lastUsed: 20, + }, + ], + activeIndex: 0, + }; + + const merged = mergeAccountStorage(existing, incoming); + expect(merged.accounts[0]?.rateLimitResetTimes).toEqual({ claude: 123456 }); + }); +}); + vi.mock("node:fs", async () => { const actual = await vi.importActual("node:fs"); return { diff --git a/src/plugin/storage.ts b/src/plugin/storage.ts index 538bcf3..1e232ca 100644 --- a/src/plugin/storage.ts +++ b/src/plugin/storage.ts @@ -445,7 +445,7 @@ async function withFileLock(path: string, fn: () => Promise): Promise { } } -function mergeAccountStorage( +export function mergeAccountStorage( existing: AccountStorageV4, incoming: AccountStorageV4, ): AccountStorageV4 { @@ -461,18 +461,23 @@ function mergeAccountStorage( if (acc.refreshToken) { const existingAcc = accountMap.get(acc.refreshToken); if (existingAcc) { + const incomingRateLimitResetTimes = acc.rateLimitResetTimes; + const mergedRateLimitResetTimes = incomingRateLimitResetTimes === undefined + ? existingAcc.rateLimitResetTimes + : Object.keys(incomingRateLimitResetTimes).length === 0 + ? undefined + : { + ...existingAcc.rateLimitResetTimes, + ...incomingRateLimitResetTimes, + }; accountMap.set(acc.refreshToken, { ...existingAcc, ...acc, - // Preserve manually configured projectId/managedProjectId from disk. - // Existing (disk) values take priority over incoming (in-memory) values - // to prevent auto-detected fallback IDs from overwriting user edits. + // Existing disk values take priority so manual overrides survive auth refresh merges. projectId: existingAcc.projectId ?? acc.projectId, managedProjectId: existingAcc.managedProjectId ?? acc.managedProjectId, - rateLimitResetTimes: { - ...existingAcc.rateLimitResetTimes, - ...acc.rateLimitResetTimes, - }, + // An explicit empty object means limits were cleared and should overwrite older disk state. + rateLimitResetTimes: mergedRateLimitResetTimes, lastUsed: Math.max(existingAcc.lastUsed || 0, acc.lastUsed || 0), }); } else {