Files
opencode-antigravity-auth/src/plugin/token.ts
Noe 2052e4d580 Add multi-account load balancing and improved OAuth UX
Adds multi-account support and round-robin load balancing for Google Antigravity OAuth to increase request throughput and resilience. Introduces an on-disk account pool with cooldowns for rate-limited accounts, automatic removal of revoked refresh tokens, and persistence of rotation state.

Improves OAuth flows and UX: CLI flow can add multiple accounts with per-account project IDs, TUI flow remains single-account, improved browser opening/fallback copy-paste handling, and clearer prompts for pasting redirect URLs or codes. Adds robust parsing of callback input and better headless handling.

Makes token refresh handling explicit and typed (throws a specific error on invalid_grant) and centralizes account management logic into an in-memory manager with persistence utilities. Adds tests for account rotation and rate-limit behavior and bumps package version.

Overall, this increases reliability under rate limits, makes multi-account configuration straightforward, and improves error handling and developer/user experience.
2025-12-16 01:35:02 +00:00

169 lines
4.8 KiB
TypeScript

import { ANTIGRAVITY_CLIENT_ID, ANTIGRAVITY_CLIENT_SECRET } from "../constants";
import { formatRefreshParts, parseRefreshParts } from "./auth";
import { clearCachedAuth, storeCachedAuth } from "./cache";
import { invalidateProjectContextCache } from "./project";
import type { OAuthAuthDetails, PluginClient, RefreshParts } from "./types";
interface OAuthErrorPayload {
error?:
| string
| {
code?: string;
status?: string;
message?: string;
};
error_description?: string;
}
/**
* Parses OAuth error payloads returned by Google token endpoints, tolerating varied shapes.
*/
function parseOAuthErrorPayload(text: string | undefined): { code?: string; description?: string } {
if (!text) {
return {};
}
try {
const payload = JSON.parse(text) as OAuthErrorPayload;
if (!payload || typeof payload !== "object") {
return { description: text };
}
let code: string | undefined;
if (typeof payload.error === "string") {
code = payload.error;
} else if (payload.error && typeof payload.error === "object") {
code = payload.error.status ?? payload.error.code;
if (!payload.error_description && payload.error.message) {
return { code, description: payload.error.message };
}
}
const description = payload.error_description;
if (description) {
return { code, description };
}
if (payload.error && typeof payload.error === "object" && payload.error.message) {
return { code, description: payload.error.message };
}
return { code };
} catch {
return { description: text };
}
}
export class AntigravityTokenRefreshError extends Error {
code?: string;
description?: string;
status: number;
statusText: string;
constructor(options: {
message: string;
code?: string;
description?: string;
status: number;
statusText: string;
}) {
super(options.message);
this.name = "AntigravityTokenRefreshError";
this.code = options.code;
this.description = options.description;
this.status = options.status;
this.statusText = options.statusText;
}
}
/**
* Refreshes an Antigravity OAuth access token, updates persisted credentials, and handles revocation.
*/
export async function refreshAccessToken(
auth: OAuthAuthDetails,
client: PluginClient,
providerId: string,
): Promise<OAuthAuthDetails | undefined> {
const parts = parseRefreshParts(auth.refresh);
if (!parts.refreshToken) {
return undefined;
}
try {
const response = await fetch("https://oauth2.googleapis.com/token", {
method: "POST",
headers: {
"Content-Type": "application/x-www-form-urlencoded",
},
body: new URLSearchParams({
grant_type: "refresh_token",
refresh_token: parts.refreshToken,
client_id: ANTIGRAVITY_CLIENT_ID,
client_secret: ANTIGRAVITY_CLIENT_SECRET,
}),
});
if (!response.ok) {
let errorText: string | undefined;
try {
errorText = await response.text();
} catch {
errorText = undefined;
}
const { code, description } = parseOAuthErrorPayload(errorText);
const details = [code, description ?? errorText].filter(Boolean).join(": ");
const baseMessage = `Antigravity token refresh failed (${response.status} ${response.statusText})`;
const message = details ? `${baseMessage} - ${details}` : baseMessage;
console.warn(`[Antigravity OAuth] ${message}`);
if (code === "invalid_grant") {
console.warn(
"[Antigravity OAuth] Google revoked the stored refresh token for this account. Reauthenticate it via `opencode auth login`.",
);
invalidateProjectContextCache(auth.refresh);
clearCachedAuth(auth.refresh);
}
throw new AntigravityTokenRefreshError({
message,
code,
description: description ?? errorText,
status: response.status,
statusText: response.statusText,
});
}
const payload = (await response.json()) as {
access_token: string;
expires_in: number;
refresh_token?: string;
};
const refreshedParts: RefreshParts = {
refreshToken: payload.refresh_token ?? parts.refreshToken,
projectId: parts.projectId,
managedProjectId: parts.managedProjectId,
};
const updatedAuth: OAuthAuthDetails = {
...auth,
access: payload.access_token,
expires: Date.now() + payload.expires_in * 1000,
refresh: formatRefreshParts(refreshedParts),
};
storeCachedAuth(updatedAuth);
invalidateProjectContextCache(auth.refresh);
return updatedAuth;
} catch (error) {
if (error instanceof AntigravityTokenRefreshError) {
throw error;
}
console.error("Failed to refresh Antigravity access token due to an unexpected error:", error);
return undefined;
}
}