diff --git a/src/adapters/google-antigravity-hosts.ts b/src/adapters/google-antigravity-hosts.ts new file mode 100644 index 0000000000..a2c1fe096d --- /dev/null +++ b/src/adapters/google-antigravity-hosts.ts @@ -0,0 +1,26 @@ +const DAILY_ANTIGRAVITY_HOST = "https://daily-cloudcode-pa.googleapis.com"; +const PROD_ANTIGRAVITY_HOST = "https://cloudcode-pa.googleapis.com"; + +/** + * 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(/\/+$/, ""); + 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. */ +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..c4f0cb58f2 --- /dev/null +++ b/src/providers/antigravity-quota.ts @@ -0,0 +1,216 @@ +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); + 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(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.record).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 { + 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 { + const candidates = collectCandidates(payload); + for (const candidate of candidates) { + if (!isWeeklyPath(candidate.path)) continue; + 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 { + const payload = await readProviderQuotaJsonForTests(response, timeoutMs); + if (payload === null) throw new Error("Antigravity quota RPC returned unreadable JSON"); + return payload; +} + +export class AntigravityQuotaRpcError extends Error { + constructor(readonly status: number) { + super(`Antigravity quota RPC failed: ${status}`); + } +} + +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( + 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), + ]); + 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; + 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 (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 db0202161d..c459c8f2a3 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -28,6 +28,12 @@ import { type CodexCapacityAggregation, type CodexCapacityQuota, } from "./codex-capacity"; +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. */ const ACCOUNT_TOKEN_SKEW_MS = 60_000; @@ -2065,13 +2071,13 @@ 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); } -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; @@ -2081,39 +2087,81 @@ 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), - }); - if (!response.ok) return null; - const body = asRecord(await readQuotaJson(response)); - const models = asRecord(body?.models); - if (!models) return null; + 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 TERMINAL_QUOTA_FAILURE; + } + liveQuota = 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 }), + redirect: "error", + 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..0a910d959d --- /dev/null +++ b/tests/antigravity-quota.test.ts @@ -0,0 +1,523 @@ +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 }); +}); + +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({ + 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 () => { + 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: [ + { 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")); + expect(requestOptions.every(({ init }) => init?.redirect === "error")).toBe(true); + }); + + 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 production or catalog 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 jsonResponse({}, 404); + 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("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) => { + 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 jsonResponse({}, 404); + 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 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 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; + + 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 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"); + }); + + 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).toBeNull(); + }); + + 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); + }); +}); diff --git a/tests/provider-quota.test.ts b/tests/provider-quota.test.ts index dffa089c56..5b0ea53d82 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",