From f725711435b182dbdfde86d73f3fd3ee345a7fe0 Mon Sep 17 00:00:00 2001 From: SB Yoon <44089734+yansigit@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:08:34 -0600 Subject: [PATCH 01/10] feat(antigravity): live quota RPC and geoblock classification Probe retrieveUserQuota with catalog fallback, skip http Bearer destinations, and surface Cloud Code Assist location blocks before generic 403s. Co-authored-by: Cursor --- src/adapters/google-antigravity-hosts.ts | 24 ++ src/adapters/google-errors.ts | 7 + src/providers/antigravity-quota.ts | 212 +++++++++++++++ src/providers/quota.ts | 87 +++++-- tests/antigravity-quota.test.ts | 319 +++++++++++++++++++++++ tests/google-antigravity-errors.test.ts | 59 +++++ 6 files changed, 682 insertions(+), 26 deletions(-) create mode 100644 src/adapters/google-antigravity-hosts.ts create mode 100644 src/providers/antigravity-quota.ts create mode 100644 tests/antigravity-quota.test.ts create mode 100644 tests/google-antigravity-errors.test.ts diff --git a/src/adapters/google-antigravity-hosts.ts b/src/adapters/google-antigravity-hosts.ts new file mode 100644 index 0000000000..b9dc041b17 --- /dev/null +++ b/src/adapters/google-antigravity-hosts.ts @@ -0,0 +1,24 @@ +const DAILY_ANTIGRAVITY_HOST = "https://daily-cloudcode-pa.googleapis.com"; +const PROD_ANTIGRAVITY_HOST = "https://cloudcode-pa.googleapis.com"; + +/** + * Return the configured Antigravity endpoint followed by its daily/production peer. + * The configured value is preserved so tests and future pinned environments keep their + * explicit first choice; the fallback is always one of Google's two known hosts. + */ +export function antigravityHostCandidates(configuredBase: string): string[] { + const configured = configuredBase.replace(/\/+$/, ""); + const other = configured === DAILY_ANTIGRAVITY_HOST + ? PROD_ANTIGRAVITY_HOST + : DAILY_ANTIGRAVITY_HOST; + return [...new Set([configured, other])]; +} + +/** OAuth bearer requests must not use a cleartext host, even if generic baseUrl config allows http. */ +export function isAntigravityHttpsHost(host: string): boolean { + try { + return new URL(host).protocol === "https:"; + } catch { + return false; + } +} diff --git a/src/adapters/google-errors.ts b/src/adapters/google-errors.ts index 69e6d0cef5..371d2b70b9 100644 --- a/src/adapters/google-errors.ts +++ b/src/adapters/google-errors.ts @@ -15,6 +15,12 @@ function googleErrorDetail(payloadText: string): { message?: string; status?: st }; } +const ANTIGRAVITY_GEO_BLOCKED_MARKER = "user location is not supported for the api use"; + +export function isAntigravityGeoBlockedBody(payloadText: string): boolean { + return payloadText.toLowerCase().includes(ANTIGRAVITY_GEO_BLOCKED_MARKER); +} + function classifyGoogle(label: string, status: number | undefined, enumStatus: string | undefined, text: string): string { const lower = `${enumStatus ?? ""} ${text}`.toLowerCase(); const quotaExhausted = @@ -29,6 +35,7 @@ function classifyGoogle(label: string, status: number | undefined, enumStatus: s if (status === 401 || enumStatus === "UNAUTHENTICATED" || lower.includes("unauthenticated") || lower.includes("invalid authentication") || lower.includes("expired")) { return `${label} authentication failed`; } + if (isAntigravityGeoBlockedBody(lower)) return `${label} location not supported`; if (status === 403 || enumStatus === "PERMISSION_DENIED" || lower.includes("permission denied") || lower.includes("access denied")) { return `${label} access denied`; } diff --git a/src/providers/antigravity-quota.ts b/src/providers/antigravity-quota.ts new file mode 100644 index 0000000000..339df293ab --- /dev/null +++ b/src/providers/antigravity-quota.ts @@ -0,0 +1,212 @@ +import { antigravityUserAgent } from "../adapters/client-fingerprint"; +import { antigravityHostCandidates, isAntigravityHttpsHost } from "../adapters/google-antigravity-hosts"; +import { readProviderQuotaJsonForTests } from "./quota"; +import type { ProviderQuota, ProviderQuotaWindow } from "./quota"; + +const LIVE_QUOTA_PATH = "/v1internal:retrieveUserQuota"; +const LIVE_SUMMARY_PATH = "/v1internal:retrieveUserQuotaSummary"; + +type FetchImpl = typeof fetch; + +export interface AntigravityLiveQuotaArgs { + accessToken: string; + projectId: string; + baseUrl: string; + timeoutMs: number; + fetchImpl?: FetchImpl; +} + +interface QuotaCandidate { + record: Record; + path: string[]; +} + +function asRecord(value: unknown): Record | null { + return value && typeof value === "object" && !Array.isArray(value) + ? value as Record + : null; +} + +function finiteNumber(value: unknown): number | undefined { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim()) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : undefined; + } + return undefined; +} + +function normalizePercent(value: unknown): number | undefined { + const numeric = finiteNumber(value); + return numeric === undefined ? undefined : Math.max(0, Math.min(100, numeric)); +} + +function resetAt(value: unknown): number | undefined { + const numeric = finiteNumber(value); + if (numeric !== undefined && numeric > 0) return numeric > 10_000_000_000 ? numeric : numeric * 1000; + if (typeof value !== "string" || !value.trim()) return undefined; + const parsed = Date.parse(value); + return Number.isFinite(parsed) && parsed > 0 ? parsed : undefined; +} + +function remainingPercent(record: Record): number | undefined { + const fraction = finiteNumber(record.remainingFraction); + if (fraction !== undefined) return normalizePercent(fraction * 100); + const percentage = finiteNumber( + record.remainingPercentage + ?? record.remainingPercent + ?? record.remaining_percent, + ); + if (percentage !== undefined) return normalizePercent(percentage <= 1 ? percentage * 100 : percentage); + return undefined; +} + +function usedPercent(record: Record): number | undefined { + const remaining = remainingPercent(record); + return remaining === undefined ? undefined : normalizePercent(100 - remaining); +} + +function recordResetAt(record: Record): number | undefined { + return resetAt(record.resetTime ?? record.resetAt ?? record.resetsAt ?? record.reset_time ?? record.nextReset); +} + +function collectCandidates(value: unknown, path: string[] = [], output: QuotaCandidate[] = []): QuotaCandidate[] { + if (Array.isArray(value)) { + for (const [index, item] of value.entries()) collectCandidates(item, [...path, String(index)], output); + return output; + } + const record = asRecord(value); + if (!record) return output; + output.push({ record, path }); + for (const [key, child] of Object.entries(record)) { + if (child && typeof child === "object") collectCandidates(child, [...path, key], output); + } + return output; +} + +function candidateModelName(candidate: QuotaCandidate): string { + const { record, path } = candidate; + const explicit = record.modelId ?? record.model_id ?? record.modelName ?? record.model ?? record.name; + return `${typeof explicit === "string" ? explicit : ""} ${path.join(" ")}`.toLowerCase(); +} + +function parseGeminiWindow(payload: unknown): ProviderQuotaWindow | undefined { + for (const candidate of collectCandidates(payload)) { + if (!candidateModelName(candidate).includes("gemini")) continue; + const percent = usedPercent(candidate.record); + if (percent === undefined) continue; + const reset = recordResetAt(candidate.record); + return { + label: "Gem", + percent, + ...(reset !== undefined ? { resetAt: reset } : {}), + }; + } + return undefined; +} + +function isWeeklyPath(path: string[]): boolean { + return path.some(part => /weekly|week|seven[_-]?day/i.test(part)); +} + +function parseWeeklyWindow(payload: unknown): { percent: number; resetAt?: number } | undefined { + const candidates = collectCandidates(payload); + const ordered = [ + ...candidates.filter(candidate => isWeeklyPath(candidate.path)), + ...candidates.filter(candidate => !isWeeklyPath(candidate.path)), + ]; + for (const candidate of ordered) { + const percent = usedPercent(candidate.record); + if (percent === undefined) continue; + const reset = recordResetAt(candidate.record); + return { percent, ...(reset !== undefined ? { resetAt: reset } : {}) }; + } + return undefined; +} + +async function readJson(response: Response, timeoutMs: number): Promise { + return await readProviderQuotaJsonForTests(response, timeoutMs); +} + +class AntigravityQuotaRpcError extends Error { + constructor(readonly status: number) { + super(`Antigravity quota RPC failed: ${status}`); + } +} + +function shouldRetryPeer(status: number): boolean { + return status === 404 || status === 503; +} + +async function fetchRpc( + fetchImpl: FetchImpl, + host: string, + method: "retrieveUserQuota" | "retrieveUserQuotaSummary", + args: AntigravityLiveQuotaArgs, +): Promise { + const path = method === "retrieveUserQuota" ? LIVE_QUOTA_PATH : LIVE_SUMMARY_PATH; + const response = await fetchImpl(`${host}${path}`, { + method: "POST", + headers: { + Accept: "application/json", + "Content-Type": "application/json", + "User-Agent": antigravityUserAgent(), + Authorization: `Bearer ${args.accessToken}`, + }, + body: JSON.stringify({ project: args.projectId }), + redirect: "error", + signal: AbortSignal.timeout(args.timeoutMs), + }); + if (!response.ok) throw new AntigravityQuotaRpcError(response.status); + return readJson(response, args.timeoutMs); +} + +async function fetchHostQuota( + fetchImpl: FetchImpl, + host: string, + args: AntigravityLiveQuotaArgs, +): Promise { + const [quotaResult, summaryResult] = await Promise.allSettled([ + fetchRpc(fetchImpl, host, "retrieveUserQuota", args), + fetchRpc(fetchImpl, host, "retrieveUserQuotaSummary", args), + ]); + for (const result of [quotaResult, summaryResult]) { + if ( + result.status === "rejected" + && result.reason instanceof AntigravityQuotaRpcError + && !shouldRetryPeer(result.reason.status) + ) { + throw result.reason; + } + } + if (quotaResult.status === "rejected" || summaryResult.status === "rejected") return null; + const quotaPayload = quotaResult.value; + const summaryPayload = summaryResult.value; + const gem = parseGeminiWindow(quotaPayload); + const weekly = parseWeeklyWindow(summaryPayload); + if (!gem && !weekly) return null; + return { + ...(gem ? { customWindows: [gem] } : {}), + ...(weekly ? { + weeklyPercent: weekly.percent, + ...(weekly.resetAt !== undefined ? { weeklyResetAt: weekly.resetAt } : {}), + } : {}), + updatedAt: Date.now(), + }; +} + +export async function fetchAntigravityLiveQuota( + args: AntigravityLiveQuotaArgs, +): Promise { + const fetchImpl = args.fetchImpl ?? fetch; + for (const host of antigravityHostCandidates(args.baseUrl)) { + if (!isAntigravityHttpsHost(host)) continue; + try { + const quota = await fetchHostQuota(fetchImpl, host, args); + if (quota) return quota; + } catch { + return null; + } + } + return null; +} diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 24b06ef7a9..7bc646a988 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -28,6 +28,8 @@ import { type CodexCapacityAggregation, type CodexCapacityQuota, } from "./codex-capacity"; +import { fetchAntigravityLiveQuota } from "./antigravity-quota"; +import { antigravityHostCandidates, isAntigravityHttpsHost } from "../adapters/google-antigravity-hosts"; /** Match oauth/index REFRESH_SKEW_MS — use stored access without refresh when still fresh. */ const ACCOUNT_TOKEN_SKEW_MS = 60_000; @@ -2009,39 +2011,72 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig return null; } const baseUrl = (config.baseUrl || "https://daily-cloudcode-pa.googleapis.com").replace(/\/+$/, ""); - const response = await fetch(`${baseUrl}/v1internal:fetchAvailableModels`, { - method: "POST", - headers: { - Accept: "application/json", - "Content-Type": "application/json", - "User-Agent": antigravityUserAgent(), - Authorization: `Bearer ${accessToken}`, - }, - body: JSON.stringify({ project: credential.projectId }), - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + const liveQuota = await fetchAntigravityLiveQuota({ + accessToken, + projectId: credential.projectId, + baseUrl, + timeoutMs: REQUEST_TIMEOUT_MS, }); - if (!response.ok) return null; - const body = asRecord(await readQuotaJson(response)); - const models = asRecord(body?.models); - if (!models) return null; const windows = new Map(); - for (const [modelId, rawModelInfo] of Object.entries(models)) { - const modelInfo = asRecord(rawModelInfo); - if (!modelInfo) continue; - for (const quotaInfo of quotaInfoEntries(modelInfo)) { - const label = classifyAntigravityFamily(modelId, modelInfo, quotaInfo); - if (!label || windows.has(label)) continue; - const percent = antigravityUsedPercent(quotaInfo); - if (percent === undefined) continue; - windows.set(label, { - label, - percent, - ...(normalizeResetAt(quotaInfo.resetTime) !== undefined ? { resetAt: normalizeResetAt(quotaInfo.resetTime) } : {}), + for (const [index, host] of antigravityHostCandidates(baseUrl).entries()) { + if (!isAntigravityHttpsHost(host)) continue; + try { + const response = await fetch(`${host}/v1internal:fetchAvailableModels`, { + method: "POST", + headers: { + Accept: "application/json", + "Content-Type": "application/json", + "User-Agent": antigravityUserAgent(), + Authorization: `Bearer ${accessToken}`, + }, + body: JSON.stringify({ project: credential.projectId }), + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), }); + if (!response.ok) { + if (index === 0 && (response.status === 404 || response.status === 503)) continue; + break; + } + const body = asRecord(await readQuotaJson(response)); + const models = asRecord(body?.models); + if (models) { + for (const [modelId, rawModelInfo] of Object.entries(models)) { + const modelInfo = asRecord(rawModelInfo); + if (!modelInfo) continue; + for (const quotaInfo of quotaInfoEntries(modelInfo)) { + const label = classifyAntigravityFamily(modelId, modelInfo, quotaInfo); + if (!label || windows.has(label)) continue; + const percent = antigravityUsedPercent(quotaInfo); + if (percent === undefined) continue; + windows.set(label, { + label, + percent, + ...(normalizeResetAt(quotaInfo.resetTime) !== undefined ? { resetAt: normalizeResetAt(quotaInfo.resetTime) } : {}), + }); + } + } + } + break; + } catch { + if (index === 0) continue; + break; } } + if (liveQuota) { + const liveWindows = liveQuota.customWindows ?? []; + const catalogClaude = windows.get("Cla"); + const customWindows = [ + ...liveWindows, + ...(liveWindows.some(window => window.label === "Cla") || !catalogClaude ? [] : [catalogClaude]), + ]; + return report(provider, "google-antigravity:retrieveUserQuota", { + ...liveQuota, + ...(customWindows.length > 0 ? { customWindows } : {}), + updatedAt: Date.now(), + }); + } + const customWindows = ["Gem", "Cla"].flatMap(label => { const window = windows.get(label); return window ? [window] : []; diff --git a/tests/antigravity-quota.test.ts b/tests/antigravity-quota.test.ts new file mode 100644 index 0000000000..2564a90924 --- /dev/null +++ b/tests/antigravity-quota.test.ts @@ -0,0 +1,319 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { fetchAntigravityLiveQuota } from "../src/providers/antigravity-quota"; +import { + clearProviderQuotaCache, + fetchProviderQuotaReports, + QUOTA_RESPONSE_MAX_BYTES, +} from "../src/providers/quota"; +import { saveCredential } from "../src/oauth/store"; +import type { OcxConfig } from "../src/types"; + +const originalFetch = globalThis.fetch; +const previousOpencodexHome = process.env.OPENCODEX_HOME; +let opencodexHome: string; + +const DAILY_HOST = "https://daily-cloudcode-pa.googleapis.com"; +const PROD_HOST = "https://cloudcode-pa.googleapis.com"; +const TOKEN = "antigravity-access-token"; +const PROJECT = "antigravity-project"; + +function liveGeminiQuota(): Response { + return jsonResponse({ + buckets: [ + { modelId: "gemini-3.6-pro", remainingFraction: 0.4, resetTime: "2026-08-19T12:00:00Z" }, + ], + }); +} + +function liveWeeklySummary(): Response { + return jsonResponse({ + weekly: { remainingPercentage: 75, resetTime: "2026-08-25T00:00:00Z" }, + }); +} + +function config(baseUrl = DAILY_HOST): OcxConfig { + return { + defaultProvider: "google-antigravity", + providers: { + "google-antigravity": { adapter: "google", authMode: "oauth", baseUrl }, + }, + } as OcxConfig; +} + +function jsonResponse(value: unknown, status = 200): Response { + return new Response(JSON.stringify(value), { + status, + headers: { "content-type": "application/json" }, + }); +} + +function catalogResponse(): Response { + return jsonResponse({ + models: { + "gemini-3.6-flash-medium": { + displayName: "Gemini 3.6 Flash (Medium)", + quotaInfo: { remainingFraction: 0.64, resetTime: "2026-08-20T14:00:00Z" }, + }, + "claude-sonnet-4.6": { + displayName: "Claude Sonnet", + quotaInfo: { remainingFraction: 0.21, resetTime: "2026-08-21T15:00:00Z" }, + }, + }, + }); +} + +function oversizedJsonResponse(value: Record): Response { + return new Response(JSON.stringify({ + ...value, + padding: "x".repeat(QUOTA_RESPONSE_MAX_BYTES), + }), { + status: 200, + headers: { "content-type": "application/json" }, + }); +} + +beforeEach(async () => { + opencodexHome = mkdtempSync(join(tmpdir(), "ocx-antigravity-quota-")); + process.env.OPENCODEX_HOME = opencodexHome; + await saveCredential("google-antigravity", { + access: TOKEN, + refresh: "antigravity-refresh-token", + expires: Date.now() + 3_600_000, + projectId: PROJECT, + }); + clearProviderQuotaCache(); +}); + +afterEach(() => { + globalThis.fetch = originalFetch; + clearProviderQuotaCache(); + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; + rmSync(opencodexHome, { recursive: true, force: true }); +}); + +describe("Antigravity live quota", () => { + test("merges live Gemini and weekly quota with catalog-only Claude windows", async () => { + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return jsonResponse({ + buckets: [ + { modelId: "gemini-3.6-pro", remainingFraction: 0.4, resetTime: "2026-08-19T12:00:00Z" }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return jsonResponse({ + weekly: { remainingPercentage: 75, resetTime: "2026-08-25T00:00:00Z" }, + }); + } + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + const report = result.reports[0]; + + expect(report?.source).toBe("google-antigravity:retrieveUserQuota"); + expect(report?.quota.customWindows).toEqual([ + { label: "Gem", percent: 60, resetAt: Date.parse("2026-08-19T12:00:00Z") }, + { label: "Cla", percent: 79, resetAt: Date.parse("2026-08-21T15:00:00Z") }, + ]); + expect(report?.quota.weeklyPercent).toBe(25); + expect(report?.quota.weeklyResetAt).toBe(Date.parse("2026-08-25T00:00:00Z")); + }); + + test("retries the production host after daily retrieveUserQuota returns 404", async () => { + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url.startsWith(DAILY_HOST) && url.includes(":retrieveUserQuota")) return jsonResponse({}, 404); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + + expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); + expect(requested).toContain(`${PROD_HOST}/v1internal:retrieveUserQuota`); + expect(result.reports[0]?.source).toBe("google-antigravity:retrieveUserQuota"); + }); + + test("falls back to the catalog when both live RPCs return 404", async () => { + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes(":retrieveUserQuota")) return jsonResponse({}, 404); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + const report = result.reports[0]; + + expect(report?.source).toBe("google-antigravity:fetchAvailableModels"); + expect(report?.quota.customWindows).toEqual([ + { label: "Gem", percent: 36, resetAt: Date.parse("2026-08-20T14:00:00Z") }, + { label: "Cla", percent: 79, resetAt: Date.parse("2026-08-21T15:00:00Z") }, + ]); + expect(report?.quota.weeklyPercent).toBeUndefined(); + }); + + test("falls back to the catalog when live RPC fetch throws", async () => { + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes(":retrieveUserQuota")) throw new Error("simulated timeout"); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + await expect(fetchProviderQuotaReports(config(), true)).resolves.toMatchObject({ + reports: [{ + source: "google-antigravity:fetchAvailableModels", + quota: { customWindows: expect.any(Array) }, + }], + }); + }); + + test("fails open to the catalog when live RPC bodies exceed the quota JSON limit", async () => { + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return oversizedJsonResponse({ + buckets: [ + { modelId: "gemini-3.6-pro", remainingFraction: 0.01 }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return oversizedJsonResponse({ + weekly: { remainingPercentage: 1 }, + }); + } + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + const report = result.reports[0]; + + expect(report?.source).toBe("google-antigravity:fetchAvailableModels"); + expect(report?.quota.customWindows).toEqual([ + { label: "Gem", percent: 36, resetAt: Date.parse("2026-08-20T14:00:00Z") }, + { label: "Cla", percent: 79, resetAt: Date.parse("2026-08-21T15:00:00Z") }, + ]); + expect(report?.quota.weeklyPercent).toBeUndefined(); + }); + + test("does not fetch the production host after daily retrieveUserQuota returns 401", async () => { + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url === `${DAILY_HOST}/v1internal:retrieveUserQuota`) return jsonResponse({}, 401); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + + expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); + expect(requested.filter(url => url.startsWith(PROD_HOST))).toEqual([]); + expect(result.reports[0]?.source).toBe("google-antigravity:fetchAvailableModels"); + }); + + test("does not fetch the production host when daily retrieveUserQuota 401 races a 404 summary", async () => { + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url === `${DAILY_HOST}/v1internal:retrieveUserQuota`) { + await Bun.sleep(20); + return jsonResponse({}, 401); + } + if (url === `${DAILY_HOST}/v1internal:retrieveUserQuotaSummary`) return jsonResponse({}, 404); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + + expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); + expect(requested.filter(url => url.startsWith(PROD_HOST))).toEqual([]); + expect(result.reports[0]?.source).toBe("google-antigravity:fetchAvailableModels"); + }); + + test("does not fetch the production host after daily retrieveUserQuota returns 429", async () => { + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url === `${DAILY_HOST}/v1internal:retrieveUserQuota`) return jsonResponse({}, 429); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + + expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); + expect(requested.filter(url => url.startsWith(PROD_HOST))).toEqual([]); + expect(result.reports[0]?.source).toBe("google-antigravity:fetchAvailableModels"); + }); + + test("does not POST retrieveUserQuota or retrieveUserQuotaSummary to an http host", async () => { + const httpHost = "http://daily-cloudcode-pa.googleapis.com"; + const requested: string[] = []; + const fetchImpl = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const quota = await fetchAntigravityLiveQuota({ + accessToken: TOKEN, + projectId: PROJECT, + baseUrl: httpHost, + timeoutMs: 8_000, + fetchImpl, + }); + + expect(requested.filter(url => url.startsWith("http://"))).toEqual([]); + expect(requested).not.toContain(`${httpHost}/v1internal:retrieveUserQuota`); + expect(requested).not.toContain(`${httpHost}/v1internal:retrieveUserQuotaSummary`); + expect(quota?.customWindows).toEqual([ + { label: "Gem", percent: 60, resetAt: Date.parse("2026-08-19T12:00:00Z") }, + ]); + }); + + test("does not POST fetchAvailableModels to an http host", async () => { + const httpHost = "http://daily-cloudcode-pa.googleapis.com"; + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + await fetchProviderQuotaReports(config(httpHost), true); + + expect(requested.filter(url => url.startsWith("http://"))).toEqual([]); + expect(requested).not.toContain(`${httpHost}/v1internal:fetchAvailableModels`); + }); +}); diff --git a/tests/google-antigravity-errors.test.ts b/tests/google-antigravity-errors.test.ts new file mode 100644 index 0000000000..fccbbf63ad --- /dev/null +++ b/tests/google-antigravity-errors.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, test } from "bun:test"; +import { + isAntigravityGeoBlockedBody, + isQuotaExhaustedBody, + retryableGoogleStatus, + safeAntigravityHttpErrorMessage, +} from "../src/adapters/google-errors"; + +const GEO_BLOCKED_DETAIL = "User location is not supported for the API use"; + +describe("Antigravity Google error classification", () => { + test("classifies geo-blocked 403 responses and redacts echoed credentials", () => { + const leakedToken = "geo-access-token-value-123456"; + const payload = JSON.stringify({ + error: { + status: "PERMISSION_DENIED", + message: `${GEO_BLOCKED_DETAIL}; accessToken=${leakedToken}`, + }, + }); + + expect(isAntigravityGeoBlockedBody(payload)).toBe(true); + const message = safeAntigravityHttpErrorMessage(403, payload); + expect(message).toContain("Antigravity location not supported"); + expect(message).not.toContain(leakedToken); + }); + + test("keeps ordinary permission-denied 403 responses as access denied", () => { + const payload = JSON.stringify({ + error: { + status: "PERMISSION_DENIED", + message: "The caller does not have permission to use this resource", + }, + }); + + expect(isAntigravityGeoBlockedBody(payload)).toBe(false); + expect(safeAntigravityHttpErrorMessage(403, payload)).toContain("Antigravity access denied"); + }); + + test("preserves quota and rate-limit classification for 429 responses", () => { + const quotaPayload = JSON.stringify({ + error: { + status: "RESOURCE_EXHAUSTED", + message: "Quota exceeded for this project", + }, + }); + const rateLimitPayload = JSON.stringify({ + error: { + status: "RESOURCE_EXHAUSTED", + message: "Rate limit exceeded", + }, + }); + + expect(safeAntigravityHttpErrorMessage(429, quotaPayload)).toContain("Antigravity quota exhausted"); + expect(safeAntigravityHttpErrorMessage(429, rateLimitPayload)).toContain("Antigravity rate limit exceeded"); + expect(isQuotaExhaustedBody(quotaPayload)).toBe(true); + expect(isQuotaExhaustedBody(rateLimitPayload)).toBe(false); + expect(retryableGoogleStatus(403)).toBe(false); + }); +}); From f3f8a5115cc66ff7009aea9fdfb94776c2199f51 Mon Sep 17 00:00:00 2001 From: SB Yoon <44089734+yansigit@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:32:11 -0600 Subject: [PATCH 02/10] fix(antigravity): harden quota interpretation and catalog fetch Prevent quota probes from following redirects or promoting daily summaries as weekly usage, and interpret explicit percentage fields without converting small percentages into fractions. Co-authored-by: Cursor --- src/providers/antigravity-quota.ts | 9 ++---- src/providers/quota.ts | 3 +- tests/antigravity-quota.test.ts | 30 +++++++++++++++++++ tests/provider-quota.test.ts | 48 ++++++++++++++++++++++++++++-- 4 files changed, 81 insertions(+), 9 deletions(-) diff --git a/src/providers/antigravity-quota.ts b/src/providers/antigravity-quota.ts index 339df293ab..91d1d75cc9 100644 --- a/src/providers/antigravity-quota.ts +++ b/src/providers/antigravity-quota.ts @@ -57,7 +57,7 @@ function remainingPercent(record: Record): number | undefined { ?? record.remainingPercent ?? record.remaining_percent, ); - if (percentage !== undefined) return normalizePercent(percentage <= 1 ? percentage * 100 : percentage); + if (percentage !== undefined) return normalizePercent(percentage); return undefined; } @@ -111,11 +111,8 @@ function isWeeklyPath(path: string[]): boolean { function parseWeeklyWindow(payload: unknown): { percent: number; resetAt?: number } | undefined { const candidates = collectCandidates(payload); - const ordered = [ - ...candidates.filter(candidate => isWeeklyPath(candidate.path)), - ...candidates.filter(candidate => !isWeeklyPath(candidate.path)), - ]; - for (const candidate of ordered) { + for (const candidate of candidates) { + if (!isWeeklyPath(candidate.path)) continue; const percent = usedPercent(candidate.record); if (percent === undefined) continue; const reset = recordResetAt(candidate.record); diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 7bc646a988..5844407c70 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -1995,7 +1995,7 @@ function antigravityUsedPercent(quotaInfo: Record): number | un const remaining = normalizePercent(toFiniteNumber(quotaInfo.remainingFraction) !== undefined ? toFiniteNumber(quotaInfo.remainingFraction)! * 100 : toFiniteNumber(quotaInfo.remainingPercentage) !== undefined - ? toFiniteNumber(quotaInfo.remainingPercentage)! * 100 + ? toFiniteNumber(quotaInfo.remainingPercentage)! : undefined); if (remaining === undefined) return undefined; return normalizePercent(100 - remaining); @@ -2031,6 +2031,7 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig Authorization: `Bearer ${accessToken}`, }, body: JSON.stringify({ project: credential.projectId }), + redirect: "error", signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), }); if (!response.ok) { diff --git a/tests/antigravity-quota.test.ts b/tests/antigravity-quota.test.ts index 2564a90924..518248c7aa 100644 --- a/tests/antigravity-quota.test.ts +++ b/tests/antigravity-quota.test.ts @@ -1,3 +1,33 @@ +import { expect, test } from "bun:test"; +import { fetchAntigravityLiveQuota } from "../src/providers/antigravity-quota"; + +test("does not classify an unlabelled daily summary window as weekly", async () => { + const result = await fetchAntigravityLiveQuota({ + accessToken: "agy-access-secret", + projectId: "agy-project-secret", + baseUrl: "https://daily-cloudcode-pa.googleapis.com", + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return new Response(JSON.stringify({ + models: { + "gemini-test": { quotaInfo: { remainingFraction: 0.5 } }, + }, + }), { status: 200 }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return new Response(JSON.stringify({ + daily: { remainingFraction: 0.75 }, + }), { status: 200 }); + } + return new Response("not found", { status: 404 }); + }, + }); + + expect(result?.customWindows?.[0]?.label).toBe("Gem"); + expect(result?.weeklyPercent).toBeUndefined(); +}); import { afterEach, beforeEach, describe, expect, test } from "bun:test"; import { mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; diff --git a/tests/provider-quota.test.ts b/tests/provider-quota.test.ts index 1c4e733683..db1bad354a 100644 --- a/tests/provider-quota.test.ts +++ b/tests/provider-quota.test.ts @@ -121,11 +121,16 @@ describe("fetchProviderQuotaReports", () => { await saveCredential("google-antigravity", { access: "agy-access-secret", refresh: "agy-refresh-secret", expires: Date.now() + 3600_000, projectId: "agy-project-secret" }); await saveCredential("kimi", { access: "kimi-access-secret", refresh: "kimi-refresh-secret", expires: Date.now() + 3600_000 }); - const seen: { url: string; authorization?: string; body?: string }[] = []; + const seen: { url: string; authorization?: string; body?: string; redirect?: RequestRedirect }[] = []; globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { const url = String(input); const headers = init?.headers as Record | undefined; - seen.push({ url, authorization: headers?.Authorization, body: typeof init?.body === "string" ? init.body : undefined }); + seen.push({ + url, + authorization: headers?.Authorization, + body: typeof init?.body === "string" ? init.body : undefined, + redirect: init?.redirect, + }); if (url === "https://chatgpt.com/backend-api/wham/usage") { return new Response(JSON.stringify({ email: "person@example.com", @@ -255,9 +260,48 @@ describe("fetchProviderQuotaReports", () => { expect(seen.find(row => row.url.includes("anthropic.com"))?.authorization).toBe("Bearer claude-access-secret"); expect(seen.find(row => row.url.includes("cloudcode-pa.googleapis.com"))?.authorization).toBe("Bearer agy-access-secret"); expect(seen.find(row => row.url.includes("cloudcode-pa.googleapis.com"))?.body).toBe(JSON.stringify({ project: "agy-project-secret" })); + expect(seen.find(row => row.url.endsWith("/v1internal:fetchAvailableModels"))?.redirect).toBe("error"); expect(seen.find(row => row.url === "https://api.kimi.com/coding/v1/usages")?.authorization).toBe("Bearer kimi-access-secret"); }); + test("treats remainingPercentage as a percentage at values one and below", async () => { + await saveCredential("google-antigravity", { + access: "agy-access-secret", + refresh: "agy-refresh-secret", + expires: Date.now() + 3600_000, + projectId: "agy-project-secret", + }); + const config = { + defaultProvider: "google-antigravity", + providers: { + "google-antigravity": { + adapter: "google", + authMode: "oauth", + baseUrl: "https://daily-cloudcode-pa.googleapis.com", + }, + }, + } as OcxConfig; + + for (const [remainingPercentage, expectedUsedPercentage] of [[1, 99], [0.75, 99.25], [75, 25]]) { + clearProviderQuotaCache(); + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes(":retrieveUserQuota")) return new Response("not found", { status: 404 }); + if (url.endsWith("/v1internal:fetchAvailableModels")) { + return new Response(JSON.stringify({ + models: { + "gemini-test": { quotaInfo: { remainingPercentage } }, + }, + }), { status: 200 }); + } + return new Response("not found", { status: 404 }); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config, true); + expect(result.reports[0]?.quota.customWindows?.[0]?.percent).toBe(expectedUsedPercentage); + } + }); + function kimiOnlyConfig(baseUrl = "https://api.kimi.com/coding/v1"): OcxConfig { return { defaultProvider: "kimi", From 23ce6d2b82166150d81d3a878ef92642eeca75ae Mon Sep 17 00:00:00 2001 From: SB Yoon <44089734+yansigit@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:09:10 -0600 Subject: [PATCH 03/10] fix(antigravity): preserve usable quota when summary fails Treat unreadable daily quota JSON as an RPC failure while allowing the optional summary call to fail without discarding the daily window. Co-authored-by: Cursor --- src/providers/antigravity-quota.ts | 22 ++++++------- tests/antigravity-quota.test.ts | 51 ++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 11 deletions(-) diff --git a/src/providers/antigravity-quota.ts b/src/providers/antigravity-quota.ts index 91d1d75cc9..0e4a66d2ff 100644 --- a/src/providers/antigravity-quota.ts +++ b/src/providers/antigravity-quota.ts @@ -122,7 +122,9 @@ function parseWeeklyWindow(payload: unknown): { percent: number; resetAt?: numbe } async function readJson(response: Response, timeoutMs: number): Promise { - return await readProviderQuotaJsonForTests(response, timeoutMs); + const payload = await readProviderQuotaJsonForTests(response, timeoutMs); + if (payload === null) throw new Error("Antigravity quota RPC returned unreadable JSON"); + return payload; } class AntigravityQuotaRpcError extends Error { @@ -167,18 +169,16 @@ async function fetchHostQuota( fetchRpc(fetchImpl, host, "retrieveUserQuota", args), fetchRpc(fetchImpl, host, "retrieveUserQuotaSummary", args), ]); - for (const result of [quotaResult, summaryResult]) { - if ( - result.status === "rejected" - && result.reason instanceof AntigravityQuotaRpcError - && !shouldRetryPeer(result.reason.status) - ) { - throw result.reason; - } + if ( + quotaResult.status === "rejected" + && quotaResult.reason instanceof AntigravityQuotaRpcError + && !shouldRetryPeer(quotaResult.reason.status) + ) { + throw quotaResult.reason; } - if (quotaResult.status === "rejected" || summaryResult.status === "rejected") return null; + if (quotaResult.status === "rejected") return null; const quotaPayload = quotaResult.value; - const summaryPayload = summaryResult.value; + const summaryPayload = summaryResult.status === "fulfilled" ? summaryResult.value : null; const gem = parseGeminiWindow(quotaPayload); const weekly = parseWeeklyWindow(summaryPayload); if (!gem && !weekly) return null; diff --git a/tests/antigravity-quota.test.ts b/tests/antigravity-quota.test.ts index 518248c7aa..aef37eb3fc 100644 --- a/tests/antigravity-quota.test.ts +++ b/tests/antigravity-quota.test.ts @@ -1,5 +1,6 @@ import { expect, test } from "bun:test"; import { fetchAntigravityLiveQuota } from "../src/providers/antigravity-quota"; +import { QUOTA_RESPONSE_MAX_BYTES } from "../src/providers/quota"; test("does not classify an unlabelled daily summary window as weekly", async () => { const result = await fetchAntigravityLiveQuota({ @@ -28,6 +29,56 @@ test("does not classify an unlabelled daily summary window as weekly", async () expect(result?.customWindows?.[0]?.label).toBe("Gem"); expect(result?.weeklyPercent).toBeUndefined(); }); + +test("keeps the daily quota when the summary RPC fails", async () => { + const result = await fetchAntigravityLiveQuota({ + accessToken: "agy-access-secret", + projectId: "agy-project-secret", + baseUrl: "https://daily-cloudcode-pa.googleapis.com", + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return new Response(JSON.stringify({ + buckets: [ + { modelId: "gemini-test", remainingFraction: 0.5 }, + ], + }), { status: 200 }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) return new Response("unavailable", { status: 503 }); + return new Response("not found", { status: 404 }); + }, + }); + + expect(result?.customWindows).toEqual([{ label: "Gem", percent: 50 }]); + expect(result?.weeklyPercent).toBeUndefined(); +}); + +test("treats a daily quota JSON read failure as an RPC failure", async () => { + const result = await fetchAntigravityLiveQuota({ + accessToken: "agy-access-secret", + projectId: "agy-project-secret", + baseUrl: "https://daily-cloudcode-pa.googleapis.com", + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return new Response(JSON.stringify({ + padding: "x".repeat(QUOTA_RESPONSE_MAX_BYTES), + }), { status: 200 }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return new Response(JSON.stringify({ + weekly: { remainingPercentage: 75 }, + }), { status: 200 }); + } + return new Response("not found", { status: 404 }); + }, + }); + + expect(result).toBeNull(); +}); + import { afterEach, beforeEach, describe, expect, test } from "bun:test"; import { mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; From 396bd294b56bd0079e07ff041ed143c417dd6caf Mon Sep 17 00:00:00 2001 From: SB Yoon <44089734+yansigit@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:08:26 -0600 Subject: [PATCH 04/10] test(antigravity): assert quota RPC redirect policy Co-authored-by: Cursor --- tests/antigravity-quota.test.ts | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/antigravity-quota.test.ts b/tests/antigravity-quota.test.ts index aef37eb3fc..b41cd2bd7b 100644 --- a/tests/antigravity-quota.test.ts +++ b/tests/antigravity-quota.test.ts @@ -3,13 +3,15 @@ import { fetchAntigravityLiveQuota } from "../src/providers/antigravity-quota"; import { QUOTA_RESPONSE_MAX_BYTES } from "../src/providers/quota"; test("does not classify an unlabelled daily summary window as weekly", async () => { + const requestOptions: Array<{ url: string; init?: RequestInit }> = []; const result = await fetchAntigravityLiveQuota({ accessToken: "agy-access-secret", projectId: "agy-project-secret", baseUrl: "https://daily-cloudcode-pa.googleapis.com", timeoutMs: 1_000, - fetchImpl: async (input: RequestInfo | URL) => { + fetchImpl: async (input: RequestInfo | URL, init?: RequestInit) => { const url = String(input); + requestOptions.push({ url, init }); if (url.endsWith(":retrieveUserQuota")) { return new Response(JSON.stringify({ models: { @@ -28,6 +30,11 @@ test("does not classify an unlabelled daily summary window as weekly", async () expect(result?.customWindows?.[0]?.label).toBe("Gem"); expect(result?.weeklyPercent).toBeUndefined(); + expect( + requestOptions + .filter(({ url }) => url.includes(":retrieveUserQuota")) + .map(({ init }) => init?.redirect), + ).toEqual(["error", "error"]); }); test("keeps the daily quota when the summary RPC fails", async () => { From 2c459372c4171f7892d7a84f15d4523ddd7a3ab2 Mon Sep 17 00:00:00 2001 From: SB Yoon <44089734+yansigit@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:17:27 -0600 Subject: [PATCH 05/10] fix(antigravity): stop quota probe on terminal RPC errors Restrict host failover to known Google daily/prod endpoints, require explicit model identifiers for Gemini window parsing, and abort live quota plus catalog probing on 401/403/429 instead of leaking tokens to production peers. Co-authored-by: Cursor --- src/adapters/google-antigravity-hosts.ts | 16 +- src/providers/antigravity-quota.ts | 40 ++-- src/providers/quota.ts | 26 ++- tests/antigravity-quota.test.ts | 266 ++++++++++++++--------- 4 files changed, 218 insertions(+), 130 deletions(-) diff --git a/src/adapters/google-antigravity-hosts.ts b/src/adapters/google-antigravity-hosts.ts index b9dc041b17..a2c1fe096d 100644 --- a/src/adapters/google-antigravity-hosts.ts +++ b/src/adapters/google-antigravity-hosts.ts @@ -2,16 +2,18 @@ const DAILY_ANTIGRAVITY_HOST = "https://daily-cloudcode-pa.googleapis.com"; const PROD_ANTIGRAVITY_HOST = "https://cloudcode-pa.googleapis.com"; /** - * Return the configured Antigravity endpoint followed by its daily/production peer. - * The configured value is preserved so tests and future pinned environments keep their - * explicit first choice; the fallback is always one of Google's two known hosts. + * Return the configured Antigravity endpoint and, for Google's known daily/prod hosts + * only, its daily/production peer. Custom baseUrl values stay single-host. */ export function antigravityHostCandidates(configuredBase: string): string[] { const configured = configuredBase.replace(/\/+$/, ""); - const other = configured === DAILY_ANTIGRAVITY_HOST - ? PROD_ANTIGRAVITY_HOST - : DAILY_ANTIGRAVITY_HOST; - return [...new Set([configured, other])]; + if (configured === DAILY_ANTIGRAVITY_HOST) { + return [DAILY_ANTIGRAVITY_HOST, PROD_ANTIGRAVITY_HOST]; + } + if (configured === PROD_ANTIGRAVITY_HOST) { + return [PROD_ANTIGRAVITY_HOST, DAILY_ANTIGRAVITY_HOST]; + } + return [configured]; } /** OAuth bearer requests must not use a cleartext host, even if generic baseUrl config allows http. */ diff --git a/src/providers/antigravity-quota.ts b/src/providers/antigravity-quota.ts index 0e4a66d2ff..8bddf0b5e3 100644 --- a/src/providers/antigravity-quota.ts +++ b/src/providers/antigravity-quota.ts @@ -84,15 +84,19 @@ function collectCandidates(value: unknown, path: string[] = [], output: QuotaCan return output; } -function candidateModelName(candidate: QuotaCandidate): string { - const { record, path } = candidate; - const explicit = record.modelId ?? record.model_id ?? record.modelName ?? record.model ?? record.name; - return `${typeof explicit === "string" ? explicit : ""} ${path.join(" ")}`.toLowerCase(); +function candidateModelName(record: Record): string { + const explicit = record.modelId + ?? record.model_id + ?? record.modelName + ?? record.model + ?? record.name + ?? record.displayName; + return typeof explicit === "string" ? explicit.toLowerCase() : ""; } function parseGeminiWindow(payload: unknown): ProviderQuotaWindow | undefined { for (const candidate of collectCandidates(payload)) { - if (!candidateModelName(candidate).includes("gemini")) continue; + if (!candidateModelName(candidate.record).includes("gemini")) continue; const percent = usedPercent(candidate.record); if (percent === undefined) continue; const reset = recordResetAt(candidate.record); @@ -127,14 +131,19 @@ async function readJson(response: Response, timeoutMs: number): Promise return payload; } -class AntigravityQuotaRpcError extends Error { +export class AntigravityQuotaRpcError extends Error { constructor(readonly status: number) { super(`Antigravity quota RPC failed: ${status}`); } } -function shouldRetryPeer(status: number): boolean { - return status === 404 || status === 503; +export function isTerminalAntigravityQuotaStatus(status: number): boolean { + return status === 401 || status === 403 || status === 429; +} + +function terminalRpcError(result: PromiseSettledResult): AntigravityQuotaRpcError | null { + if (result.status !== "rejected" || !(result.reason instanceof AntigravityQuotaRpcError)) return null; + return isTerminalAntigravityQuotaStatus(result.reason.status) ? result.reason : null; } async function fetchRpc( @@ -169,13 +178,8 @@ async function fetchHostQuota( fetchRpc(fetchImpl, host, "retrieveUserQuota", args), fetchRpc(fetchImpl, host, "retrieveUserQuotaSummary", args), ]); - if ( - quotaResult.status === "rejected" - && quotaResult.reason instanceof AntigravityQuotaRpcError - && !shouldRetryPeer(quotaResult.reason.status) - ) { - throw quotaResult.reason; - } + const terminalError = terminalRpcError(quotaResult) ?? terminalRpcError(summaryResult); + if (terminalError) throw terminalError; if (quotaResult.status === "rejected") return null; const quotaPayload = quotaResult.value; const summaryPayload = summaryResult.status === "fulfilled" ? summaryResult.value : null; @@ -201,8 +205,10 @@ export async function fetchAntigravityLiveQuota( try { const quota = await fetchHostQuota(fetchImpl, host, args); if (quota) return quota; - } catch { - return null; + } catch (error) { + if (error instanceof AntigravityQuotaRpcError && isTerminalAntigravityQuotaStatus(error.status)) { + throw error; + } } } return null; diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 5844407c70..c7d68228c8 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -28,7 +28,11 @@ import { type CodexCapacityAggregation, type CodexCapacityQuota, } from "./codex-capacity"; -import { fetchAntigravityLiveQuota } from "./antigravity-quota"; +import { + AntigravityQuotaRpcError, + fetchAntigravityLiveQuota, + isTerminalAntigravityQuotaStatus, +} from "./antigravity-quota"; import { antigravityHostCandidates, isAntigravityHttpsHost } from "../adapters/google-antigravity-hosts"; /** Match oauth/index REFRESH_SKEW_MS — use stored access without refresh when still fresh. */ @@ -2011,12 +2015,20 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig return null; } const baseUrl = (config.baseUrl || "https://daily-cloudcode-pa.googleapis.com").replace(/\/+$/, ""); - const liveQuota = await fetchAntigravityLiveQuota({ - accessToken, - projectId: credential.projectId, - baseUrl, - timeoutMs: REQUEST_TIMEOUT_MS, - }); + let liveQuota: ProviderQuota | null; + try { + liveQuota = await fetchAntigravityLiveQuota({ + accessToken, + projectId: credential.projectId, + baseUrl, + timeoutMs: REQUEST_TIMEOUT_MS, + }); + } catch (error) { + if (error instanceof AntigravityQuotaRpcError && isTerminalAntigravityQuotaStatus(error.status)) { + return null; + } + liveQuota = null; + } const windows = new Map(); for (const [index, host] of antigravityHostCandidates(baseUrl).entries()) { diff --git a/tests/antigravity-quota.test.ts b/tests/antigravity-quota.test.ts index b41cd2bd7b..c7249dd621 100644 --- a/tests/antigravity-quota.test.ts +++ b/tests/antigravity-quota.test.ts @@ -1,91 +1,3 @@ -import { expect, test } from "bun:test"; -import { fetchAntigravityLiveQuota } from "../src/providers/antigravity-quota"; -import { QUOTA_RESPONSE_MAX_BYTES } from "../src/providers/quota"; - -test("does not classify an unlabelled daily summary window as weekly", async () => { - const requestOptions: Array<{ url: string; init?: RequestInit }> = []; - const result = await fetchAntigravityLiveQuota({ - accessToken: "agy-access-secret", - projectId: "agy-project-secret", - baseUrl: "https://daily-cloudcode-pa.googleapis.com", - timeoutMs: 1_000, - fetchImpl: async (input: RequestInfo | URL, init?: RequestInit) => { - const url = String(input); - requestOptions.push({ url, init }); - if (url.endsWith(":retrieveUserQuota")) { - return new Response(JSON.stringify({ - models: { - "gemini-test": { quotaInfo: { remainingFraction: 0.5 } }, - }, - }), { status: 200 }); - } - if (url.endsWith(":retrieveUserQuotaSummary")) { - return new Response(JSON.stringify({ - daily: { remainingFraction: 0.75 }, - }), { status: 200 }); - } - return new Response("not found", { status: 404 }); - }, - }); - - expect(result?.customWindows?.[0]?.label).toBe("Gem"); - expect(result?.weeklyPercent).toBeUndefined(); - expect( - requestOptions - .filter(({ url }) => url.includes(":retrieveUserQuota")) - .map(({ init }) => init?.redirect), - ).toEqual(["error", "error"]); -}); - -test("keeps the daily quota when the summary RPC fails", async () => { - const result = await fetchAntigravityLiveQuota({ - accessToken: "agy-access-secret", - projectId: "agy-project-secret", - baseUrl: "https://daily-cloudcode-pa.googleapis.com", - timeoutMs: 1_000, - fetchImpl: async (input: RequestInfo | URL) => { - const url = String(input); - if (url.endsWith(":retrieveUserQuota")) { - return new Response(JSON.stringify({ - buckets: [ - { modelId: "gemini-test", remainingFraction: 0.5 }, - ], - }), { status: 200 }); - } - if (url.endsWith(":retrieveUserQuotaSummary")) return new Response("unavailable", { status: 503 }); - return new Response("not found", { status: 404 }); - }, - }); - - expect(result?.customWindows).toEqual([{ label: "Gem", percent: 50 }]); - expect(result?.weeklyPercent).toBeUndefined(); -}); - -test("treats a daily quota JSON read failure as an RPC failure", async () => { - const result = await fetchAntigravityLiveQuota({ - accessToken: "agy-access-secret", - projectId: "agy-project-secret", - baseUrl: "https://daily-cloudcode-pa.googleapis.com", - timeoutMs: 1_000, - fetchImpl: async (input: RequestInfo | URL) => { - const url = String(input); - if (url.endsWith(":retrieveUserQuota")) { - return new Response(JSON.stringify({ - padding: "x".repeat(QUOTA_RESPONSE_MAX_BYTES), - }), { status: 200 }); - } - if (url.endsWith(":retrieveUserQuotaSummary")) { - return new Response(JSON.stringify({ - weekly: { remainingPercentage: 75 }, - }), { status: 200 }); - } - return new Response("not found", { status: 404 }); - }, - }); - - expect(result).toBeNull(); -}); - import { afterEach, beforeEach, describe, expect, test } from "bun:test"; import { mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; @@ -183,10 +95,121 @@ afterEach(() => { rmSync(opencodexHome, { recursive: true, force: true }); }); +test("does not classify an unlabelled daily summary window as weekly", async () => { + const requestOptions: Array<{ url: string; init?: RequestInit }> = []; + const result = await fetchAntigravityLiveQuota({ + accessToken: "agy-access-secret", + projectId: "agy-project-secret", + baseUrl: DAILY_HOST, + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL, init?: RequestInit) => { + const url = String(input); + requestOptions.push({ url, init }); + if (url.endsWith(":retrieveUserQuota")) { + return jsonResponse({ + buckets: [ + { modelId: "gemini-test", remainingFraction: 0.5 }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return jsonResponse({ + daily: { remainingFraction: 0.75 }, + }); + } + return jsonResponse({}, 404); + }, + }); + + expect(result?.customWindows?.[0]?.label).toBe("Gem"); + expect(result?.weeklyPercent).toBeUndefined(); + expect( + requestOptions + .filter(({ url }) => url.includes(":retrieveUserQuota")) + .map(({ init }) => init?.redirect), + ).toEqual(["error", "error"]); +}); + +test("keeps the daily quota when the summary RPC fails", async () => { + const result = await fetchAntigravityLiveQuota({ + accessToken: "agy-access-secret", + projectId: "agy-project-secret", + baseUrl: DAILY_HOST, + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return jsonResponse({ + buckets: [ + { modelId: "gemini-test", remainingFraction: 0.5 }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) return jsonResponse({}, 503); + return jsonResponse({}, 404); + }, + }); + + expect(result?.customWindows).toEqual([{ label: "Gem", percent: 50 }]); + expect(result?.weeklyPercent).toBeUndefined(); +}); + +test("treats a daily quota JSON read failure as an RPC failure", async () => { + const result = await fetchAntigravityLiveQuota({ + accessToken: "agy-access-secret", + projectId: "agy-project-secret", + baseUrl: DAILY_HOST, + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return oversizedJsonResponse({ + buckets: [ + { modelId: "gemini-3.6-pro", remainingFraction: 0.01 }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return jsonResponse({ + weekly: { remainingPercentage: 75 }, + }); + } + return jsonResponse({}, 404); + }, + }); + + expect(result).toBeNull(); +}); + +test("does not parse gemini from ancestor JSON path keys", async () => { + const result = await fetchAntigravityLiveQuota({ + accessToken: TOKEN, + projectId: PROJECT, + baseUrl: DAILY_HOST, + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return jsonResponse({ + "gemini-quotas": { + items: [{ remainingFraction: 0.5 }], + }, + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) return jsonResponse({}, 404); + return jsonResponse({}, 404); + }, + }); + + expect(result?.customWindows).toBeUndefined(); +}); + describe("Antigravity live quota", () => { test("merges live Gemini and weekly quota with catalog-only Claude windows", async () => { - globalThis.fetch = (async (input: RequestInfo | URL) => { + const requestOptions: Array<{ url: string; init?: RequestInit }> = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { const url = String(input); + requestOptions.push({ url, init }); if (url.endsWith(":retrieveUserQuota")) { return jsonResponse({ buckets: [ @@ -213,6 +236,7 @@ describe("Antigravity live quota", () => { ]); expect(report?.quota.weeklyPercent).toBe(25); expect(report?.quota.weeklyResetAt).toBe(Date.parse("2026-08-25T00:00:00Z")); + expect(requestOptions.every(({ init }) => init?.redirect === "error")).toBe(true); }); test("retries the production host after daily retrieveUserQuota returns 404", async () => { @@ -299,7 +323,7 @@ describe("Antigravity live quota", () => { expect(report?.quota.weeklyPercent).toBeUndefined(); }); - test("does not fetch the production host after daily retrieveUserQuota returns 401", async () => { + test("does not fetch production or catalog after daily retrieveUserQuota returns 401", async () => { const requested: string[] = []; globalThis.fetch = (async (input: RequestInfo | URL) => { const url = String(input); @@ -307,7 +331,7 @@ describe("Antigravity live quota", () => { if (url === `${DAILY_HOST}/v1internal:retrieveUserQuota`) return jsonResponse({}, 401); if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); - if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + if (url.endsWith(":fetchAvailableModels")) return jsonResponse({}, 404); return jsonResponse({}, 404); }) as typeof fetch; @@ -315,10 +339,11 @@ describe("Antigravity live quota", () => { expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); expect(requested.filter(url => url.startsWith(PROD_HOST))).toEqual([]); - expect(result.reports[0]?.source).toBe("google-antigravity:fetchAvailableModels"); + expect(requested.filter(url => url.endsWith(":fetchAvailableModels"))).toEqual([]); + expect(result.reports).toEqual([]); }); - test("does not fetch the production host when daily retrieveUserQuota 401 races a 404 summary", async () => { + test("does not fetch production or catalog when daily retrieveUserQuota 401 races a 404 summary", async () => { const requested: string[] = []; globalThis.fetch = (async (input: RequestInfo | URL) => { const url = String(input); @@ -330,7 +355,7 @@ describe("Antigravity live quota", () => { if (url === `${DAILY_HOST}/v1internal:retrieveUserQuotaSummary`) return jsonResponse({}, 404); if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); - if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + if (url.endsWith(":fetchAvailableModels")) return jsonResponse({}, 404); return jsonResponse({}, 404); }) as typeof fetch; @@ -338,10 +363,11 @@ describe("Antigravity live quota", () => { expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); expect(requested.filter(url => url.startsWith(PROD_HOST))).toEqual([]); - expect(result.reports[0]?.source).toBe("google-antigravity:fetchAvailableModels"); + expect(requested.filter(url => url.endsWith(":fetchAvailableModels"))).toEqual([]); + expect(result.reports).toEqual([]); }); - test("does not fetch the production host after daily retrieveUserQuota returns 429", async () => { + test("does not fetch production or catalog after daily retrieveUserQuota returns 429", async () => { const requested: string[] = []; globalThis.fetch = (async (input: RequestInfo | URL) => { const url = String(input); @@ -349,6 +375,26 @@ describe("Antigravity live quota", () => { if (url === `${DAILY_HOST}/v1internal:retrieveUserQuota`) return jsonResponse({}, 429); if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return jsonResponse({}, 503); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(), true); + + expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); + expect(requested.filter(url => url.startsWith(PROD_HOST))).toEqual([]); + expect(requested.filter(url => url.endsWith(":fetchAvailableModels"))).toEqual([]); + expect(result.reports).toEqual([]); + }); + + test("does not fetch production or catalog after daily retrieveUserQuota returns 403", async () => { + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url === `${DAILY_HOST}/v1internal:retrieveUserQuota`) return jsonResponse({}, 403); + if (url.endsWith(":retrieveUserQuota")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); return jsonResponse({}, 404); }) as typeof fetch; @@ -357,6 +403,30 @@ describe("Antigravity live quota", () => { expect(requested).toContain(`${DAILY_HOST}/v1internal:retrieveUserQuota`); expect(requested.filter(url => url.startsWith(PROD_HOST))).toEqual([]); + expect(requested.filter(url => url.endsWith(":fetchAvailableModels"))).toEqual([]); + expect(result.reports).toEqual([]); + }); + + test("does not fail over to daily or prod for a custom baseUrl on 404/503", async () => { + const customHost = "https://custom-proxy.example.com"; + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url.startsWith(customHost) && url.includes(":retrieveUserQuota")) return jsonResponse({}, 404); + if (url.startsWith(customHost) && url.includes(":retrieveUserQuotaSummary")) return jsonResponse({}, 503); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const result = await fetchProviderQuotaReports(config(customHost), true); + + expect(requested.filter(url => url.startsWith(DAILY_HOST) || url.startsWith(PROD_HOST))).toEqual([]); + expect(requested.filter(url => url.startsWith(customHost))).toEqual([ + `${customHost}/v1internal:retrieveUserQuota`, + `${customHost}/v1internal:retrieveUserQuotaSummary`, + `${customHost}/v1internal:fetchAvailableModels`, + ]); expect(result.reports[0]?.source).toBe("google-antigravity:fetchAvailableModels"); }); @@ -382,9 +452,7 @@ describe("Antigravity live quota", () => { expect(requested.filter(url => url.startsWith("http://"))).toEqual([]); expect(requested).not.toContain(`${httpHost}/v1internal:retrieveUserQuota`); expect(requested).not.toContain(`${httpHost}/v1internal:retrieveUserQuotaSummary`); - expect(quota?.customWindows).toEqual([ - { label: "Gem", percent: 60, resetAt: Date.parse("2026-08-19T12:00:00Z") }, - ]); + expect(quota).toBeNull(); }); test("does not POST fetchAvailableModels to an http host", async () => { From 2946115007593444b7fb1163c49065ba5933333a Mon Sep 17 00:00:00 2001 From: SB Yoon <44089734+yansigit@users.noreply.github.com> Date: Wed, 19 Aug 2026 22:48:27 -0600 Subject: [PATCH 06/10] fix(antigravity): drop last-good quota on terminal RPC and classify weekly on the leaf path Co-authored-by: Cursor --- src/providers/antigravity-quota.ts | 3 +- src/providers/quota.ts | 4 +-- tests/antigravity-quota.test.ts | 48 ++++++++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 3 deletions(-) diff --git a/src/providers/antigravity-quota.ts b/src/providers/antigravity-quota.ts index 8bddf0b5e3..c4f0cb58f2 100644 --- a/src/providers/antigravity-quota.ts +++ b/src/providers/antigravity-quota.ts @@ -110,7 +110,8 @@ function parseGeminiWindow(payload: unknown): ProviderQuotaWindow | undefined { } function isWeeklyPath(path: string[]): boolean { - return path.some(part => /weekly|week|seven[_-]?day/i.test(part)); + const leaf = path.at(-1); + return typeof leaf === "string" && /weekly|week|seven[_-]?day/i.test(leaf); } function parseWeeklyWindow(payload: unknown): { percent: number; resetAt?: number } | undefined { diff --git a/src/providers/quota.ts b/src/providers/quota.ts index c7d68228c8..c03e0b804a 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -2005,7 +2005,7 @@ function antigravityUsedPercent(quotaInfo: Record): number | un return normalizePercent(100 - remaining); } -async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig): Promise { +async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig): Promise { const credential = getCredential("google-antigravity"); if (!credential?.projectId) return null; let accessToken: string; @@ -2025,7 +2025,7 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig }); } catch (error) { if (error instanceof AntigravityQuotaRpcError && isTerminalAntigravityQuotaStatus(error.status)) { - return null; + return TERMINAL_QUOTA_FAILURE; } liveQuota = null; } diff --git a/tests/antigravity-quota.test.ts b/tests/antigravity-quota.test.ts index c7249dd621..0a910d959d 100644 --- a/tests/antigravity-quota.test.ts +++ b/tests/antigravity-quota.test.ts @@ -95,6 +95,33 @@ afterEach(() => { rmSync(opencodexHome, { recursive: true, force: true }); }); +test("does not classify a daily bucket nested under a weekly ancestor as weekly", async () => { + const result = await fetchAntigravityLiveQuota({ + accessToken: TOKEN, + projectId: PROJECT, + baseUrl: DAILY_HOST, + timeoutMs: 1_000, + fetchImpl: async (input: RequestInfo | URL) => { + const url = String(input); + if (url.endsWith(":retrieveUserQuota")) { + return jsonResponse({ + buckets: [ + { modelId: "gemini-test", remainingFraction: 0.5 }, + ], + }); + } + if (url.endsWith(":retrieveUserQuotaSummary")) { + return jsonResponse({ + weekly: { daily: { remainingPercentage: 90 } }, + }); + } + return jsonResponse({}, 404); + }, + }); + + expect(result?.weeklyPercent).toBeUndefined(); +}); + test("does not classify an unlabelled daily summary window as weekly", async () => { const requestOptions: Array<{ url: string; init?: RequestInit }> = []; const result = await fetchAntigravityLiveQuota({ @@ -343,6 +370,27 @@ describe("Antigravity live quota", () => { expect(result.reports).toEqual([]); }); + test("drops last-good Antigravity quota after a terminal 401 refresh", async () => { + let rejected = false; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (rejected && url.includes(":retrieveUserQuota") && !url.includes("Summary")) { + return jsonResponse({}, 401); + } + if (url.endsWith(":retrieveUserQuota") && !url.includes("Summary")) return liveGeminiQuota(); + if (url.endsWith(":retrieveUserQuotaSummary")) return liveWeeklySummary(); + if (url.endsWith(":fetchAvailableModels")) return catalogResponse(); + return jsonResponse({}, 404); + }) as typeof fetch; + + const valid = await fetchProviderQuotaReports(config(), true); + rejected = true; + const rejectedRefresh = await fetchProviderQuotaReports(config(), true); + + expect(valid.reports).toHaveLength(1); + expect(rejectedRefresh.reports).toEqual([]); + }); + test("does not fetch production or catalog when daily retrieveUserQuota 401 races a 404 summary", async () => { const requested: string[] = []; globalThis.fetch = (async (input: RequestInfo | URL) => { From 19a6de174b2c2d3560e047e880afecb2c3d643f6 Mon Sep 17 00:00:00 2001 From: SB Yoon <44089734+yansigit@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:08:34 -0600 Subject: [PATCH 07/10] feat(antigravity): process-local account cooldowns Key 429/quota/geoblock cooldowns by OAuth account, fail closed when a rotated Cloud Code Assist credential has no project id, and sweep expired entries with the state store. Co-authored-by: Cursor --- src/adapters/base.ts | 2 + src/lib/state-store-registrations.ts | 2 + src/oauth/antigravity-routing.ts | 132 +++++++++++++++++++++++++ src/server/responses/core.ts | 101 +++++++++++++++++-- tests/antigravity-project-bind.test.ts | 63 ++++++++++++ tests/antigravity-routing.test.ts | 79 +++++++++++++++ tests/state-store-sweeper.test.ts | 1 + 7 files changed, 372 insertions(+), 8 deletions(-) create mode 100644 src/oauth/antigravity-routing.ts create mode 100644 tests/antigravity-project-bind.test.ts create mode 100644 tests/antigravity-routing.test.ts diff --git a/src/adapters/base.ts b/src/adapters/base.ts index faeea0e959..01bad2d1e0 100644 --- a/src/adapters/base.ts +++ b/src/adapters/base.ts @@ -106,6 +106,8 @@ export interface AdapterRequest { export interface AdapterFetchContext { /** Remains attached to the returned response body after the response headers arrive. */ abortSignal?: AbortSignal; + /** OAuth account identity used for provider-local cooldown bookkeeping. */ + accountId?: string; /** Deadline for receiving response headers on each attempt, not for consuming the response body. */ timeoutMs?: number; /** Return final non-2xx responses untouched so the caller can own the error-body read. */ diff --git a/src/lib/state-store-registrations.ts b/src/lib/state-store-registrations.ts index 61aad5f7b9..eb4d9b4215 100644 --- a/src/lib/state-store-registrations.ts +++ b/src/lib/state-store-registrations.ts @@ -31,6 +31,7 @@ import { sweepExpiredXaiPermanentFailureVerdicts, } from "../oauth"; import { sweepExpiredAnthropicRoutingHealth } from "../oauth/anthropic-routing"; +import { sweepExpiredAntigravityRoutingHealth } from "../oauth/antigravity-routing"; import { listLiveOAuthAccountKeys, reconcileOAuthReauthState } from "../oauth/store"; import { reconcileGuardianBackoff } from "../oauth/token-guardian"; import { sweepExpiredApiKeyCooldowns } from "../providers/key-failover"; @@ -83,6 +84,7 @@ export const STATE_STORE_REGISTRATIONS = [ reconcileGeneration: reconcileComboTargetCooldowns, }, { name: "anthropic-routing-health", sweepExpired: sweepExpiredAnthropicRoutingHealth }, + { name: "antigravity-routing-health", sweepExpired: sweepExpiredAntigravityRoutingHealth }, { name: "xai-refresh-verdicts", sweepExpired: sweepExpiredXaiPermanentFailureVerdicts }, { name: "responses-continuation", diff --git a/src/oauth/antigravity-routing.ts b/src/oauth/antigravity-routing.ts new file mode 100644 index 0000000000..b3c0a9af9e --- /dev/null +++ b/src/oauth/antigravity-routing.ts @@ -0,0 +1,132 @@ +export type AntigravityCooldownReason = "rate_limited" | "quota_exhausted" | "geo_blocked"; + +const DEFAULT_RATE_LIMITED_COOLDOWN_MS = 5_000; +const MAX_RATE_LIMITED_COOLDOWN_MS = 60_000; +const DEFAULT_QUOTA_EXHAUSTED_COOLDOWN_MS = 24 * 60 * 60_000; +const MAX_QUOTA_EXHAUSTED_COOLDOWN_MS = 7 * 24 * 60 * 60_000; +const GEO_BLOCKED_COOLDOWN_MS = 24 * 60 * 60_000; + +type AntigravityAccountHealth = { + cooldownUntil: number; +}; + +const accountHealth = new Map(); + +function positiveDurationOrDefault( + durationMs: number | undefined, + defaultMs: number, + maxMs?: number, +): number { + if (typeof durationMs !== "number" || !Number.isFinite(durationMs) || durationMs <= 0) { + return defaultMs; + } + return maxMs === undefined ? durationMs : Math.min(durationMs, maxMs); +} + +function cooldownDurationMs( + reason: AntigravityCooldownReason, + retryAfterMs: number | undefined, +): number { + switch (reason) { + case "rate_limited": + return positiveDurationOrDefault( + retryAfterMs, + DEFAULT_RATE_LIMITED_COOLDOWN_MS, + MAX_RATE_LIMITED_COOLDOWN_MS, + ); + case "quota_exhausted": + return positiveDurationOrDefault( + retryAfterMs, + DEFAULT_QUOTA_EXHAUSTED_COOLDOWN_MS, + MAX_QUOTA_EXHAUSTED_COOLDOWN_MS, + ); + case "geo_blocked": + return GEO_BLOCKED_COOLDOWN_MS; + } +} + +export function recordAntigravityCooldown( + accountId: string, + reason: AntigravityCooldownReason, + retryAfterMs?: number, + now = Date.now(), +): void { + const cooldownUntil = now + cooldownDurationMs(reason, retryAfterMs); + const current = accountHealth.get(accountId); + if (!current || current.cooldownUntil < cooldownUntil) { + accountHealth.set(accountId, { cooldownUntil }); + } +} + +export function isAntigravityAccountInCooldown(accountId: string, now = Date.now()): boolean { + const health = accountHealth.get(accountId); + if (!health) return false; + if (health.cooldownUntil <= now) { + accountHealth.delete(accountId); + return false; + } + return true; +} + +export function nextAntigravityAccount( + accountIds: string[], + activeId: string | undefined, + now = Date.now(), +): string | undefined { + if (accountIds.length === 0) return undefined; + + const activeIndex = activeId === undefined ? -1 : accountIds.indexOf(activeId); + const startIndex = activeIndex < 0 ? 0 : activeIndex + 1; + for (let offset = 0; offset < accountIds.length; offset += 1) { + const accountId = accountIds[(startIndex + offset) % accountIds.length]!; + if (activeId !== undefined && accountId === activeId) continue; + if (!isAntigravityAccountInCooldown(accountId, now)) return accountId; + } + return undefined; +} + +export function sweepExpiredAntigravityRoutingHealth(now = Date.now()): number { + let removed = 0; + for (const [accountId, health] of accountHealth) { + if (health.cooldownUntil > now) continue; + accountHealth.delete(accountId); + removed += 1; + } + return removed; +} + +export function clearAntigravityAccountCooldown(accountId: string): void { + accountHealth.delete(accountId); +} + +export const ANTIGRAVITY_MISSING_PROJECT_MESSAGE = + "Antigravity requires a discovered Cloud Code Assist project id (re-run `ocx login google-antigravity`)."; + +export type BindAntigravityProjectFailure = { + ok: false; + status: 400; + type: "invalid_request_error"; + message: string; +}; + +export type BindAntigravityProjectSuccess = { + ok: true; + provider: T & { project: string }; +}; + +/** Pair Cloud Code Assist `project` with the credential in use. Never keep a previous account's id. */ +export function bindAntigravityProject( + provider: T, + projectId: string | undefined, +): BindAntigravityProjectSuccess | BindAntigravityProjectFailure { + const project = typeof projectId === "string" ? projectId.trim() : ""; + if (!project) { + return { + ok: false, + status: 400, + type: "invalid_request_error", + message: ANTIGRAVITY_MISSING_PROJECT_MESSAGE, + }; + } + return { ok: true, provider: { ...provider, project } }; +} diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 3c7773e4e4..bc58d56e12 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -84,6 +84,12 @@ import { resolveAnthropicAccountForSession, rotateAnthropicAccountOn429, } from "../../oauth/anthropic-routing"; +import { + bindAntigravityProject, + isAntigravityAccountInCooldown, + nextAntigravityAccount, +} from "../../oauth/antigravity-routing"; +import { getAccountCredential, getAccountSet, setActiveAccount } from "../../oauth/store"; import { buildWebSearchTool, planWebSearch, runWithWebSearch, shouldResolveOpenAiWebSearchSidecar } from "../../web-search"; import { buildImageTool, buildVideoTool, planImageBridge, planVideoBridge, runWithImageBridge, clampImageMaxRounds, IMAGE_GEN_TOOL_NAME, VIDEO_GEN_TOOL_NAME } from "../../images"; import { describeImagesInPlace, isModelTextOnly, planVisionSidecar, resolveOpenAiVisionModel, shouldResolveOpenAiVisionSidecar, stripImagesInPlace } from "../../vision"; @@ -2246,6 +2252,8 @@ async function handleResponsesInner( let replayOAuthCredentialSnapshot: Pick | undefined; let anthropicPoolAccountId: string | null = null; let anthropicPoolFailovers = 0; + let antigravityAccountId: string | undefined; + let antigravityFailovers = 0; const anthropicSessionKey = route.providerName === "anthropic" && route.provider.authMode === "oauth" ? anthropicSessionKeyFromParts({ sessionIdHeader: sessionIdHeaderFromRequest(req.headers), @@ -2278,25 +2286,50 @@ async function handleResponsesInner( route.provider = { ...route.provider, apiKey: accessToken }; logCtx.provider = formatAnthropicProviderForLog("anthropic", selection.accountId, config); } else { - const resolved = await getValidAccessTokenSnapshot(route.providerName); + let resolved = await getValidAccessTokenSnapshot(route.providerName); + let skippedAntigravityCooldown = false; + if (route.providerName === "google-antigravity" + && route.provider.googleMode === "cloud-code-assist" + && isAntigravityAccountInCooldown(resolved.accountId)) { + const accountIds = getAccountSet("google-antigravity")?.accounts.map(account => account.id) ?? []; + const nextAccountId = nextAntigravityAccount(accountIds, resolved.accountId); + if (!nextAccountId) { + return formatErrorResponse(429, "rate_limit_error", "All Google Antigravity OAuth accounts are temporarily unavailable"); + } + const accessToken = await getValidAccessTokenForAccount("google-antigravity", nextAccountId); + const nextCredential = getAccountCredential("google-antigravity", nextAccountId); + resolved = { + ...resolved, + accountId: nextAccountId, + accessToken, + // Always replace; omitting a missing id would keep the previous account's projectId. + projectId: nextCredential?.projectId, + }; + skippedAntigravityCooldown = true; + void setActiveAccount("google-antigravity", nextAccountId).catch(() => { /* best-effort promotion */ }); + } replayOAuthCredentialSnapshot = { accountId: resolved.accountId, generation: resolved.generation, }; + if (skippedAntigravityCooldown) replayOAuthCredentialSnapshot = undefined; if (isOAuth401ReplayProvider) sentOAuthSnapshot = resolved; route.provider = { ...route.provider, apiKey: resolved.accessToken }; + if (route.providerName === "google-antigravity" && route.provider.googleMode === "cloud-code-assist") { + antigravityAccountId = resolved.accountId; + // Always overwrite `project` from the credential in use. A missing id fails closed + // so a rotated account cannot inherit the previous account's Cloud Code Assist project. + const bound = bindAntigravityProject(route.provider, resolved.projectId); + if (!bound.ok) { + return formatErrorResponse(bound.status, bound.type, bound.message); + } + route.provider = bound.provider; + } if (route.providerName === "kiro") { // `{}` is intentional: this is an account-scoped request with no stored routing metadata. // Only genuinely accountless adapter calls leave the context undefined and use local/env fallback. parsed._kiroAuthContext = { ...(resolved.kiro ?? {}) }; } - // Antigravity (cloud-code-assist) needs the discovered Cloud Code Assist project id in the - // CCA envelope. Keep it paired with the token snapshot so an account rotation cannot mix - // a fresh token with project metadata re-read from a different credential generation. - if (route.provider.googleMode === "cloud-code-assist" && !route.provider.project) { - const projectId = resolved.projectId; - if (projectId) route.provider = { ...route.provider, project: projectId }; - } } } catch (err) { if (err instanceof UnsupportedOAuthProviderError) { @@ -3934,6 +3967,7 @@ async function handleResponsesInner( providerName: route.providerName, modelId: route.modelId, }), + ...(antigravityAccountId ? { accountId: antigravityAccountId } : {}), }); } else { // #1851 scope guard: transient-5xx retry on this generic adapter path is opt-in for @@ -4038,6 +4072,7 @@ async function handleResponsesInner( providerName: route.providerName, modelId: route.modelId, }), + ...(antigravityAccountId ? { accountId: antigravityAccountId } : {}), }); } return await fetchWithHeaderTimeout(retryRequest.url, { @@ -4215,6 +4250,55 @@ async function handleResponsesInner( break; } } + // Antigravity OAuth accounts use the same bounded pre-stream carousel as Anthropic, but + // their process-local routing module also excludes accounts cooled by quota/rate-limit + // responses. Geoblocked 403s intentionally never enter this loop. + while ( + upstreamResponse.status === 429 + && route.providerName === "google-antigravity" + && route.provider.googleMode === "cloud-code-assist" + && antigravityAccountId + && antigravityFailovers < 3 + ) { + const accountIds = getAccountSet("google-antigravity")?.accounts.map(account => account.id) ?? []; + const nextAccountId = nextAntigravityAccount(accountIds, antigravityAccountId); + if (!nextAccountId) break; + try { void upstreamResponse.body?.cancel().catch(() => {}); } catch { /* already consumed/closed */ } + try { + const accessToken = await getValidAccessTokenForAccount("google-antigravity", nextAccountId); + const nextCredential = getAccountCredential("google-antigravity", nextAccountId); + const bound = bindAntigravityProject( + { ...route.provider, apiKey: accessToken }, + nextCredential?.projectId, + ); + if (!bound.ok) { + return formatErrorResponse(bound.status, bound.type, bound.message); + } + antigravityAccountId = nextAccountId; + antigravityFailovers += 1; + route.provider = bound.provider; + replayOAuthCredentialSnapshot = undefined; + invalidateSameTargetRequest(); + activeAdapter = resolveAdapter( + resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), + config.cacheRetention, + ); + bindRouteReasoningReplayScope({ + parsed, + providerName: route.providerName, + provider: route.provider, + adapterName: activeAdapter.name, + codexAuthContext: authCtx, + forwardHeaders: selectedForwardHeaders, + }); + void setActiveAccount("google-antigravity", nextAccountId).catch(() => { /* best-effort promotion */ }); + const result = await rebuildAndRefetch("rate-limit-429"); + if ("failed" in result) return result.failed; + upstreamResponse = result; + } catch { + break; + } + } // Anthropic 413 request_too_large: rebuild once with every image one tier lower // (spiral guard: single attempt). The biased response re-enters the 429 check above. if (shouldAttemptImageTierRetry({ @@ -4369,6 +4453,7 @@ async function handleResponsesInner( providerName: route.providerName, modelId: nextParsed.modelId, }), + ...(antigravityAccountId ? { accountId: antigravityAccountId } : {}), }); } // Same #1851 scope guard as the initial send: transient-5xx retry only for direct diff --git a/tests/antigravity-project-bind.test.ts b/tests/antigravity-project-bind.test.ts new file mode 100644 index 0000000000..3238594108 --- /dev/null +++ b/tests/antigravity-project-bind.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, test } from "bun:test"; +import { bindAntigravityProject } from "../src/oauth/antigravity-routing"; + +const MISSING_PROJECT_MESSAGE = + "Antigravity requires a discovered Cloud Code Assist project id (re-run `ocx login google-antigravity`)."; + +describe("bindAntigravityProject", () => { + test("fails closed when the current credential has no projectId", () => { + const previous = { + apiKey: "token-a", + googleMode: "cloud-code-assist" as const, + project: "project-from-previous-account", + }; + + const bound = bindAntigravityProject(previous, undefined); + + expect(bound.ok).toBe(false); + if (bound.ok) throw new Error("expected fail-closed bind"); + expect(bound.status).toBe(400); + expect(bound.type).toBe("invalid_request_error"); + expect(bound.message).toBe(MISSING_PROJECT_MESSAGE); + expect(previous.project).toBe("project-from-previous-account"); + }); + + test("fails closed for an empty projectId instead of keeping the previous project", () => { + const previous = { project: "project-from-previous-account" }; + + const bound = bindAntigravityProject(previous, ""); + + expect(bound.ok).toBe(false); + if (bound.ok) throw new Error("expected fail-closed bind"); + expect(bound.status).toBe(400); + expect(bound.type).toBe("invalid_request_error"); + expect(bound.message).toBe(MISSING_PROJECT_MESSAGE); + expect(previous.project).toBe("project-from-previous-account"); + }); + + test("overwrites a previous account project with the current credential project", () => { + const previous = { + apiKey: "token-b", + googleMode: "cloud-code-assist" as const, + project: "project-from-previous-account", + }; + + const bound = bindAntigravityProject(previous, "project-from-current-account"); + + expect(bound.ok).toBe(true); + if (!bound.ok) throw new Error("expected successful bind"); + expect(bound.provider.project).toBe("project-from-current-account"); + expect(bound.provider.apiKey).toBe("token-b"); + expect(previous.project).toBe("project-from-previous-account"); + }); + + test("assigns the current credential project when the provider had none", () => { + const previous = { apiKey: "token-c", googleMode: "cloud-code-assist" as const }; + + const bound = bindAntigravityProject(previous, "project-from-current-account"); + + expect(bound.ok).toBe(true); + if (!bound.ok) throw new Error("expected successful bind"); + expect(bound.provider.project).toBe("project-from-current-account"); + }); +}); diff --git a/tests/antigravity-routing.test.ts b/tests/antigravity-routing.test.ts new file mode 100644 index 0000000000..19fe8506c3 --- /dev/null +++ b/tests/antigravity-routing.test.ts @@ -0,0 +1,79 @@ +import { afterEach, describe, expect, test } from "bun:test"; +import { + clearAntigravityAccountCooldown, + isAntigravityAccountInCooldown, + nextAntigravityAccount, + recordAntigravityCooldown, + sweepExpiredAntigravityRoutingHealth, +} from "../src/oauth/antigravity-routing"; + +const NOW = 1_700_000_000_000; +const ACCOUNT_IDS = ["account-a", "account-b", "account-c"]; + +afterEach(() => { + for (const accountId of ACCOUNT_IDS) clearAntigravityAccountCooldown(accountId); +}); + +describe("Antigravity account cooldowns", () => { + test("records a rate-limit cooldown and sweeps it after expiry", () => { + recordAntigravityCooldown("account-a", "rate_limited", undefined, NOW); + + expect(isAntigravityAccountInCooldown("account-a", NOW)).toBe(true); + expect(isAntigravityAccountInCooldown("account-a", NOW + 4_999)).toBe(true); + expect(isAntigravityAccountInCooldown("account-a", NOW + 5_000)).toBe(false); + recordAntigravityCooldown("account-b", "rate_limited", undefined, NOW); + expect(sweepExpiredAntigravityRoutingHealth(NOW + 5_000)).toBe(1); + expect(sweepExpiredAntigravityRoutingHealth(NOW + 5_000)).toBe(0); + }); + + test("caps rate-limit Retry-After and uses a parsed quota reset", () => { + recordAntigravityCooldown("account-a", "rate_limited", 120_000, NOW); + recordAntigravityCooldown("account-b", "quota_exhausted", 30_000, NOW); + + expect(isAntigravityAccountInCooldown("account-a", NOW + 60_000)).toBe(false); + expect(isAntigravityAccountInCooldown("account-b", NOW + 29_999)).toBe(true); + expect(isAntigravityAccountInCooldown("account-b", NOW + 30_000)).toBe(false); + }); + + test("honors a quota reset longer than 24 hours", () => { + const resetDurationMs = 48 * 60 * 60_000; + recordAntigravityCooldown("account-a", "quota_exhausted", resetDurationMs, NOW); + + expect(isAntigravityAccountInCooldown("account-a", NOW + 24 * 60 * 60_000 + 1)).toBe(true); + expect(isAntigravityAccountInCooldown("account-a", NOW + resetDurationMs)).toBe(false); + }); + + test("caps quota-exhausted Retry-After at 7 days", () => { + const tenYearsMs = 10 * 365 * 24 * 60 * 60_000; + const sevenDaysMs = 7 * 24 * 60 * 60_000; + recordAntigravityCooldown("account-a", "quota_exhausted", tenYearsMs, NOW); + + expect(isAntigravityAccountInCooldown("account-a", NOW + sevenDaysMs - 1)).toBe(true); + expect(isAntigravityAccountInCooldown("account-a", NOW + sevenDaysMs)).toBe(false); + }); + + test("skips cooled accounts when selecting the next account", () => { + recordAntigravityCooldown("account-b", "rate_limited", undefined, NOW); + + expect(nextAntigravityAccount(ACCOUNT_IDS, "account-a", NOW)).toBe("account-c"); + expect(nextAntigravityAccount(ACCOUNT_IDS, "account-c", NOW)).toBe("account-a"); + expect(nextAntigravityAccount(ACCOUNT_IDS, undefined, NOW)).toBe("account-a"); + }); + + test("keeps a geo block out of the short retry-limit path", () => { + recordAntigravityCooldown("account-a", "geo_blocked", undefined, NOW); + + expect(nextAntigravityAccount(["account-a"], "account-a", NOW + 60_000)).toBeUndefined(); + expect(isAntigravityAccountInCooldown("account-a", NOW + 60_000)).toBe(true); + expect(isAntigravityAccountInCooldown("account-a", NOW + 24 * 60 * 60_000)).toBe(false); + }); + + test("retains the longest expiry from concurrent cooldown records", () => { + recordAntigravityCooldown("account-a", "rate_limited", 60_000, NOW); + recordAntigravityCooldown("account-a", "rate_limited", undefined, NOW + 1); + + expect(isAntigravityAccountInCooldown("account-a", NOW + 5_001)).toBe(true); + expect(isAntigravityAccountInCooldown("account-a", NOW + 59_999)).toBe(true); + expect(isAntigravityAccountInCooldown("account-a", NOW + 60_001)).toBe(false); + }); +}); diff --git a/tests/state-store-sweeper.test.ts b/tests/state-store-sweeper.test.ts index ea2d473b27..364faeb623 100644 --- a/tests/state-store-sweeper.test.ts +++ b/tests/state-store-sweeper.test.ts @@ -100,6 +100,7 @@ describe("state-store sweeper", () => { "provider-request-pacing", "combo-target-cooldowns", "anthropic-routing-health", + "antigravity-routing-health", "xai-refresh-verdicts", "responses-continuation", "antigravity-replay", From 6a1ff3f536d163fb7e4775461724336ce7353156 Mon Sep 17 00:00:00 2001 From: SB Yoon <44089734+yansigit@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:33:05 -0600 Subject: [PATCH 08/10] fix(antigravity): bind replacement before active promotion Validate a cooldown-selected account's Cloud Code Assist project before promoting it, so missing project metadata cannot change the active account state. Co-authored-by: Cursor --- src/server/responses/core.ts | 4 +++- tests/antigravity-project-bind.test.ts | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index bc58d56e12..e129de0d07 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -2306,7 +2306,6 @@ async function handleResponsesInner( projectId: nextCredential?.projectId, }; skippedAntigravityCooldown = true; - void setActiveAccount("google-antigravity", nextAccountId).catch(() => { /* best-effort promotion */ }); } replayOAuthCredentialSnapshot = { accountId: resolved.accountId, @@ -2324,6 +2323,9 @@ async function handleResponsesInner( return formatErrorResponse(bound.status, bound.type, bound.message); } route.provider = bound.provider; + if (skippedAntigravityCooldown) { + void setActiveAccount("google-antigravity", resolved.accountId).catch(() => { /* best-effort promotion */ }); + } } if (route.providerName === "kiro") { // `{}` is intentional: this is an account-scoped request with no stored routing metadata. diff --git a/tests/antigravity-project-bind.test.ts b/tests/antigravity-project-bind.test.ts index 3238594108..8c4e0df216 100644 --- a/tests/antigravity-project-bind.test.ts +++ b/tests/antigravity-project-bind.test.ts @@ -1,4 +1,5 @@ import { describe, expect, test } from "bun:test"; +import { readFileSync } from "node:fs"; import { bindAntigravityProject } from "../src/oauth/antigravity-routing"; const MISSING_PROJECT_MESSAGE = @@ -60,4 +61,17 @@ describe("bindAntigravityProject", () => { if (!bound.ok) throw new Error("expected successful bind"); expect(bound.provider.project).toBe("project-from-current-account"); }); + + test("binds a cooldown replacement before promoting it active", () => { + const source = readFileSync(new URL("../src/server/responses/core.ts", import.meta.url), "utf8"); + const start = source.indexOf("let skippedAntigravityCooldown = false;"); + const end = source.indexOf('if (route.providerName === "kiro")', start); + const initialSelection = source.slice(start, end); + + expect(start).toBeGreaterThanOrEqual(0); + expect(end).toBeGreaterThan(start); + expect(initialSelection.indexOf("bindAntigravityProject")).toBeLessThan( + initialSelection.indexOf('setActiveAccount("google-antigravity"'), + ); + }); }); From 8188771e9850a8d7bb2b092434e20865695afedc Mon Sep 17 00:00:00 2001 From: SB Yoon <44089734+yansigit@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:12:20 -0600 Subject: [PATCH 09/10] fix(antigravity): make account cooldowns effective for HTTP failures Record provider-attributed geo blocks, quota exhaustion, and rate limits at the retry boundary so account rotation can actually avoid unhealthy accounts. Co-authored-by: Cursor --- src/adapters/google-http.ts | 39 +++++++++++++++++++++++++++++++- tests/google-vertex-http.test.ts | 24 +++++++++++++++++++- 2 files changed, 61 insertions(+), 2 deletions(-) diff --git a/src/adapters/google-http.ts b/src/adapters/google-http.ts index f7b90de87e..f60c15aec0 100644 --- a/src/adapters/google-http.ts +++ b/src/adapters/google-http.ts @@ -1,7 +1,13 @@ import type { AdapterFetchContext, AdapterRequest } from "./base"; -import { isQuotaExhaustedBody, retryableGoogleStatus, safeGoogleHttpErrorMessage } from "./google-errors"; +import { + isAntigravityGeoBlockedBody, + isQuotaExhaustedBody, + retryableGoogleStatus, + safeGoogleHttpErrorMessage, +} from "./google-errors"; import { repairGoogleInvalidRequestBody } from "./google-wire-compiler"; import { normalizeUpstreamHttpErrorResponse, readDisplaySafeErrorPayloadText } from "./upstream-http-error"; +import { recordAntigravityCooldown } from "../oauth/antigravity-routing"; import { abortError, cancelResponseBodyBestEffort, @@ -26,6 +32,34 @@ async function normalizeFinalGoogleError(label: string, res: Response, signal?: }); } +function retryAfterMs(value: string | null, now = Date.now()): number | undefined { + const text = value?.trim(); + if (!text) return undefined; + if (/^\d+(?:\.\d+)?$/.test(text)) { + const seconds = Number(text); + return Number.isFinite(seconds) && seconds > 0 ? Math.ceil(seconds * 1000) : undefined; + } + const timestamp = Date.parse(text); + return Number.isFinite(timestamp) && timestamp > now ? timestamp - now : undefined; +} + +async function recordAntigravityHttpCooldown( + response: Response, + accountId: string | undefined, +): Promise { + if (!accountId || (response.status !== 429 && response.status !== 403)) return; + const payloadText = await readDisplaySafeErrorPayloadText(response.clone()); + if (response.status === 429) { + recordAntigravityCooldown( + accountId, + isQuotaExhaustedBody(payloadText) ? "quota_exhausted" : "rate_limited", + retryAfterMs(response.headers.get("retry-after")), + ); + } else if (isAntigravityGeoBlockedBody(payloadText)) { + recordAntigravityCooldown(accountId, "geo_blocked"); + } +} + /** * Fetch a Google-family upstream with Kiro-style hardening: per-attempt timeout * (`AbortSignal.any([parent, timeout])`), bounded retry on transient status / network errors, @@ -53,6 +87,9 @@ export async function fetchGoogleWithRetry( headers: activeRequest.headers, body: activeRequest.body, }, timeoutMs, ctx.abortSignal, ctx.stream, executor); + if (label === "Antigravity") { + await recordAntigravityHttpCooldown(res, ctx.accountId); + } if (res.status === 400 && repairInvalid400 && !compatibilityReplayUsed) { let payloadText = ""; try { diff --git a/tests/google-vertex-http.test.ts b/tests/google-vertex-http.test.ts index 496fb56a1a..9c0cdb206a 100644 --- a/tests/google-vertex-http.test.ts +++ b/tests/google-vertex-http.test.ts @@ -2,9 +2,13 @@ import { afterEach, describe, expect, test } from "bun:test"; import type { AdapterRequest } from "../src/adapters/base"; import { fetchAntigravityWithRetry, fetchDirectGeminiWithRetry, fetchVertexWithRetry } from "../src/adapters/google-http"; import { safeVertexHttpErrorMessage, retryableGoogleStatus } from "../src/adapters/google-errors"; +import { clearAntigravityAccountCooldown, isAntigravityAccountInCooldown } from "../src/oauth/antigravity-routing"; const realFetch = globalThis.fetch; -afterEach(() => { globalThis.fetch = realFetch; }); +afterEach(() => { + globalThis.fetch = realFetch; + clearAntigravityAccountCooldown("account-http-geo"); +}); const request: AdapterRequest = { url: "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/publishers/google/models/gemini-3-pro:streamGenerateContent?alt=sse", @@ -205,6 +209,24 @@ describe("vertex retry fetch", () => { expect(await res403.text()).toContain("Vertex AI access denied"); }); + test("records an Antigravity geo-block cooldown from its HTTP response", async () => { + const mock = mockFetch([ + new Response( + vertexError(403, "PERMISSION_DENIED", "User location is not supported for the API use"), + { status: 403 }, + ), + ]); + + const res = await fetchAntigravityWithRetry(request, { + timeoutMs: 5_000, + accountId: "account-http-geo", + }); + + expect(res.status).toBe(403); + expect(mock.calls).toHaveLength(1); + expect(isAntigravityAccountInCooldown("account-http-geo")).toBe(true); + }); + test("aborts promptly when the caller signal fires", async () => { mockFetch([new Response(vertexError(503, "UNAVAILABLE", "x"), { status: 503, headers: { "Retry-After": "30" } }), new Response("ok", { status: 200 })]); const controller = new AbortController(); From 442cb2d23ae5555e6989c6df033d5ecb721291d9 Mon Sep 17 00:00:00 2001 From: SB Yoon <44089734+yansigit@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:56:57 -0600 Subject: [PATCH 10/10] fix(antigravity): record cooldown reasons and fail fast on HTTP 429/403 Store AntigravityCooldownReason on account health, expose getAntigravityAccountCooldown, and stop retrying the same account inside fetchGoogleWithRetry after recording cooldowns so outer account rotation can switch credentials immediately. Co-authored-by: Cursor --- src/adapters/google-http.ts | 15 +++- src/oauth/antigravity-routing.ts | 21 ++++- tests/antigravity-routing.test.ts | 130 ++++++++++++++++++++++++++++++ 3 files changed, 161 insertions(+), 5 deletions(-) diff --git a/src/adapters/google-http.ts b/src/adapters/google-http.ts index f60c15aec0..9650a172ce 100644 --- a/src/adapters/google-http.ts +++ b/src/adapters/google-http.ts @@ -46,8 +46,8 @@ function retryAfterMs(value: string | null, now = Date.now()): number | undefine async function recordAntigravityHttpCooldown( response: Response, accountId: string | undefined, -): Promise { - if (!accountId || (response.status !== 429 && response.status !== 403)) return; +): Promise { + if (!accountId || (response.status !== 429 && response.status !== 403)) return false; const payloadText = await readDisplaySafeErrorPayloadText(response.clone()); if (response.status === 429) { recordAntigravityCooldown( @@ -55,9 +55,13 @@ async function recordAntigravityHttpCooldown( isQuotaExhaustedBody(payloadText) ? "quota_exhausted" : "rate_limited", retryAfterMs(response.headers.get("retry-after")), ); - } else if (isAntigravityGeoBlockedBody(payloadText)) { + return true; + } + if (isAntigravityGeoBlockedBody(payloadText)) { recordAntigravityCooldown(accountId, "geo_blocked"); + return true; } + return false; } /** @@ -88,7 +92,10 @@ export async function fetchGoogleWithRetry( body: activeRequest.body, }, timeoutMs, ctx.abortSignal, ctx.stream, executor); if (label === "Antigravity") { - await recordAntigravityHttpCooldown(res, ctx.accountId); + const cooldownRecorded = await recordAntigravityHttpCooldown(res, ctx.accountId); + if (cooldownRecorded) { + return ctx.returnRawErrors ? res : normalizeFinalGoogleError(label, res, ctx.abortSignal); + } } if (res.status === 400 && repairInvalid400 && !compatibilityReplayUsed) { let payloadText = ""; diff --git a/src/oauth/antigravity-routing.ts b/src/oauth/antigravity-routing.ts index b3c0a9af9e..b5f5959caa 100644 --- a/src/oauth/antigravity-routing.ts +++ b/src/oauth/antigravity-routing.ts @@ -1,3 +1,8 @@ +/** + * Process-local Antigravity account health (cooldowns). Stored in an in-memory + * `Map` for the lifetime of this process only — + * cooldowns reset on restart and are not shared across workers. + */ export type AntigravityCooldownReason = "rate_limited" | "quota_exhausted" | "geo_blocked"; const DEFAULT_RATE_LIMITED_COOLDOWN_MS = 5_000; @@ -8,6 +13,7 @@ const GEO_BLOCKED_COOLDOWN_MS = 24 * 60 * 60_000; type AntigravityAccountHealth = { cooldownUntil: number; + reason: AntigravityCooldownReason; }; const accountHealth = new Map(); @@ -54,10 +60,23 @@ export function recordAntigravityCooldown( const cooldownUntil = now + cooldownDurationMs(reason, retryAfterMs); const current = accountHealth.get(accountId); if (!current || current.cooldownUntil < cooldownUntil) { - accountHealth.set(accountId, { cooldownUntil }); + accountHealth.set(accountId, { cooldownUntil, reason }); } } +export function getAntigravityAccountCooldown( + accountId: string, + now = Date.now(), +): { cooldownUntil: number; reason: AntigravityCooldownReason } | undefined { + const health = accountHealth.get(accountId); + if (!health) return undefined; + if (health.cooldownUntil <= now) { + accountHealth.delete(accountId); + return undefined; + } + return { cooldownUntil: health.cooldownUntil, reason: health.reason }; +} + export function isAntigravityAccountInCooldown(accountId: string, now = Date.now()): boolean { const health = accountHealth.get(accountId); if (!health) return false; diff --git a/tests/antigravity-routing.test.ts b/tests/antigravity-routing.test.ts index 19fe8506c3..8ea0cec498 100644 --- a/tests/antigravity-routing.test.ts +++ b/tests/antigravity-routing.test.ts @@ -1,19 +1,49 @@ import { afterEach, describe, expect, test } from "bun:test"; +import type { AdapterRequest } from "../src/adapters/base"; +import { fetchAntigravityWithRetry } from "../src/adapters/google-http"; import { clearAntigravityAccountCooldown, + getAntigravityAccountCooldown, isAntigravityAccountInCooldown, nextAntigravityAccount, recordAntigravityCooldown, sweepExpiredAntigravityRoutingHealth, } from "../src/oauth/antigravity-routing"; +const realFetch = globalThis.fetch; + const NOW = 1_700_000_000_000; const ACCOUNT_IDS = ["account-a", "account-b", "account-c"]; afterEach(() => { + globalThis.fetch = realFetch; for (const accountId of ACCOUNT_IDS) clearAntigravityAccountCooldown(accountId); + clearAntigravityAccountCooldown("account-http-retry"); }); +const antigravityRequest: AdapterRequest = { + url: "https://daily-cloudcode-pa.googleapis.com/v1/projects/p:generateContent", + method: "POST", + headers: { authorization: "Bearer tok", "content-type": "application/json" }, + body: "{}", +}; + +function mockFetch(responses: Array): { calls: RequestInit[] } { + const calls: RequestInit[] = []; + let i = 0; + globalThis.fetch = (async (_url: string | URL | Request, init?: RequestInit) => { + calls.push(init ?? {}); + const next = responses[i++] ?? responses[responses.length - 1]; + if (next instanceof Error) throw next; + return next; + }) as typeof fetch; + return { calls }; +} + +function googleError(code: number, status: string, message: string): string { + return JSON.stringify({ error: { code, status, message } }); +} + describe("Antigravity account cooldowns", () => { test("records a rate-limit cooldown and sweeps it after expiry", () => { recordAntigravityCooldown("account-a", "rate_limited", undefined, NOW); @@ -76,4 +106,104 @@ describe("Antigravity account cooldowns", () => { expect(isAntigravityAccountInCooldown("account-a", NOW + 59_999)).toBe(true); expect(isAntigravityAccountInCooldown("account-a", NOW + 60_001)).toBe(false); }); + + test("records and returns the cooldown reason via getAntigravityAccountCooldown", () => { + recordAntigravityCooldown("account-a", "rate_limited", undefined, NOW); + expect(getAntigravityAccountCooldown("account-a", NOW)).toEqual({ + cooldownUntil: NOW + 5_000, + reason: "rate_limited", + }); + + recordAntigravityCooldown("account-b", "quota_exhausted", 30_000, NOW); + expect(getAntigravityAccountCooldown("account-b", NOW)).toEqual({ + cooldownUntil: NOW + 30_000, + reason: "quota_exhausted", + }); + + recordAntigravityCooldown("account-c", "geo_blocked", undefined, NOW); + expect(getAntigravityAccountCooldown("account-c", NOW)).toEqual({ + cooldownUntil: NOW + 24 * 60 * 60_000, + reason: "geo_blocked", + }); + expect(getAntigravityAccountCooldown("account-c", NOW + 24 * 60 * 60_000)).toBeUndefined(); + }); + + test("keeps the longer cooldown and its reason when a shorter one is recorded", () => { + recordAntigravityCooldown("account-a", "geo_blocked", undefined, NOW); + recordAntigravityCooldown("account-a", "rate_limited", undefined, NOW + 1); + + expect(getAntigravityAccountCooldown("account-a", NOW)).toEqual({ + cooldownUntil: NOW + 24 * 60 * 60_000, + reason: "geo_blocked", + }); + }); + + test("updates reason when a longer cooldown replaces a shorter one", () => { + recordAntigravityCooldown("account-a", "rate_limited", undefined, NOW); + recordAntigravityCooldown("account-a", "quota_exhausted", 60_000, NOW); + + expect(getAntigravityAccountCooldown("account-a", NOW)).toEqual({ + cooldownUntil: NOW + 60_000, + reason: "quota_exhausted", + }); + }); +}); + +describe("Antigravity HTTP cooldown fail-fast", () => { + test("does not retry a rate-limited 429 after recording cooldown", async () => { + const mock = mockFetch([ + new Response(googleError(429, "RESOURCE_EXHAUSTED", "rate limit, try again"), { + status: 429, + headers: { "Retry-After": "0" }, + }), + new Response("ok", { status: 200 }), + ]); + + const res = await fetchAntigravityWithRetry(antigravityRequest, { + timeoutMs: 5_000, + accountId: "account-http-retry", + }); + + expect(res.status).toBe(429); + expect(mock.calls).toHaveLength(1); + expect(getAntigravityAccountCooldown("account-http-retry")?.reason).toBe("rate_limited"); + }); + + test("does not retry a quota-exhausted 429 after recording cooldown", async () => { + const mock = mockFetch([ + new Response( + googleError(429, "RESOURCE_EXHAUSTED", "Quota exceeded for your current billing plan"), + { status: 429, headers: { "Retry-After": "0" } }, + ), + new Response("ok", { status: 200 }), + ]); + + const res = await fetchAntigravityWithRetry(antigravityRequest, { + timeoutMs: 5_000, + accountId: "account-http-retry", + }); + + expect(res.status).toBe(429); + expect(mock.calls).toHaveLength(1); + expect(getAntigravityAccountCooldown("account-http-retry")?.reason).toBe("quota_exhausted"); + }); + + test("does not retry a geo-blocked 403 after recording cooldown", async () => { + const mock = mockFetch([ + new Response( + googleError(403, "PERMISSION_DENIED", "User location is not supported for the API use"), + { status: 403 }, + ), + new Response("ok", { status: 200 }), + ]); + + const res = await fetchAntigravityWithRetry(antigravityRequest, { + timeoutMs: 5_000, + accountId: "account-http-retry", + }); + + expect(res.status).toBe(403); + expect(mock.calls).toHaveLength(1); + expect(getAntigravityAccountCooldown("account-http-retry")?.reason).toBe("geo_blocked"); + }); });