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/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/adapters/google-http.ts b/src/adapters/google-http.ts index f7b90de87e..9650a172ce 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,38 @@ 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 false; + const payloadText = await readDisplaySafeErrorPayloadText(response.clone()); + if (response.status === 429) { + recordAntigravityCooldown( + accountId, + isQuotaExhaustedBody(payloadText) ? "quota_exhausted" : "rate_limited", + retryAfterMs(response.headers.get("retry-after")), + ); + return true; + } + if (isAntigravityGeoBlockedBody(payloadText)) { + recordAntigravityCooldown(accountId, "geo_blocked"); + return true; + } + return false; +} + /** * 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 +91,12 @@ export async function fetchGoogleWithRetry( headers: activeRequest.headers, body: activeRequest.body, }, timeoutMs, ctx.abortSignal, ctx.stream, executor); + if (label === "Antigravity") { + 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 = ""; try { 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..b5f5959caa --- /dev/null +++ b/src/oauth/antigravity-routing.ts @@ -0,0 +1,151 @@ +/** + * 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; +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; + reason: AntigravityCooldownReason; +}; + +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, 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; + 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/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 24b06ef7a9..c03e0b804a 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; @@ -1993,13 +1999,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; @@ -2009,39 +2015,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/src/server/responses/core.ts b/src/server/responses/core.ts index 3c7773e4e4..e129de0d07 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,52 @@ 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; + } 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 (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. // 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 +3969,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 +4074,7 @@ async function handleResponsesInner( providerName: route.providerName, modelId: route.modelId, }), + ...(antigravityAccountId ? { accountId: antigravityAccountId } : {}), }); } return await fetchWithHeaderTimeout(retryRequest.url, { @@ -4215,6 +4252,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 +4455,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..8c4e0df216 --- /dev/null +++ b/tests/antigravity-project-bind.test.ts @@ -0,0 +1,77 @@ +import { describe, expect, test } from "bun:test"; +import { readFileSync } from "node:fs"; +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"); + }); + + 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"'), + ); + }); +}); 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/antigravity-routing.test.ts b/tests/antigravity-routing.test.ts new file mode 100644 index 0000000000..8ea0cec498 --- /dev/null +++ b/tests/antigravity-routing.test.ts @@ -0,0 +1,209 @@ +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); + + 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); + }); + + 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"); + }); +}); 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/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(); 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", 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",