opencode-antigravity-auth 2.0.0 → 2.1.1

package/dist/src/plugin.js

@@ -12,7 +12,7 @@ import { isEmptyResponseBody, createSyntheticErrorResponse, } from "./plugin/req
 import { EmptyResponseError } from "./plugin/errors.js";
 import { AntigravityTokenRefreshError, refreshAccessToken } from "./plugin/token.js";
 import { startOAuthListener } from "./plugin/server.js";
-import { clearAccounts, loadAccounts, saveAccounts, saveAccountsReplace } from "./plugin/storage.js";
+import { clearAccounts, loadAccounts, saveAccountCredentials, saveAccounts, saveAccountsReplace, updateAccountCredentials, updateAccounts } from "./plugin/storage.js";
 import { AccountManager, parseRateLimitReason, calculateBackoffMs, computeSoftQuotaCacheTtlMs } from "./plugin/accounts.js";
 import { createAutoUpdateCheckerHook } from "./hooks/auto-update-checker/index.js";
 import { loadConfig, initRuntimeConfig } from "./plugin/config/index.js";
@@ -98,7 +98,19 @@ async function triggerAsyncQuotaRefreshForAccount(accountManager, accountIndex,
             quotaRefreshInProgressByEmail.delete(accountKey);
             return;
         }
-        const results = await checkAccountsQuota([singleAccount], client, providerId);
+        const results = await checkAccountsQuota([singleAccount], client, providerId, { onAccountUpdated: saveAccountCredentials });
+        const updatedAccount = results[0]?.updatedAccount;
+        if (updatedAccount) {
+            const saved = await updateAccounts(storage => {
+                const current = storage.accounts.find(item => singleAccount.id ? item.id === singleAccount.id : item.refreshToken === singleAccount.refreshToken);
+                if (current?.refreshToken === singleAccount.refreshToken)
+                    updateAccountCredentials(current, updatedAccount);
+            });
+            const current = saved?.accounts.find(item => singleAccount.id ? item.id === singleAccount.id : item.refreshToken === updatedAccount.refreshToken);
+            if (current && account.parts.refreshToken === singleAccount.refreshToken) {
+                accountManager.updateFromAuth(account, { ...accountManager.toAuthDetails(account), refresh: formatRefreshParts(current) }, singleAccount.refreshToken);
+            }
+        }
         if (results[0]?.status === "ok" && results[0]?.quota?.groups) {
             accountManager.updateQuotaCache(accountIndex, results[0].quota.groups);
             accountManager.requestSaveToDisk();
@@ -993,15 +1005,17 @@ function sleep(ms, signal) {
 /**
  * Creates an Antigravity OAuth plugin for a specific provider ID.
  */
-export const createAntigravityPlugin = (providerId) => async ({ client, directory, runtimeVersion = 1, fetch: sendFetch = globalThis.fetch }) => {
+export const createAntigravityPlugin = (providerId) => async ({ client, directory, runtimeVersion = 1, fetch: sendFetch = globalThis.fetch, canUseAccount }) => {
     // Load configuration from files and environment variables
     const config = loadConfig(directory);
     initRuntimeConfig(config);
     // Cached getAuth function for tool access
     let cachedGetAuth = null;
     const refreshQueues = new Set();
+    const quotaRefreshes = new Set();
     const transportManagers = new Set();
     let disposed = false;
+    let managerGeneration = 0;
     // Initialize debug with config
     initializeDebug(config);
     // Initialize structured logger for TUI integration
@@ -1144,19 +1158,23 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
     });
     return {
         async flush() {
-            await Promise.allSettled([...transportManagers].map((manager) => manager.flushSaveToDisk()));
+            managerGeneration += 1;
+            await Promise.allSettled([...refreshQueues].map((queue) => queue.stop()));
+            refreshQueues.clear();
+            await Promise.allSettled([...transportManagers].map((manager) => manager.retire()));
             transportManagers.clear();
         },
         async dispose() {
             if (disposed)
                 return;
             disposed = true;
+            managerGeneration += 1;
             releaseSignatureCache?.();
             clearInterval(versionRefreshTimer);
-            for (const queue of refreshQueues)
-                queue.stop();
+            await Promise.allSettled([...refreshQueues].map((queue) => queue.stop()));
             refreshQueues.clear();
-            await Promise.allSettled([...transportManagers].map((manager) => manager.flushSaveToDisk()));
+            await Promise.allSettled([...quotaRefreshes]);
+            await Promise.allSettled([...transportManagers].map((manager) => manager.retire()));
             transportManagers.clear();
             cachedGetAuth = null;
         },
@@ -1169,11 +1187,11 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
             loader: async (getAuth, provider) => {
                 if (disposed)
                     throw new Error("Antigravity plugin has been disposed");
+                const expectedGeneration = managerGeneration;
                 if (runtimeVersion === 2) {
-                    for (const queue of refreshQueues)
-                        queue.stop();
+                    await Promise.allSettled([...refreshQueues].map((queue) => queue.stop()));
                     refreshQueues.clear();
-                    await Promise.allSettled([...transportManagers].map((manager) => manager.flushSaveToDisk()));
+                    await Promise.allSettled([...transportManagers].map((manager) => manager.retire()));
                     transportManagers.clear();
                 }
                 // Cache getAuth for tool access
@@ -1196,11 +1214,22 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
                 const authParts = parseRefreshParts(auth.refresh);
                 const storedAccounts = await loadAccounts();
                 // Note: AccountManager now ensures the current auth is always included in accounts
-                const accountManager = await AccountManager.loadFromDisk(auth);
-                if (disposed)
-                    throw new Error("Antigravity plugin has been disposed");
+                const accountManager = await AccountManager.loadFromDisk(auth, runtimeVersion === 2);
+                if (runtimeVersion === 2)
+                    accountManager.useManagedStorage();
+                const abandonManager = async () => {
+                    await accountManager.retire();
+                    transportManagers.delete(accountManager);
+                    throw new Error(disposed ? "Antigravity plugin has been disposed" : "Antigravity accounts changed");
+                };
+                if (disposed || managerGeneration !== expectedGeneration)
+                    return abandonManager();
                 activeAccountManager = accountManager;
                 transportManagers.add(accountManager);
+                if (runtimeVersion === 2)
+                    await accountManager.saveToDisk();
+                if (disposed || managerGeneration !== expectedGeneration)
+                    return abandonManager();
                 if (accountManager.getAccountCount() > 0) {
                     accountManager.requestSaveToDisk();
                 }
@@ -1242,7 +1271,13 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
                         if (!isGenerativeLanguageRequest(input)) {
                             return sendFetch(input, init);
                         }
+                        const checkLifetime = () => {
+                            if (disposed || (runtimeVersion === 2 && managerGeneration !== expectedGeneration))
+                                throw new Error("Antigravity accounts changed; retry the request");
+                        };
+                        checkLifetime();
                         const latestAuth = await getAuth();
+                        checkLifetime();
                         if (!isOAuthAuth(latestAuth)) {
                             return sendFetch(input, init);
                         }
@@ -1252,6 +1287,10 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
                         const urlString = toUrlString(input);
                         const family = getModelFamilyFromUrl(urlString);
                         const model = extractModelFromUrl(urlString);
+                        const eligible = (account) => !canUseAccount || canUseAccount({ id: account.id, refreshToken: account.parts.refreshToken }, model ?? "");
+                        if (!accountManager.getAccounts().some(account => account.enabled !== false && eligible(account))) {
+                            throw new Error("No enabled Antigravity account has access to this model. Refresh available models from /antigravity.");
+                        }
                         const debugLines = [];
                         const pushDebug = (line) => {
                             if (!isDebugEnabled())
@@ -1303,30 +1342,30 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
                             if (family !== "gemini")
                                 return false;
                             // Use AccountManager method which properly checks for disabled/cooling-down accounts
-                            return accountManager.hasOtherAccountWithAntigravityAvailable(currentAccount.index, family, model);
+                            return accountManager.hasOtherAccountWithAntigravityAvailable(currentAccount.index, family, model, eligible);
                         };
                         while (true) {
                             // Check for abort at the start of each iteration
                             checkAborted();
-                            const accountCount = accountManager.getAccountCount();
+                            const accountCount = accountManager.getEnabledAccounts().filter(eligible).length;
                             const routingDecision = resolveHeaderRoutingDecision(urlString, family, config);
                             const { cliFirst, preferredHeaderStyle, explicitQuota, allowQuotaFallback, } = routingDecision;
                             if (accountCount === 0) {
                                 throw new Error("No Antigravity accounts available. Run `opencode auth login`.");
                             }
                             const softQuotaCacheTtlMs = computeSoftQuotaCacheTtlMs(config.soft_quota_cache_ttl_minutes, config.quota_refresh_interval_minutes);
-                            let account = accountManager.getCurrentOrNextForFamily(family, model, config.account_selection_strategy, preferredHeaderStyle, config.pid_offset_enabled, config.soft_quota_threshold_percent, softQuotaCacheTtlMs);
+                            let account = accountManager.getCurrentOrNextForFamily(family, model, config.account_selection_strategy, preferredHeaderStyle, config.pid_offset_enabled, config.soft_quota_threshold_percent, softQuotaCacheTtlMs, eligible);
                             if (!account && allowQuotaFallback) {
                                 const alternateHeaderStyle = preferredHeaderStyle === "antigravity" ? "gemini-cli" : "antigravity";
-                                account = accountManager.getCurrentOrNextForFamily(family, model, config.account_selection_strategy, alternateHeaderStyle, config.pid_offset_enabled, config.soft_quota_threshold_percent, softQuotaCacheTtlMs);
+                                account = accountManager.getCurrentOrNextForFamily(family, model, config.account_selection_strategy, alternateHeaderStyle, config.pid_offset_enabled, config.soft_quota_threshold_percent, softQuotaCacheTtlMs, eligible);
                                 if (account) {
                                     pushDebug(`selected-by-fallback idx=${account.index} preferred=${preferredHeaderStyle} alternate=${alternateHeaderStyle}`);
                                 }
                             }
                             if (!account) {
-                                if (accountManager.areAllAccountsOverSoftQuota(family, config.soft_quota_threshold_percent, softQuotaCacheTtlMs, model)) {
+                                if (accountManager.areAllAccountsOverSoftQuota(family, config.soft_quota_threshold_percent, softQuotaCacheTtlMs, model, eligible)) {
                                     const threshold = config.soft_quota_threshold_percent;
-                                    const softQuotaWaitMs = accountManager.getMinWaitTimeForSoftQuota(family, threshold, softQuotaCacheTtlMs, model);
+                                    const softQuotaWaitMs = accountManager.getMinWaitTimeForSoftQuota(family, threshold, softQuotaCacheTtlMs, model, eligible);
                                     const maxWaitMs = (config.max_rate_limit_wait_seconds ?? 300) * 1000;
                                     if (softQuotaWaitMs === null || (maxWaitMs > 0 && softQuotaWaitMs > maxWaitMs)) {
                                         const waitTimeFormatted = softQuotaWaitMs ? formatWaitTime(softQuotaWaitMs) : "unknown";
@@ -1346,7 +1385,7 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
                                 }
                                 const strictWait = !allowQuotaFallback;
                                 // All accounts are rate-limited - wait and retry
-                                const waitMs = accountManager.getMinWaitTimeForFamily(family, model, preferredHeaderStyle, strictWait) || 60_000;
+                                const waitMs = accountManager.getMinWaitTimeForFamily(family, model, preferredHeaderStyle, strictWait, eligible) || 60_000;
                                 const waitSecValue = Math.max(1, Math.ceil(waitMs / 1000));
                                 pushDebug(`all-rate-limited family=${family} accounts=${accountCount} waitMs=${waitMs}`);
                                 if (isDebugEnabled()) {
@@ -1414,7 +1453,7 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
                                         continue;
                                     }
                                     resetAccountFailureState(account.index);
-                                    accountManager.updateFromAuth(account, refreshed);
+                                    accountManager.updateFromAuth(account, refreshed, parseRefreshParts(authRecord.refresh).refreshToken);
                                     authRecord = refreshed;
                                     try {
                                         await accountManager.saveToDisk();
@@ -1425,9 +1464,24 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
                                 }
                                 catch (error) {
                                     if (error instanceof AntigravityTokenRefreshError && error.code === "invalid_grant") {
-                                        const removed = accountManager.removeAccount(account);
+                                        if (runtimeVersion === 2) {
+                                            const rejectedToken = parseRefreshParts(authRecord.refresh).refreshToken;
+                                            const saved = await updateAccounts(storage => {
+                                                const current = storage.accounts.find(item => account.id ? item.id === account.id : item.refreshToken === rejectedToken);
+                                                if (current?.refreshToken === rejectedToken)
+                                                    current.enabled = false;
+                                            });
+                                            const current = saved?.accounts.find(item => account.id ? item.id === account.id : item.refreshToken === rejectedToken);
+                                            if (current && current.refreshToken !== rejectedToken) {
+                                                accountManager.updateFromAuth(account, { type: "oauth", refresh: formatRefreshParts(current) });
+                                                continue;
+                                            }
+                                        }
+                                        const removed = runtimeVersion === 2
+                                            ? accountManager.setAccountEnabled(account.index, false)
+                                            : accountManager.removeAccount(account);
                                         if (removed) {
-                                            log.warn("Removed revoked account from pool - reauthenticate via `opencode auth login`");
+                                            log.warn("Revoked account excluded from rotation - reauthenticate via `opencode auth login`");
                                             try {
                                                 await accountManager.saveToDisk();
                                             }
@@ -1487,7 +1541,7 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
                             }
                             if (projectContext.auth.refresh !== authRecord.refresh ||
                                 projectContext.auth.access !== authRecord.access) {
-                                accountManager.updateFromAuth(account, projectContext.auth);
+                                accountManager.updateFromAuth(account, projectContext.auth, parseRefreshParts(authRecord.refresh).refreshToken);
                                 authRecord = projectContext.auth;
                                 try {
                                     await accountManager.saveToDisk();
@@ -1554,7 +1608,7 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
                                 // Antigravity-first fallback: exhaust antigravity across ALL accounts before gemini-cli
                                 if (allowQuotaFallback && family === "gemini" && headerStyle === "antigravity") {
                                     // Check if ANY other account has antigravity available
-                                    if (accountManager.hasOtherAccountWithAntigravityAvailable(account.index, family, model)) {
+                                    if (accountManager.hasOtherAccountWithAntigravityAvailable(account.index, family, model, eligible)) {
                                         // Switch to another account with antigravity (preserve antigravity priority)
                                         pushDebug(`antigravity rate-limited on account ${account.index}, but available on other accounts. Switching.`);
                                         shouldSwitchAccount = true;
@@ -1867,7 +1921,11 @@ export const createAntigravityPlugin = (providerId) => async ({ client, director
                                             account.consecutiveFailures = 0;
                                             getHealthTracker().recordSuccess(account.index);
                                             accountManager.markAccountUsed(account.index);
-                                            void triggerAsyncQuotaRefreshForAccount(accountManager, account.index, client, providerId, config.quota_refresh_interval_minutes);
+                                            if (!disposed) {
+                                                const quotaRefresh = triggerAsyncQuotaRefreshForAccount(accountManager, account.index, client, providerId, config.quota_refresh_interval_minutes);
+                                                quotaRefreshes.add(quotaRefresh);
+                                                void quotaRefresh.then(() => quotaRefreshes.delete(quotaRefresh), () => quotaRefreshes.delete(quotaRefresh));
+                                            }
                                         }
                                         logAntigravityDebugResponse(debugContext, response, {
                                             note: response.ok ? "Success" : `Error ${response.status}`,