diff --git a/src/lib/state-store-registrations.ts b/src/lib/state-store-registrations.ts index 61aad5f7b9..b10ca4cd50 100644 --- a/src/lib/state-store-registrations.ts +++ b/src/lib/state-store-registrations.ts @@ -37,7 +37,7 @@ import { sweepExpiredApiKeyCooldowns } from "../providers/key-failover"; import { reconcileProviderRequestPacing } from "../providers/request-pacing"; import { sweepAbandonedResponseStateTemps, sweepExpiredResponseStates } from "../responses/state"; import { sweepExpiredAntigravityReplay } from "../adapters/google-antigravity-replay"; -import { reconcileProviderAccountQuotaRows } from "../providers/quota"; +import { listLiveProviderAccountQuotaKeys, reconcileProviderAccountQuotaRows } from "../providers/quota"; import { reconcileRouterWarningMemos } from "../router"; import type { OcxConfig } from "../types"; import { @@ -62,13 +62,15 @@ export function reconcileLiveStateStores() { export function buildGenerationContext(): GenerationContext { if (!liveServerConfig) throw new Error("live server config is not installed"); const providerNames = new Set(Object.keys(liveServerConfig.providers)); + const oauthAccountKeys = listLiveOAuthAccountKeys(providerNames); return { generation: 0, providerNames, comboIds: new Set(Object.keys(liveServerConfig.combos ?? {})), comboTargets: listLiveComboTargetKeys(liveServerConfig), codexAccountIds: listLiveCodexAccountIds(liveServerConfig), - oauthAccountKeys: listLiveOAuthAccountKeys(providerNames), + oauthAccountKeys, + providerAccountQuotaKeys: listLiveProviderAccountQuotaKeys(liveServerConfig.providers, oauthAccountKeys), configRoots: listLiveConfigOwnershipRoots(getConfigDir()), }; } diff --git a/src/lib/state-store-sweeper.ts b/src/lib/state-store-sweeper.ts index 3f687d84fc..7ee0e6268f 100644 --- a/src/lib/state-store-sweeper.ts +++ b/src/lib/state-store-sweeper.ts @@ -7,6 +7,7 @@ export interface GenerationContext { comboTargets: ReadonlySet; codexAccountIds: ReadonlySet; oauthAccountKeys: ReadonlySet; + providerAccountQuotaKeys: ReadonlySet; configRoots: ReadonlySet; } diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 63c74b9cd6..9e9e118b53 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -1,3 +1,7 @@ +const originalFetch = globalThis.fetch; +import { providerDestinationConfigError } from "../lib/destination-policy"; +import { providerOutboundPost, providerRedirectError } from "../lib/provider-outbound"; + import { createHash } from "node:crypto"; import { effectiveCodexAuthAccountId, @@ -1432,11 +1436,13 @@ type AccountQuotaCacheEntry = { const accountQuotaCache = new Map(); const accountQuotaInflight = new Map>(); let lastReconciledGeneration = 0; -let liveAccountQuotaKeys = new Set(); +let liveAccountQuotaKeys = new Map(); let liveProviderQuotaKeys = new Set(); function mayCommitAccountQuotaKey(key: string, writerGeneration: number): boolean { - return writerGeneration >= lastReconciledGeneration || liveAccountQuotaKeys.has(key); + const liveSinceGeneration = liveAccountQuotaKeys.get(key); + return writerGeneration >= lastReconciledGeneration + || (liveSinceGeneration !== undefined && writerGeneration >= liveSinceGeneration); } function mayCommitProviderQuotaKey(key: string, writerGeneration: number): boolean { @@ -1452,19 +1458,51 @@ export interface ProviderAccountQuota { /** Providers whose per-account quota can be probed. Extend as other OAuth APIs are covered. */ export function supportsPerAccountQuota(provider: string): boolean { - return provider === "anthropic"; + return provider === "anthropic" || provider === "google-antigravity"; +} + +function normalizeAntigravityDestination(baseUrl?: string): string { + const raw = baseUrl?.trim() || "https://daily-cloudcode-pa.googleapis.com"; + try { + const u = new URL(raw); + return u.origin.toLowerCase() + u.pathname.replace(/\/+$/, ""); + } catch { + return raw.replace(/\/+$/, "").toLowerCase(); + } +} + +function accountCacheKey(provider: string, accountId: string, destination = ""): string { + return destination ? provider + "\u0000" + accountId + "\u0000" + destination : provider + "\u0000" + accountId; } -function accountCacheKey(provider: string, accountId: string): string { - return `${provider}\u0000${accountId}`; +export function listLiveProviderAccountQuotaKeys( + providers: OcxConfig["providers"], + oauthAccountKeys: ReadonlySet, +): Set { + const keys = new Set(); + for (const canonical of oauthAccountKeys) { + const separator = canonical.indexOf("\0"); + const provider = canonical.slice(0, separator); + const accountId = canonical.slice(separator + 1); + const destination = provider === "google-antigravity" + ? normalizeAntigravityDestination(providers[provider]?.baseUrl) + : ""; + keys.add(accountCacheKey(provider, accountId, destination)); + } + return keys; } /** * Synchronous last-good per-account quota read for routing. Never probes the network. * Returns null when nothing is cached (or the cached row has no bars). */ -export function getCachedProviderAccountQuota(provider: string, accountId: string): ProviderQuota | null { - const entry = accountQuotaCache.get(accountCacheKey(provider, accountId)); +export function getCachedProviderAccountQuota( + provider: string, + accountId: string, + baseUrl?: string, +): ProviderQuota | null { + const destination = provider === "google-antigravity" ? normalizeAntigravityDestination(baseUrl) : ""; + const entry = accountQuotaCache.get(accountCacheKey(provider, accountId, destination)); return entry?.quota ?? null; } @@ -1494,9 +1532,16 @@ export function sweepExpiredProviderAccountQuotaRows(now = Date.now()): number { export function reconcileProviderAccountQuotaRows(context: GenerationContext): number { if (context.generation <= lastReconciledGeneration) return 0; + const nextLiveAccountQuotaKeys = new Map(); + for (const key of context.providerAccountQuotaKeys) { + nextLiveAccountQuotaKeys.set( + key, + liveAccountQuotaKeys.get(key) ?? (lastReconciledGeneration === 0 ? 0 : context.generation), + ); + } let removed = 0; for (const key of accountQuotaCache.keys()) { - if (context.oauthAccountKeys.has(key)) continue; + if (context.providerAccountQuotaKeys.has(key)) continue; accountQuotaCache.delete(key); removed += 1; } @@ -1505,7 +1550,7 @@ export function reconcileProviderAccountQuotaRows(context: GenerationContext): n removed += cache.response.reports.length - reports.length; cache = { ...cache, response: { ...cache.response, reports } }; } - liveAccountQuotaKeys = new Set(context.oauthAccountKeys); + liveAccountQuotaKeys = nextLiveAccountQuotaKeys; liveProviderQuotaKeys = new Set(context.providerNames); lastReconciledGeneration = context.generation; return removed; @@ -1514,7 +1559,7 @@ export function reconcileProviderAccountQuotaRows(context: GenerationContext): n /** Test-only reset so a direct reconcile call in one file cannot leak across files. */ export function resetProviderQuotaReconcileStateForTests(): void { lastReconciledGeneration = 0; - liveAccountQuotaKeys = new Set(); + liveAccountQuotaKeys = new Map(); liveProviderQuotaKeys = new Set(); } @@ -1560,9 +1605,12 @@ async function getTokenForAccountQuotaProbe(provider: string, accountId: string) async function fetchAccountQuota( provider: string, accountId: string, - forceRefresh: boolean, + forceRefresh = false, + baseUrl?: string, + allowPrivateNetwork?: boolean, ): Promise { - const key = accountCacheKey(provider, accountId); + const normalizedDest = provider === "google-antigravity" ? normalizeAntigravityDestination(baseUrl) : ""; + const key = accountCacheKey(provider, accountId, normalizedDest); const writerGeneration = captureConfigGeneration(); const cached = accountQuotaCache.get(key); if (!forceRefresh && cached && Date.now() - cached.ts < ACCOUNT_QUOTA_TTL_MS) return cached; @@ -1571,8 +1619,39 @@ async function fetchAccountQuota( const probe = (async (): Promise => { try { - const token = await getTokenForAccountQuotaProbe(provider, accountId); - const quota = await fetchAnthropicUsageQuota(token); + let quota: ProviderQuota | null = null; + if (provider === "anthropic") { + const token = await getTokenForAccountQuotaProbe(provider, accountId); + quota = await fetchAnthropicUsageQuota(token); + } else if (provider === "google-antigravity") { + const stored = getAccountCredential(provider, accountId); + if (!stored?.projectId) { + // Permanent configuration gap (projectId only appears after a fresh login): + // skip silently without marking the row unavailable, so the GUI does not + // surface a spurious quota error on every poll. + const entry: AccountQuotaCacheEntry = { ts: Date.now(), quota: null }; + if (mayCommitAccountQuotaKey(key, writerGeneration)) { + accountQuotaCache.set(key, entry); + sweepExpiredOnWrite(entry.ts); + } + return entry; + } + // Pre-flight destination policy gate before acquiring or refreshing account token + const destError = providerDestinationConfigError("google-antigravity", { + baseUrl: baseUrl || "https://daily-cloudcode-pa.googleapis.com", + allowPrivateNetwork: allowPrivateNetwork ?? false, + }); + if (destError) { + const entry: AccountQuotaCacheEntry = { ts: Date.now(), quota: null, unavailable: true }; + if (mayCommitAccountQuotaKey(key, writerGeneration)) { + accountQuotaCache.set(key, entry); + sweepExpiredOnWrite(entry.ts); + } + return entry; + } + const token = await getTokenForAccountQuotaProbe(provider, accountId); + quota = await fetchAntigravityUsageQuota(token, stored.projectId, baseUrl, { allowPrivateNetwork }); + } if (!quota) { // Preserve last-good bars and mark unavailable; advance TTL so failures // negative-cache instead of re-probing on every GUI poll. @@ -1619,12 +1698,14 @@ async function fetchAccountQuota( export async function fetchProviderAccountQuotas( provider: string, forceRefresh = false, + baseUrl?: string, + allowPrivateNetwork?: boolean, ): Promise { if (!supportsPerAccountQuota(provider)) return []; const set = getAccountSet(provider); if (!set) return []; return await Promise.all(set.accounts.map(async account => { - const entry = await fetchAccountQuota(provider, account.id, forceRefresh); + const entry = await fetchAccountQuota(provider, account.id, forceRefresh, baseUrl, allowPrivateNetwork); return { accountId: account.id, quota: entry.quota, @@ -2126,31 +2207,49 @@ function antigravityUsedPercent(quotaInfo: Record): number | un ? toFiniteNumber(quotaInfo.remainingPercentage)! * 100 : undefined); if (remaining === undefined) return undefined; - return normalizePercent(100 - remaining); -} - -async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig): Promise { - const credential = getCredential("google-antigravity"); - if (!credential?.projectId) return null; - let accessToken: string; + return Math.min(100, Math.max(0, normalizePercent(100 - remaining) ?? 0)); +} + +export async function fetchAntigravityUsageQuota( + accessToken: string, + projectId: string, + baseUrl = "https://daily-cloudcode-pa.googleapis.com", + options?: { + allowPrivateNetwork?: boolean; + fetch?: typeof globalThis.fetch; + outboundPost?: typeof providerOutboundPost; + }, +): Promise { + const outboundPost = options?.outboundPost ?? providerOutboundPost; + const normalizedUrl = (baseUrl || "https://daily-cloudcode-pa.googleapis.com").replace(/\/+$/, ""); + const targetUrl = normalizedUrl + "/v1internal:fetchAvailableModels"; + let response: Response; try { - accessToken = await getValidAccessToken("google-antigravity"); + const activeFetch = options?.fetch ?? (globalThis.fetch !== originalFetch ? globalThis.fetch : undefined); + response = await outboundPost( + "google-antigravity", + { + baseUrl: normalizedUrl, + allowPrivateNetwork: options?.allowPrivateNetwork ?? false, + ...(activeFetch ? { fetch: activeFetch } : {}), + }, + targetUrl, + { + headers: { + Accept: "application/json", + "Content-Type": "application/json", + "User-Agent": antigravityUserAgent(), + Authorization: "Bearer " + accessToken, + }, + body: JSON.stringify({ project: projectId }), + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }, + ); } catch { 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 redirectError = await providerRedirectError(response, targetUrl); + if (redirectError || !response.ok) return null; const body = asRecord(await readQuotaJson(response)); const models = asRecord(body?.models); if (!models) return null; @@ -2177,8 +2276,34 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig return window ? [window] : []; }); if (customWindows.length === 0) return null; + return { customWindows, updatedAt: Date.now() }; +} + +async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig): Promise { + const destError = providerDestinationConfigError("google-antigravity", config); + if (destError) return null; + const probedAccountId = getAccountSet("google-antigravity")?.activeAccountId; + const normalizedDest = normalizeAntigravityDestination(config.baseUrl); + const probedAccountKey = probedAccountId ? accountCacheKey("google-antigravity", probedAccountId, normalizedDest) : null; + const writerGeneration = captureConfigGeneration(); + const credential = getCredential("google-antigravity"); + if (!credential?.projectId) return null; + let accessToken: string; + try { + accessToken = await getValidAccessToken("google-antigravity"); + } catch { + return null; + } + const quota = await fetchAntigravityUsageQuota(accessToken, credential.projectId, config.baseUrl, { allowPrivateNetwork: config.allowPrivateNetwork }); + if (!quota) return null; + if (probedAccountId && probedAccountKey) { + const stillOwnsToken = getAccountCredential("google-antigravity", probedAccountId)?.access === accessToken; + if (stillOwnsToken && mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) { + accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota }); + } + } return report(provider, "google-antigravity:fetchAvailableModels", { - customWindows, + ...quota, updatedAt: Date.now(), }); } diff --git a/src/server/management/oauth-account-routes.ts b/src/server/management/oauth-account-routes.ts index ee57515b41..2f01ca0060 100644 --- a/src/server/management/oauth-account-routes.ts +++ b/src/server/management/oauth-account-routes.ts @@ -279,7 +279,7 @@ export async function handleOauthAccountRoutes(ctx: ManagementContext): Promise< const forceRefresh = url.searchParams.get("refresh") === "1"; // Probing may refresh the active credential and mark needsReauth — project health // from the post-probe store so the response is not stale. - const rows = await fetchProviderAccountQuotas(provider, forceRefresh); + const rows = await fetchProviderAccountQuotas(provider, forceRefresh, ctx.config.providers[provider]?.baseUrl, ctx.config.providers[provider]?.allowPrivateNetwork); const byId = new Map(rows.map(row => [row.accountId, row])); const projected = projectAccounts(); return jsonResponse({ diff --git a/tests/combos.test.ts b/tests/combos.test.ts index 166a738554..27bee2776a 100644 --- a/tests/combos.test.ts +++ b/tests/combos.test.ts @@ -809,6 +809,7 @@ describe("combo generation reconciliation", () => { comboTargets: new Set(["free::a/m1", "free::c/m3"]), codexAccountIds: new Set(), oauthAccountKeys: new Set(), + providerAccountQuotaKeys: new Set(), configRoots: new Set(), }); expect(removed).toBeGreaterThan(0); @@ -841,6 +842,7 @@ describe("combo generation reconciliation", () => { comboTargets: new Set(["free::a/m1", "free::c/m3"]), codexAccountIds: new Set(), oauthAccountKeys: new Set(), + providerAccountQuotaKeys: new Set(), configRoots: new Set(), }); diff --git a/tests/oauth-store-multi.test.ts b/tests/oauth-store-multi.test.ts index 50893daedc..52fe6630c1 100644 --- a/tests/oauth-store-multi.test.ts +++ b/tests/oauth-store-multi.test.ts @@ -327,6 +327,7 @@ describe("multi-account auth store", () => { comboTargets: new Set(), codexAccountIds: new Set(), oauthAccountKeys: new Set(), + providerAccountQuotaKeys: new Set(), configRoots: new Set(), }); release(); diff --git a/tests/provider-account-quota.test.ts b/tests/provider-account-quota.test.ts index 634c663ca2..b0f894417f 100644 --- a/tests/provider-account-quota.test.ts +++ b/tests/provider-account-quota.test.ts @@ -3,6 +3,8 @@ import { mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { saveCredential } from "../src/oauth/store"; +import { PROXY_ENV_KEYS } from "../src/lib/proxy-env"; +import { reconcileStateGeneration, registerStateStore } from "../src/lib/state-store-sweeper"; import type { OcxConfig } from "../src/types"; import { clearAccountQuotaCache, @@ -10,6 +12,7 @@ import { fetchProviderAccountQuotas, fetchProviderQuotaReports, getCachedProviderAccountQuota, + listLiveProviderAccountQuotaKeys, reconcileProviderAccountQuotaRows, resetProviderQuotaReconcileStateForTests, supportsPerAccountQuota, @@ -17,6 +20,8 @@ import { const originalFetch = globalThis.fetch; const previousOpencodexHome = process.env.OPENCODEX_HOME; +const proxyKeys = PROXY_ENV_KEYS.flatMap(key => [key, key.toLowerCase()]); +const previousProxyEnv = Object.fromEntries(proxyKeys.map(key => [key, process.env[key]])); let opencodexHome: string; const FIRST = { accountId: "acct-first", email: "first@example.com" }; @@ -39,6 +44,7 @@ function usageBody(fiveHour: number, sevenDay: number): string { beforeEach(() => { opencodexHome = mkdtempSync(join(tmpdir(), "ocx-account-quota-")); process.env.OPENCODEX_HOME = opencodexHome; + for (const key of proxyKeys) delete process.env[key]; clearAccountQuotaCache(); clearProviderQuotaCache(); }); @@ -47,6 +53,11 @@ afterEach(() => { globalThis.fetch = originalFetch; if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = previousOpencodexHome; + for (const key of proxyKeys) { + const previous = previousProxyEnv[key]; + if (previous === undefined) delete process.env[key]; + else process.env[key] = previous; + } rmSync(opencodexHome, { recursive: true, force: true }); clearAccountQuotaCache(); clearProviderQuotaCache(); @@ -411,6 +422,7 @@ describe("fetchProviderAccountQuotas", () => { comboTargets: new Set(), codexAccountIds: new Set(), oauthAccountKeys: new Set(), + providerAccountQuotaKeys: new Set(), configRoots: new Set(), }); releaseUsage(); @@ -418,4 +430,574 @@ describe("fetchProviderAccountQuotas", () => { expect(getCachedProviderAccountQuota("anthropic", first!.id)).toBeNull(); }); + + test("reports Google Antigravity per-account Gem/Cla quota keyed by account credentials (#1082)", async () => { + const { saveCredential } = await import("../src/oauth/store"); + await saveCredential("google-antigravity", { + access: "token-agy-1", + refresh: "refresh-agy-1", + expires: Date.now() + 3600_000, + projectId: "project-1", + accountId: "acct-agy-1", + email: "agy1@example.com", + }); + await saveCredential("google-antigravity", { + access: "token-agy-2", + refresh: "refresh-agy-2", + expires: Date.now() + 3600_000, + projectId: "project-2", + accountId: "acct-agy-2", + email: "agy2@example.com", + }); + + const seenAuthorizationByProject = new Map(); + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + expect(String(input)).toContain("/v1internal:fetchAvailableModels"); + const body = JSON.parse(String(init?.body)) as { project?: string }; + seenAuthorizationByProject.set( + body.project ?? "", + new Headers(init?.headers).get("Authorization"), + ); + if (body.project === "project-1") { + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { + quotaInfo: { remainingFraction: 0.64, resetTime: "2026-07-05T14:00:00Z" }, + }, + "claude-sonnet-4-6": { + quotaInfo: { remainingFraction: 0.21, resetTime: "2026-07-05T15:00:00Z" }, + }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + } + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { + quotaInfo: { remainingFraction: 0.90, resetTime: "2026-07-05T14:00:00Z" }, + }, + "claude-sonnet-4-6": { + quotaInfo: { remainingFraction: 0.85, resetTime: "2026-07-05T15:00:00Z" }, + }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + const rows = await fetchProviderAccountQuotas("google-antigravity"); + expect(rows.length).toBe(2); + expect(seenAuthorizationByProject.get("project-1")).toBe("Bearer token-agy-1"); + expect(seenAuthorizationByProject.get("project-2")).toBe("Bearer token-agy-2"); + + const window = (accountId: string, label: string) => rows + .find(r => r.accountId === accountId) + ?.quota?.customWindows?.find(w => w.label === label)?.percent; + const { getAccountSet } = await import("../src/oauth/store"); + const ids = getAccountSet("google-antigravity")!.accounts.map(a => a.id); + // 1 - 0.64 = 36% used, 1 - 0.21 = 79% used for account 1 + // 1 - 0.90 = 10% used, 1 - 0.85 = 15% used for account 2 + expect([window(ids[0]!, "Gem"), window(ids[0]!, "Cla")]).toEqual([36, 79]); + expect([window(ids[1]!, "Gem"), window(ids[1]!, "Cla")]).toEqual([10, 15]); + }); + + test("uses the configured provider baseUrl for Antigravity account quota probes", async () => { + const { saveCredential } = await import("../src/oauth/store"); + await saveCredential("google-antigravity", { + access: "token-agy-1", + refresh: "refresh-agy-1", + expires: Date.now() + 3600_000, + projectId: "project-1", + accountId: "acct-agy-1", + email: "agy1@example.com", + }); + + const seenUrls: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + seenUrls.push(String(input)); + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { quotaInfo: { remainingFraction: 0.5, resetTime: "2026-07-05T14:00:00Z" } }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + const rows = await fetchProviderAccountQuotas("google-antigravity", true, "https://custom-antigravity.example.com"); + expect(rows.length).toBe(1); + expect(seenUrls).toEqual(["https://custom-antigravity.example.com/v1internal:fetchAvailableModels"]); + }); + + test("skips Antigravity accounts without projectId instead of throwing or flagging unavailable", async () => { + const { saveCredential } = await import("../src/oauth/store"); + await saveCredential("google-antigravity", { + access: "token-with-project", + refresh: "refresh-with-project", + expires: Date.now() + 3600_000, + projectId: "project-1", + accountId: "acct-agy-1", + email: "agy1@example.com", + }); + await saveCredential("google-antigravity", { + access: "token-no-project", + refresh: "refresh-no-project", + expires: Date.now() + 3600_000, + accountId: "acct-agy-2", + email: "agy2@example.com", + }); + + let fetchCalls = 0; + globalThis.fetch = (async () => { + fetchCalls += 1; + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { quotaInfo: { remainingFraction: 0.5, resetTime: "2026-07-05T14:00:00Z" } }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + const rows = await fetchProviderAccountQuotas("google-antigravity", true); + expect(rows.length).toBe(2); + expect(fetchCalls).toBe(1); + const nullQuotaRows = rows.filter(r => r.quota === null); + expect(nullQuotaRows.length).toBe(1); + expect(nullQuotaRows[0]!.unavailable).toBeUndefined(); + const quotaRows = rows.filter(r => r.quota !== null); + expect(quotaRows.length).toBe(1); + expect(quotaRows[0]!.quota?.customWindows?.length).toBeGreaterThan(0); + }); + + test("negative-caches Antigravity upstream 404 so probes are not repeated within the TTL", async () => { + const { saveCredential } = await import("../src/oauth/store"); + await saveCredential("google-antigravity", { + access: "token-agy-1", + refresh: "refresh-agy-1", + expires: Date.now() + 3600_000, + projectId: "project-1", + accountId: "acct-agy-1", + email: "agy1@example.com", + }); + + let fetchCalls = 0; + globalThis.fetch = (async () => { + fetchCalls += 1; + return new Response("not found", { status: 404, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + const first = await fetchProviderAccountQuotas("google-antigravity", true); + expect(first[0]!.unavailable).toBe(true); + expect(fetchCalls).toBe(1); + + const second = await fetchProviderAccountQuotas("google-antigravity", false); + expect(fetchCalls).toBe(1); + }); + + test("clamps invalid or out-of-range remainingFraction values safely", async () => { + const { fetchAntigravityUsageQuota } = await import("../src/providers/quota"); + globalThis.fetch = (async () => { + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { quotaInfo: { remainingFraction: 1.5, resetTime: "2026-07-05T14:00:00Z" } }, + "claude-sonnet-4-6": { quotaInfo: { remainingFraction: -0.2, resetTime: "2026-07-05T15:00:00Z" } }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + const quota = await fetchAntigravityUsageQuota("token", "project-1"); + expect(quota).not.toBeNull(); + const gem = quota?.customWindows?.find(w => w.label === "Gem"); + const cla = quota?.customWindows?.find(w => w.label === "Cla"); + expect(gem?.percent).toBe(0); + expect(cla?.percent).toBe(100); + }); + + test("gracefully handles non-JSON upstream 502 error responses without uncaught exceptions", async () => { + const { fetchAntigravityUsageQuota } = await import("../src/providers/quota"); + globalThis.fetch = (async () => { + return new Response("502 Bad Gateway", { + status: 502, + headers: { "content-type": "text/html" }, + }); + }) as typeof fetch; + + const quota = await fetchAntigravityUsageQuota("token", "project-1"); + expect(quota).toBeNull(); + }); + test("destination-bound cache isolates quota across baseUrl changes and prevents cross-endpoint contamination", async () => { + const { saveCredential } = await import("../src/oauth/store"); + await saveCredential("google-antigravity", { + access: "token-agy-1", + refresh: "refresh-agy-1", + expires: Date.now() + 3600_000, + projectId: "project-1", + accountId: "acct-agy-1", + email: "agy1@example.com", + }); + + let requestedUrl = ""; + globalThis.fetch = (async (input: RequestInfo | URL) => { + requestedUrl = String(input); + if (requestedUrl.includes("dest-1.example.com")) { + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { quotaInfo: { remainingFraction: 0.9, resetTime: "2026-07-05T14:00:00Z" } }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + } + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { quotaInfo: { remainingFraction: 0.2, resetTime: "2026-07-05T14:00:00Z" } }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + // 1. Probe dest-1 (10% used) + const rows1 = await fetchProviderAccountQuotas("google-antigravity", true, "https://dest-1.example.com"); + expect(rows1[0]!.quota?.customWindows?.find(w => w.label === "Gem")?.percent).toBe(10); + + // 2. Non-forced probe to dest-2 must NOT return dest-1 cached value (80% used) + const rows2 = await fetchProviderAccountQuotas("google-antigravity", false, "https://dest-2.example.com"); + expect(rows2[0]!.quota?.customWindows?.find(w => w.label === "Gem")?.percent).toBe(80); + }); + + test("rejects invalid destination policy targets before sending bearer tokens", async () => { + const { saveCredential } = await import("../src/oauth/store"); + await saveCredential("google-antigravity", { + access: "token-agy-1", + refresh: "refresh-agy-1", + expires: Date.now() + 3600_000, + projectId: "project-1", + accountId: "acct-agy-1", + email: "agy1@example.com", + }); + + let fetchCalled = false; + globalThis.fetch = (async () => { + fetchCalled = true; + return new Response("ok"); + }) as typeof fetch; + + // Probe blocked metadata endpoint 169.254.169.254 + const rows = await fetchProviderAccountQuotas("google-antigravity", true, "http://169.254.169.254"); + expect(rows[0]!.unavailable).toBe(true); + expect(fetchCalled).toBe(false); + }); + + test("destination-bound in-flight promise and stale writer rejection on baseUrl change", async () => { + const { saveCredential } = await import("../src/oauth/store"); + await saveCredential("google-antigravity", { + access: "token-agy-1", + refresh: "refresh-agy-1", + expires: Date.now() + 3600_000, + projectId: "project-1", + accountId: "acct-agy-1", + email: "agy1@example.com", + }); + + let resolveDest1: (value: Response) => void; + const dest1Promise = new Promise(resolve => { + resolveDest1 = resolve; + }); + + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes("dest-1.example.com")) { + return dest1Promise; + } + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { quotaInfo: { remainingFraction: 0.5, resetTime: "2026-07-05T14:00:00Z" } }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + // Start probe to dest-1 (in flight) + const probe1 = fetchProviderAccountQuotas("google-antigravity", true, "https://dest-1.example.com"); + + // Immediately start probe to dest-2 (must not join probe 1) + const probe2 = fetchProviderAccountQuotas("google-antigravity", true, "https://dest-2.example.com"); + const rows2 = await probe2; + expect(rows2[0]!.quota?.customWindows?.find(w => w.label === "Gem")?.percent).toBe(50); + + // Resolve dest-1 + resolveDest1!(new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { quotaInfo: { remainingFraction: 0.9, resetTime: "2026-07-05T14:00:00Z" } }, + }, + }), { status: 200, headers: { "content-type": "application/json" } })); + + const rows1 = await probe1; + expect(rows1[0]!.quota?.customWindows?.find(w => w.label === "Gem")?.percent).toBe(10); + }); + + test("a pre-reconcile destination writer cannot repopulate cache after an A to B to A switch", async () => { + const { getAccountSet, saveCredential } = await import("../src/oauth/store"); + await saveCredential("google-antigravity", { + access: "token-agy-aba", + refresh: "refresh-agy-aba", + expires: Date.now() + 3600_000, + projectId: "project-aba", + accountId: "acct-agy-aba", + email: "aba@example.com", + }); + + let releaseFirstA!: (value: Response) => void; + const firstAResponse = new Promise(resolve => { releaseFirstA = resolve; }); + let markFirstAStarted!: () => void; + const firstAStarted = new Promise(resolve => { markFirstAStarted = resolve; }); + let aCalls = 0; + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes("dest-a.example.com")) { + aCalls += 1; + if (aCalls === 1) { + markFirstAStarted(); + return firstAResponse; + } + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { quotaInfo: { remainingFraction: 0.6, resetTime: "2026-07-05T14:00:00Z" } }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + } + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { quotaInfo: { remainingFraction: 0.5, resetTime: "2026-07-05T14:00:00Z" } }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + const accountId = getAccountSet("google-antigravity")!.accounts[0]!.id; + const context = { + generation: 0, + providerNames: new Set(["google-antigravity"]), + comboIds: new Set(), + comboTargets: new Set(), + codexAccountIds: new Set(), + oauthAccountKeys: new Set([`google-antigravity\0${accountId}`]), + configRoots: new Set(), + }; + const destinationKeys = (baseUrl: string) => listLiveProviderAccountQuotaKeys({ + "google-antigravity": { adapter: "google-antigravity", authMode: "oauth", baseUrl }, + }, context.oauthAccountKeys); + const unregister = registerStateStore({ + name: "provider-account-quota-aba-test", + reconcileGeneration: reconcileProviderAccountQuotaRows, + }); + try { + reconcileStateGeneration({ + ...context, + providerAccountQuotaKeys: destinationKeys("https://dest-a.example.com"), + }); + const staleA = fetchProviderAccountQuotas("google-antigravity", true, "https://dest-a.example.com"); + await firstAStarted; + + reconcileStateGeneration({ + ...context, + providerAccountQuotaKeys: destinationKeys("https://dest-b.example.com"), + }); + const rowsB = await fetchProviderAccountQuotas("google-antigravity", true, "https://dest-b.example.com"); + expect(rowsB[0]!.quota?.customWindows?.find(w => w.label === "Gem")?.percent).toBe(50); + expect(getCachedProviderAccountQuota( + "google-antigravity", + accountId, + "https://dest-b.example.com", + )?.customWindows?.find(w => w.label === "Gem")?.percent).toBe(50); + expect(getCachedProviderAccountQuota( + "google-antigravity", + accountId, + "https://dest-a.example.com", + )).toBeNull(); + reconcileStateGeneration({ + ...context, + providerAccountQuotaKeys: destinationKeys("https://dest-a.example.com"), + }); + + releaseFirstA(new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { quotaInfo: { remainingFraction: 0.9, resetTime: "2026-07-05T14:00:00Z" } }, + }, + }), { status: 200, headers: { "content-type": "application/json" } })); + await staleA; + expect(getCachedProviderAccountQuota( + "google-antigravity", + accountId, + "https://dest-a.example.com", + )).toBeNull(); + + const currentA = await fetchProviderAccountQuotas("google-antigravity", false, "https://dest-a.example.com"); + expect(aCalls).toBe(2); + expect(currentA[0]!.quota?.customWindows?.find(w => w.label === "Gem")?.percent).toBe(40); + expect(getCachedProviderAccountQuota( + "google-antigravity", + accountId, + "https://dest-a.example.com", + )?.customWindows?.find(w => w.label === "Gem")?.percent).toBe(40); + } finally { + unregister(); + } + }); + + test("fail-closed redirect handling: rejected destination or redirect yields null without following", async () => { + const { fetchAntigravityUsageQuota } = await import("../src/providers/quota"); + + let redirectTargetCalled = false; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const url = String(input); + if (url.includes("redirect-sink.example.com")) { + redirectTargetCalled = true; + return new Response("sink", { status: 200 }); + } + // If fetch honors redirect: error, fetch will throw or fail on 302 + if (init?.redirect === "error") { + throw new TypeError("Failed to fetch: redirect"); + } + return new Response("", { status: 302, headers: { Location: "https://redirect-sink.example.com" } }); + }) as typeof fetch; + + const quota = await fetchAntigravityUsageQuota("token", "project-1", "https://redirecting.example.com"); + expect(quota).toBeNull(); + expect(redirectTargetCalled).toBe(false); + }); + + test("Antigravity quota blocks private DNS answers and pins the bearer POST to the validated public address", async () => { + const { fetchAntigravityUsageQuota } = await import("../src/providers/quota"); + const { providerOutboundPost } = await import("../src/lib/provider-outbound"); + + let privateResolveCalls = 0; + let privatePinnedCalls = 0; + const blocked = await fetchAntigravityUsageQuota( + "private-dns-token", + "private-dns-project", + "https://quota-private.example.com", + { + outboundPost: (name, provider, url, init) => providerOutboundPost(name, provider, url, init, { + resolveAddresses: async () => { + privateResolveCalls += 1; + throw new Error("provider URL hostname quota-private.example.com resolves to a private-network address (10.0.0.7)"); + }, + pinnedPost: async () => { + privatePinnedCalls += 1; + return new Response("unexpected", { status: 500 }); + }, + }), + }, + ); + expect(blocked).toBeNull(); + expect(privateResolveCalls).toBe(1); + expect(privatePinnedCalls).toBe(0); + + let pinnedAddress = ""; + let pinnedAuthorization = ""; + let pinnedBody = ""; + const quota = await fetchAntigravityUsageQuota( + "public-dns-token", + "public-dns-project", + "https://quota-public.example.com", + { + outboundPost: (name, provider, url, init) => providerOutboundPost(name, provider, url, init, { + resolveAddresses: async () => ({ + hostname: "quota-public.example.com", + addresses: [{ address: "93.184.216.34", family: 4 }], + privateNetwork: false, + }), + pinnedPost: async (_url, address, body, _signal, requestOptions) => { + pinnedAddress = address.address; + pinnedAuthorization = new Headers(requestOptions?.headers).get("authorization") ?? ""; + pinnedBody = body; + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { + quotaInfo: { remainingFraction: 0.9, resetTime: "2026-07-05T14:00:00Z" }, + }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + }, + }), + }, + ); + expect(quota?.customWindows?.find(window => window.label === "Gem")?.percent).toBe(10); + expect(pinnedAddress).toBe("93.184.216.34"); + expect(pinnedAuthorization).toBe("Bearer public-dns-token"); + expect(JSON.parse(pinnedBody)).toEqual({ project: "public-dns-project" }); + }); + + test("destination-qualified Antigravity cache rows survive generation reconcile when account remains live", async () => { + const { saveCredential, getAccountSet } = await import("../src/oauth/store"); + await saveCredential("google-antigravity", { + access: "token-agy-reconcile", + refresh: "refresh-agy-reconcile", + expires: Date.now() + 3600_000, + projectId: "project-reconcile", + accountId: "acct-agy-reconcile", + email: "reconcile@example.com", + }); + + let probeCount = 0; + globalThis.fetch = (async () => { + probeCount += 1; + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { quotaInfo: { remainingFraction: 0.8, resetTime: "2026-07-05T14:00:00Z" } }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + const rows1 = await fetchProviderAccountQuotas("google-antigravity", true, "https://custom-dest.example.com"); + expect(rows1.length).toBe(1); + expect(probeCount).toBe(1); + expect(rows1[0]!.quota?.customWindows?.find(w => w.label === "Gem")?.percent).toBe(20); + + const set = getAccountSet("google-antigravity"); + expect(set?.accounts.length).toBe(1); + const liveAccountId = set!.accounts[0]!.id; + const providers: OcxConfig["providers"] = { + "google-antigravity": { + adapter: "google-antigravity", + authMode: "oauth", + baseUrl: "https://custom-dest.example.com", + }, + }; + const oauthAccountKeys = new Set([`google-antigravity\0${liveAccountId}`]); + reconcileProviderAccountQuotaRows({ + generation: 50_000, + providerNames: new Set(["google-antigravity"]), + comboIds: new Set(), + comboTargets: new Set(), + codexAccountIds: new Set(), + oauthAccountKeys, + providerAccountQuotaKeys: listLiveProviderAccountQuotaKeys(providers, oauthAccountKeys), + configRoots: new Set(), + }); + + const rows2 = await fetchProviderAccountQuotas("google-antigravity", false, "https://custom-dest.example.com"); + expect(rows2.length).toBe(1); + expect(probeCount).toBe(1); + expect(rows2[0]!.quota?.customWindows?.find(w => w.label === "Gem")?.percent).toBe(20); + }); + + test("forwards allowPrivateNetwork through account quota probing", async () => { + const { saveCredential } = await import("../src/oauth/store"); + await saveCredential("google-antigravity", { + access: "token-agy-priv", + refresh: "refresh-agy-priv", + expires: Date.now() + 3600_000, + projectId: "project-priv", + accountId: "acct-agy-priv", + email: "priv@example.com", + }); + + let requestedUrl = ""; + globalThis.fetch = (async (input: RequestInfo | URL) => { + requestedUrl = String(input); + return new Response(JSON.stringify({ + models: { + "gemini-3.7-flash": { + quotaInfo: { remainingFraction: 0.50, resetTime: "2026-07-05T14:00:00Z" }, + }, + }, + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + + const rows = await fetchProviderAccountQuotas("google-antigravity", true, "https://127.0.0.1:8443", true); + expect(rows.length).toBe(1); + expect(requestedUrl).toBe("https://127.0.0.1:8443/v1internal:fetchAvailableModels"); + expect(rows[0]?.quota?.customWindows?.find(w => w.label === "Gem")?.percent).toBe(50); + }); }); diff --git a/tests/request-pacing.test.ts b/tests/request-pacing.test.ts index bf8cfbb106..692be2a5f0 100644 --- a/tests/request-pacing.test.ts +++ b/tests/request-pacing.test.ts @@ -220,6 +220,7 @@ describe("provider request pacing queue", () => { comboTargets: new Set(), codexAccountIds: new Set(), oauthAccountKeys: new Set(), + providerAccountQuotaKeys: new Set(), configRoots: new Set(), })).toBe(1); diff --git a/tests/state-store-sweeper.test.ts b/tests/state-store-sweeper.test.ts index ea2d473b27..9e3241feeb 100644 --- a/tests/state-store-sweeper.test.ts +++ b/tests/state-store-sweeper.test.ts @@ -59,6 +59,7 @@ function context( comboTargets: new Set(), codexAccountIds: new Set(), oauthAccountKeys: new Set(), + providerAccountQuotaKeys: new Set(), configRoots: new Set(), ...overrides, };