opencode-antigravity-auth 1.0.6 → 1.1.0

package/dist/src/plugin.js

@@ -1,13 +1,162 @@
 import { exec } from "node:child_process";
-import { ANTIGRAVITY_ENDPOINT_FALLBACKS, ANTIGRAVITY_PROVIDER_ID, ANTIGRAVITY_REDIRECT_URI, } from "./constants";
+import { ANTIGRAVITY_ENDPOINT_FALLBACKS, ANTIGRAVITY_PROVIDER_ID } from "./constants";
 import { authorizeAntigravity, exchangeAntigravity } from "./antigravity/oauth";
-import { accessTokenExpired, isOAuthAuth } from "./plugin/auth";
-import { promptProjectId } from "./plugin/cli";
+import { accessTokenExpired, isOAuthAuth, parseRefreshParts } from "./plugin/auth";
+import { promptAddAnotherAccount, promptProjectId } from "./plugin/cli";
 import { ensureProjectContext } from "./plugin/project";
 import { startAntigravityDebugRequest } from "./plugin/debug";
 import { isGenerativeLanguageRequest, prepareAntigravityRequest, transformAntigravityResponse, } from "./plugin/request";
-import { refreshAccessToken } from "./plugin/token";
+import { AntigravityTokenRefreshError, refreshAccessToken } from "./plugin/token";
 import { startOAuthListener } from "./plugin/server";
+import { loadAccounts, saveAccounts } from "./plugin/storage";
+import { AccountManager } from "./plugin/accounts";
+const MAX_OAUTH_ACCOUNTS = 10;
+async function openBrowser(url) {
+    try {
+        if (process.platform === "darwin") {
+            exec(`open "${url}"`);
+            return;
+        }
+        if (process.platform === "win32") {
+            exec(`start "${url}"`);
+            return;
+        }
+        exec(`xdg-open "${url}"`);
+    }
+    catch {
+        // ignore
+    }
+}
+async function promptOAuthCallbackValue(message) {
+    const { createInterface } = await import("node:readline/promises");
+    const { stdin, stdout } = await import("node:process");
+    const rl = createInterface({ input: stdin, output: stdout });
+    try {
+        return (await rl.question(message)).trim();
+    }
+    finally {
+        rl.close();
+    }
+}
+function getStateFromAuthorizationUrl(authorizationUrl) {
+    try {
+        return new URL(authorizationUrl).searchParams.get("state") ?? "";
+    }
+    catch {
+        return "";
+    }
+}
+function extractOAuthCallbackParams(url) {
+    const code = url.searchParams.get("code");
+    const state = url.searchParams.get("state");
+    if (!code || !state) {
+        return null;
+    }
+    return { code, state };
+}
+function parseOAuthCallbackInput(value, fallbackState) {
+    const trimmed = value.trim();
+    if (!trimmed) {
+        return { error: "Missing authorization code" };
+    }
+    try {
+        const url = new URL(trimmed);
+        const code = url.searchParams.get("code");
+        const state = url.searchParams.get("state") ?? fallbackState;
+        if (!code) {
+            return { error: "Missing code in callback URL" };
+        }
+        if (!state) {
+            return { error: "Missing state in callback URL" };
+        }
+        return { code, state };
+    }
+    catch {
+        if (!fallbackState) {
+            return { error: "Missing state. Paste the full redirect URL instead of only the code." };
+        }
+        return { code: trimmed, state: fallbackState };
+    }
+}
+function clampInt(value, min, max) {
+    if (!Number.isFinite(value)) {
+        return min;
+    }
+    return Math.min(max, Math.max(min, Math.floor(value)));
+}
+async function persistAccountPool(results) {
+    if (results.length === 0) {
+        return;
+    }
+    const now = Date.now();
+    const stored = await loadAccounts();
+    const accounts = stored?.accounts ? [...stored.accounts] : [];
+    const indexByRefreshToken = new Map();
+    for (let i = 0; i < accounts.length; i++) {
+        const token = accounts[i]?.refreshToken;
+        if (token) {
+            indexByRefreshToken.set(token, i);
+        }
+    }
+    for (const result of results) {
+        const parts = parseRefreshParts(result.refresh);
+        if (!parts.refreshToken) {
+            continue;
+        }
+        const existingIndex = indexByRefreshToken.get(parts.refreshToken);
+        if (existingIndex === undefined) {
+            indexByRefreshToken.set(parts.refreshToken, accounts.length);
+            accounts.push({
+                email: result.email,
+                refreshToken: parts.refreshToken,
+                projectId: parts.projectId,
+                managedProjectId: parts.managedProjectId,
+                addedAt: now,
+                lastUsed: now,
+                isRateLimited: false,
+                rateLimitResetTime: 0,
+            });
+            continue;
+        }
+        const existing = accounts[existingIndex];
+        if (!existing) {
+            continue;
+        }
+        accounts[existingIndex] = {
+            ...existing,
+            email: result.email ?? existing.email,
+            projectId: parts.projectId ?? existing.projectId,
+            managedProjectId: parts.managedProjectId ?? existing.managedProjectId,
+            lastUsed: now,
+        };
+    }
+    if (accounts.length === 0) {
+        return;
+    }
+    const activeIndex = typeof stored?.activeIndex === "number" && Number.isFinite(stored.activeIndex) ? stored.activeIndex : 0;
+    await saveAccounts({
+        version: 1,
+        accounts,
+        activeIndex: clampInt(activeIndex, 0, accounts.length - 1),
+    });
+}
+function retryAfterMsFromResponse(response) {
+    const retryAfterMsHeader = response.headers.get("retry-after-ms");
+    if (retryAfterMsHeader) {
+        const parsed = Number.parseInt(retryAfterMsHeader, 10);
+        if (!Number.isNaN(parsed) && parsed > 0) {
+            return parsed;
+        }
+    }
+    const retryAfterHeader = response.headers.get("retry-after");
+    if (retryAfterHeader) {
+        const parsed = Number.parseInt(retryAfterHeader, 10);
+        if (!Number.isNaN(parsed) && parsed > 0) {
+            return parsed * 1000;
+        }
+    }
+    return 60_000;
+}
 /**
  * Creates an Antigravity OAuth plugin for a specific provider ID.
  */
@@ -19,6 +168,15 @@ export const createAntigravityPlugin = (providerId) => async ({ client }) => ({
             if (!isOAuthAuth(auth)) {
                 return null;
             }
+            const accountManager = await AccountManager.loadFromDisk(auth);
+            if (accountManager.getAccountCount() > 0) {
+                try {
+                    await accountManager.saveToDisk();
+                }
+                catch (error) {
+                    console.error("[opencode-antigravity-auth] Failed to persist initial account pool:", error);
+                }
+            }
             if (provider.models) {
                 for (const model of Object.values(provider.models)) {
                     if (model) {
@@ -39,89 +197,170 @@ export const createAntigravityPlugin = (providerId) => async ({ client }) => ({
                     if (!isOAuthAuth(latestAuth)) {
                         return fetch(input, init);
                     }
-                    let authRecord = latestAuth;
-                    if (accessTokenExpired(authRecord)) {
-                        const refreshed = await refreshAccessToken(authRecord, client, providerId);
-                        if (!refreshed) {
-                            return fetch(input, init);
-                        }
-                        authRecord = refreshed;
-                    }
-                    const accessToken = authRecord.access;
-                    if (!accessToken) {
-                        return fetch(input, init);
+                    const accountCount = accountManager.getAccountCount();
+                    if (accountCount === 0) {
+                        throw new Error("No Antigravity accounts configured. Run `opencode auth login`.");
                     }
-                    /**
-                     * Ensures we have a usable project context for the current auth snapshot.
-                     */
-                    async function resolveProjectContext() {
+                    let lastFailure = null;
+                    let lastError = null;
+                    accountLoop: for (let attempt = 0; attempt < accountCount; attempt++) {
+                        const account = accountManager.pickNext();
+                        if (!account) {
+                            const waitMs = accountManager.getMinWaitTimeMs();
+                            const waitSec = Math.max(1, Math.ceil(waitMs / 1000));
+                            throw new Error(`All ${accountManager.getAccountCount()} account(s) are rate-limited. Retry in ${waitSec}s or add more accounts via 'opencode auth login'.`);
+                        }
                         try {
-                            return await ensureProjectContext(authRecord, client, providerId);
+                            await accountManager.saveToDisk();
                         }
                         catch (error) {
-                            throw error;
+                            console.error("[opencode-antigravity-auth] Failed to persist rotation state:", error);
                         }
-                    }
-                    const projectContext = await resolveProjectContext();
-                    // Endpoint fallback logic: try daily → autopush → prod
-                    let lastError = null;
-                    let lastResponse = null;
-                    for (let i = 0; i < ANTIGRAVITY_ENDPOINT_FALLBACKS.length; i++) {
-                        const currentEndpoint = ANTIGRAVITY_ENDPOINT_FALLBACKS[i];
-                        try {
-                            const { request, init: transformedInit, streaming, requestedModel, effectiveModel, projectId: usedProjectId, endpoint: usedEndpoint, toolDebugMissing, toolDebugSummary, toolDebugPayload, } = prepareAntigravityRequest(input, init, accessToken, projectContext.effectiveProjectId, currentEndpoint);
-                            const originalUrl = toUrlString(input);
-                            const resolvedUrl = toUrlString(request);
-                            const debugContext = startAntigravityDebugRequest({
-                                originalUrl,
-                                resolvedUrl,
-                                method: transformedInit.method,
-                                headers: transformedInit.headers,
-                                body: transformedInit.body,
-                                streaming,
-                                projectId: projectContext.effectiveProjectId,
-                            });
-                            const response = await fetch(request, transformedInit);
-                            // Check if we should retry with next endpoint
-                            const shouldRetry = (response.status === 403 || // Forbidden
-                                response.status === 404 || // Not Found
-                                response.status === 429 || // Rate Limit
-                                response.status >= 500 // Server errors
-                            );
-                            if (shouldRetry && i < ANTIGRAVITY_ENDPOINT_FALLBACKS.length - 1) {
-                                // Try next endpoint
-                                lastResponse = response;
+                        let authRecord = accountManager.toAuthDetails(account);
+                        if (accessTokenExpired(authRecord)) {
+                            try {
+                                const refreshed = await refreshAccessToken(authRecord, client, providerId);
+                                if (!refreshed) {
+                                    lastError = new Error("Antigravity token refresh failed");
+                                    continue;
+                                }
+                                accountManager.updateFromAuth(account, refreshed);
+                                authRecord = refreshed;
+                                try {
+                                    await accountManager.saveToDisk();
+                                }
+                                catch (error) {
+                                    console.error("[opencode-antigravity-auth] Failed to persist refreshed auth:", error);
+                                }
+                            }
+                            catch (error) {
+                                if (error instanceof AntigravityTokenRefreshError && error.code === "invalid_grant") {
+                                    const removed = accountManager.removeAccount(account);
+                                    if (removed) {
+                                        console.warn("[opencode-antigravity-auth] Removed revoked account from pool. Reauthenticate it via `opencode auth login` to add it back.");
+                                        try {
+                                            await accountManager.saveToDisk();
+                                        }
+                                        catch (persistError) {
+                                            console.error("[opencode-antigravity-auth] Failed to persist revoked account removal:", persistError);
+                                        }
+                                    }
+                                    if (accountManager.getAccountCount() === 0) {
+                                        try {
+                                            await client.auth.set({
+                                                path: { id: providerId },
+                                                body: { type: "oauth", refresh: "" },
+                                            });
+                                        }
+                                        catch (storeError) {
+                                            console.error("Failed to clear stored Antigravity OAuth credentials:", storeError);
+                                        }
+                                        throw new Error("All Antigravity accounts have invalid refresh tokens. Run `opencode auth login` and reauthenticate.");
+                                    }
+                                    lastError = error;
+                                    continue;
+                                }
+                                lastError = error instanceof Error ? error : new Error(String(error));
                                 continue;
                             }
-                            // Success or final attempt - return transformed response
-                            return transformAntigravityResponse(response, streaming, debugContext, requestedModel, usedProjectId, usedEndpoint, effectiveModel, toolDebugMissing, toolDebugSummary, toolDebugPayload);
+                        }
+                        const accessToken = authRecord.access;
+                        if (!accessToken) {
+                            lastError = new Error("Missing access token");
+                            continue;
+                        }
+                        let projectContext;
+                        try {
+                            projectContext = await ensureProjectContext(authRecord);
                         }
                         catch (error) {
-                            // Network error or other exception
-                            if (i < ANTIGRAVITY_ENDPOINT_FALLBACKS.length - 1) {
+                            lastError = error instanceof Error ? error : new Error(String(error));
+                            continue;
+                        }
+                        if (projectContext.auth !== authRecord) {
+                            accountManager.updateFromAuth(account, projectContext.auth);
+                            authRecord = projectContext.auth;
+                            try {
+                                await accountManager.saveToDisk();
+                            }
+                            catch (error) {
+                                console.error("[opencode-antigravity-auth] Failed to persist project context:", error);
+                            }
+                        }
+                        for (let i = 0; i < ANTIGRAVITY_ENDPOINT_FALLBACKS.length; i++) {
+                            const currentEndpoint = ANTIGRAVITY_ENDPOINT_FALLBACKS[i];
+                            try {
+                                const prepared = prepareAntigravityRequest(input, init, accessToken, projectContext.effectiveProjectId, currentEndpoint);
+                                const originalUrl = toUrlString(input);
+                                const resolvedUrl = toUrlString(prepared.request);
+                                const debugContext = startAntigravityDebugRequest({
+                                    originalUrl,
+                                    resolvedUrl,
+                                    method: prepared.init.method,
+                                    headers: prepared.init.headers,
+                                    body: prepared.init.body,
+                                    streaming: prepared.streaming,
+                                    projectId: projectContext.effectiveProjectId,
+                                });
+                                const response = await fetch(prepared.request, prepared.init);
+                                if (response.status === 429 && accountManager.getAccountCount() > 1) {
+                                    const retryAfterMs = retryAfterMsFromResponse(response);
+                                    accountManager.markRateLimited(account, retryAfterMs);
+                                    try {
+                                        await accountManager.saveToDisk();
+                                    }
+                                    catch (error) {
+                                        console.error("[opencode-antigravity-auth] Failed to persist rate-limit state:", error);
+                                    }
+                                    lastFailure = {
+                                        response,
+                                        streaming: prepared.streaming,
+                                        debugContext,
+                                        requestedModel: prepared.requestedModel,
+                                        projectId: prepared.projectId,
+                                        endpoint: prepared.endpoint,
+                                        effectiveModel: prepared.effectiveModel,
+                                        toolDebugMissing: prepared.toolDebugMissing,
+                                        toolDebugSummary: prepared.toolDebugSummary,
+                                        toolDebugPayload: prepared.toolDebugPayload,
+                                    };
+                                    continue accountLoop;
+                                }
+                                const shouldRetry = (response.status === 403 ||
+                                    response.status === 404 ||
+                                    (response.status === 429 && accountManager.getAccountCount() <= 1) ||
+                                    response.status >= 500);
+                                if (shouldRetry && i < ANTIGRAVITY_ENDPOINT_FALLBACKS.length - 1) {
+                                    lastFailure = {
+                                        response,
+                                        streaming: prepared.streaming,
+                                        debugContext,
+                                        requestedModel: prepared.requestedModel,
+                                        projectId: prepared.projectId,
+                                        endpoint: prepared.endpoint,
+                                        effectiveModel: prepared.effectiveModel,
+                                        toolDebugMissing: prepared.toolDebugMissing,
+                                        toolDebugSummary: prepared.toolDebugSummary,
+                                        toolDebugPayload: prepared.toolDebugPayload,
+                                    };
+                                    continue;
+                                }
+                                return transformAntigravityResponse(response, prepared.streaming, debugContext, prepared.requestedModel, prepared.projectId, prepared.endpoint, prepared.effectiveModel, prepared.toolDebugMissing, prepared.toolDebugSummary, prepared.toolDebugPayload);
+                            }
+                            catch (error) {
+                                if (i < ANTIGRAVITY_ENDPOINT_FALLBACKS.length - 1) {
+                                    lastError = error instanceof Error ? error : new Error(String(error));
+                                    continue;
+                                }
                                 lastError = error instanceof Error ? error : new Error(String(error));
-                                continue;
+                                continue accountLoop;
                             }
-                            // Final attempt failed, throw the error
-                            throw error;
                         }
                     }
-                    // If we get here, all endpoints failed
-                    if (lastResponse) {
-                        // Return the last response even if it was an error
-                        const { streaming, requestedModel, effectiveModel, projectId: usedProjectId, endpoint: usedEndpoint, toolDebugMissing, toolDebugSummary, toolDebugPayload, } = prepareAntigravityRequest(input, init, accessToken, projectContext.effectiveProjectId, ANTIGRAVITY_ENDPOINT_FALLBACKS[ANTIGRAVITY_ENDPOINT_FALLBACKS.length - 1]);
-                        const debugContext = startAntigravityDebugRequest({
-                            originalUrl: toUrlString(input),
-                            resolvedUrl: toUrlString(input),
-                            method: init?.method,
-                            headers: init?.headers,
-                            body: init?.body,
-                            streaming,
-                            projectId: projectContext.effectiveProjectId,
-                        });
-                        return transformAntigravityResponse(lastResponse, streaming, debugContext, requestedModel, usedProjectId, usedEndpoint, effectiveModel, toolDebugMissing, toolDebugSummary, toolDebugPayload);
+                    if (lastFailure) {
+                        return transformAntigravityResponse(lastFailure.response, lastFailure.streaming, lastFailure.debugContext, lastFailure.requestedModel, lastFailure.projectId, lastFailure.endpoint, lastFailure.effectiveModel, lastFailure.toolDebugMissing, lastFailure.toolDebugSummary, lastFailure.toolDebugPayload);
                     }
-                    throw lastError || new Error("All Antigravity endpoints failed");
+                    throw lastError || new Error("All Antigravity accounts failed");
                 },
             };
         },
@@ -129,59 +368,148 @@ export const createAntigravityPlugin = (providerId) => async ({ client }) => ({
             {
                 label: "OAuth with Google (Antigravity)",
                 type: "oauth",
-                authorize: async () => {
+                authorize: async (inputs) => {
                     const isHeadless = !!(process.env.SSH_CONNECTION ||
                         process.env.SSH_CLIENT ||
                         process.env.SSH_TTY ||
                         process.env.OPENCODE_HEADLESS);
+                    // CLI flow (`opencode auth login`) passes an inputs object.
+                    if (inputs) {
+                        const accounts = [];
+                        while (accounts.length < MAX_OAUTH_ACCOUNTS) {
+                            console.log(`\n=== Antigravity OAuth (Account ${accounts.length + 1}) ===`);
+                            const projectId = await promptProjectId();
+                            const result = await (async () => {
+                                let listener = null;
+                                if (!isHeadless) {
+                                    try {
+                                        listener = await startOAuthListener();
+                                    }
+                                    catch {
+                                        listener = null;
+                                    }
+                                }
+                                const authorization = await authorizeAntigravity(projectId);
+                                const fallbackState = getStateFromAuthorizationUrl(authorization.url);
+                                console.log("\nOAuth URL:\n" + authorization.url + "\n");
+                                if (!isHeadless) {
+                                    await openBrowser(authorization.url);
+                                }
+                                if (listener) {
+                                    try {
+                                        const callbackUrl = await listener.waitForCallback();
+                                        const params = extractOAuthCallbackParams(callbackUrl);
+                                        if (!params) {
+                                            return { type: "failed", error: "Missing code or state in callback URL" };
+                                        }
+                                        return exchangeAntigravity(params.code, params.state);
+                                    }
+                                    catch (error) {
+                                        return {
+                                            type: "failed",
+                                            error: error instanceof Error ? error.message : "Unknown error",
+                                        };
+                                    }
+                                    finally {
+                                        try {
+                                            await listener.close();
+                                        }
+                                        catch {
+                                            // ignore
+                                        }
+                                    }
+                                }
+                                console.log("1. Open the URL below in your browser and complete Google sign-in.");
+                                console.log("2. After approving, copy the full redirected localhost URL from the address bar.");
+                                console.log("3. Paste it back here.");
+                                const callbackInput = await promptOAuthCallbackValue("Paste the redirect URL (or just the code) here: ");
+                                const params = parseOAuthCallbackInput(callbackInput, fallbackState);
+                                if ("error" in params) {
+                                    return { type: "failed", error: params.error };
+                                }
+                                return exchangeAntigravity(params.code, params.state);
+                            })();
+                            if (result.type === "failed") {
+                                if (accounts.length === 0) {
+                                    return {
+                                        url: "",
+                                        instructions: `Authentication failed: ${result.error}`,
+                                        method: "auto",
+                                        callback: async () => result,
+                                    };
+                                }
+                                console.warn(`[opencode-antigravity-auth] Skipping failed account ${accounts.length + 1}: ${result.error}`);
+                                break;
+                            }
+                            accounts.push(result);
+                            try {
+                                await persistAccountPool([result]);
+                            }
+                            catch {
+                                // ignore
+                            }
+                            if (accounts.length >= MAX_OAUTH_ACCOUNTS) {
+                                break;
+                            }
+                            const addAnother = await promptAddAnotherAccount(accounts.length);
+                            if (!addAnother) {
+                                break;
+                            }
+                        }
+                        const primary = accounts[0];
+                        if (!primary) {
+                            return {
+                                url: "",
+                                instructions: "Authentication cancelled",
+                                method: "auto",
+                                callback: async () => ({ type: "failed", error: "Authentication cancelled" }),
+                            };
+                        }
+                        return {
+                            url: "",
+                            instructions: `Multi-account setup complete (${accounts.length} account(s)).`,
+                            method: "auto",
+                            callback: async () => primary,
+                        };
+                    }
+                    // TUI flow (`/connect`) does not support per-account prompts yet.
+                    const projectId = "";
                     let listener = null;
                     if (!isHeadless) {
                         try {
                             listener = await startOAuthListener();
                         }
-                        catch (error) {
-                            console.log("\nWarning: Couldn't start the local callback listener. Falling back to manual copy/paste.");
+                        catch {
+                            listener = null;
                         }
                     }
-                    const authorization = await authorizeAntigravity("");
-                    // Try to open the browser automatically
+                    const authorization = await authorizeAntigravity(projectId);
+                    const fallbackState = getStateFromAuthorizationUrl(authorization.url);
                     if (!isHeadless) {
-                        try {
-                            if (process.platform === "darwin") {
-                                exec(`open "${authorization.url}"`);
-                            }
-                            else if (process.platform === "win32") {
-                                exec(`start "${authorization.url}"`);
-                            }
-                            else {
-                                exec(`xdg-open "${authorization.url}"`);
-                            }
-                        }
-                        catch (e) {
-                            console.log("Could not open browser automatically. Please Copy/Paste the URL below.");
-                        }
+                        await openBrowser(authorization.url);
                     }
                     if (listener) {
-                        const { host } = new URL(ANTIGRAVITY_REDIRECT_URI);
                         return {
                             url: authorization.url,
-                            instructions: "Complete the sign-in flow in your browser. We'll automatically detect the redirect back to localhost.",
+                            instructions: "Complete sign-in in your browser. We'll automatically detect the redirect back to localhost.",
                             method: "auto",
                             callback: async () => {
                                 try {
-                                    // We know listener is not null here because we checked 'if (listener)'
-                                    // But TS might need a check or non-null assertion if not inferable.
-                                    // Since we are in the if (listener) block, it is safe.
                                     const callbackUrl = await listener.waitForCallback();
-                                    const code = callbackUrl.searchParams.get("code");
-                                    const state = callbackUrl.searchParams.get("state");
-                                    if (!code || !state) {
-                                        return {
-                                            type: "failed",
-                                            error: "Missing code or state in callback URL",
-                                        };
+                                    const params = extractOAuthCallbackParams(callbackUrl);
+                                    if (!params) {
+                                        return { type: "failed", error: "Missing code or state in callback URL" };
                                     }
-                                    return await exchangeAntigravity(code, state);
+                                    const result = await exchangeAntigravity(params.code, params.state);
+                                    if (result.type === "success") {
+                                        try {
+                                            await persistAccountPool([result]);
+                                        }
+                                        catch {
+                                            // ignore
+                                        }
+                                    }
+                                    return result;
                                 }
                                 catch (error) {
                                     return {
@@ -191,9 +519,10 @@ export const createAntigravityPlugin = (providerId) => async ({ client }) => ({
                                 }
                                 finally {
                                     try {
-                                        await listener?.close();
+                                        await listener.close();
                                     }
                                     catch {
+                                        // ignore
                                     }
                                 }
                             },
@@ -201,27 +530,23 @@ export const createAntigravityPlugin = (providerId) => async ({ client }) => ({
                     }
                     return {
                         url: authorization.url,
-                        instructions: "Paste the full redirected URL (e.g., http://localhost:8085/oauth2callback?code=...): ",
+                        instructions: "Visit the URL above, complete OAuth, then paste either the full redirect URL or the authorization code.",
                         method: "code",
-                        callback: async (callbackUrl) => {
-                            try {
-                                const url = new URL(callbackUrl);
-                                const code = url.searchParams.get("code");
-                                const state = url.searchParams.get("state");
-                                if (!code || !state) {
-                                    return {
-                                        type: "failed",
-                                        error: "Missing code or state in callback URL",
-                                    };
-                                }
-                                return exchangeAntigravity(code, state);
+                        callback: async (codeInput) => {
+                            const params = parseOAuthCallbackInput(codeInput, fallbackState);
+                            if ("error" in params) {
+                                return { type: "failed", error: params.error };
                             }
-                            catch (error) {
-                                return {
-                                    type: "failed",
-                                    error: error instanceof Error ? error.message : "Unknown error",
-                                };
+                            const result = await exchangeAntigravity(params.code, params.state);
+                            if (result.type === "success") {
+                                try {
+                                    await persistAccountPool([result]);
+                                }
+                                catch {
+                                    // ignore
+                                }
                             }
+                            return result;
                         },
                     };
                 },