From 1dd8d379b42a3a8d4a45398e67e4d7df330e0bd9 Mon Sep 17 00:00:00 2001 From: chilung Date: Wed, 19 Aug 2026 21:31:26 +0800 Subject: [PATCH 1/9] feat(quota): add per-account Gem/Cla quota probing for Google Antigravity (#1082) --- src/providers/quota.ts | 50 +++++++++++++++------- tests/provider-account-quota.test.ts | 62 ++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 16 deletions(-) diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 63c74b9cd6..e4bbfeda78 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -1452,7 +1452,7 @@ 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 accountCacheKey(provider: string, accountId: string): string { @@ -1571,8 +1571,16 @@ 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) throw new Error("google-antigravity account projectId missing"); + const token = await getTokenForAccountQuotaProbe(provider, accountId); + quota = await fetchAntigravityUsageQuota(token, stored.projectId); + } if (!quota) { // Preserve last-good bars and mark unavailable; advance TTL so failures // negative-cache instead of re-probing on every GUI poll. @@ -2129,17 +2137,13 @@ function antigravityUsedPercent(quotaInfo: Record): number | un 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; - try { - accessToken = await getValidAccessToken("google-antigravity"); - } catch { - return null; - } - const baseUrl = (config.baseUrl || "https://daily-cloudcode-pa.googleapis.com").replace(/\/+$/, ""); - const response = await fetch(`${baseUrl}/v1internal:fetchAvailableModels`, { +export async function fetchAntigravityUsageQuota( + accessToken: string, + projectId: string, + baseUrl = "https://daily-cloudcode-pa.googleapis.com", +): Promise { + const normalizedUrl = (baseUrl || "https://daily-cloudcode-pa.googleapis.com").replace(/\/+$/, ""); + const response = await fetch(`${normalizedUrl}/v1internal:fetchAvailableModels`, { method: "POST", headers: { Accept: "application/json", @@ -2147,7 +2151,7 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig "User-Agent": antigravityUserAgent(), Authorization: `Bearer ${accessToken}`, }, - body: JSON.stringify({ project: credential.projectId }), + body: JSON.stringify({ project: projectId }), signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), }); if (!response.ok) return null; @@ -2177,8 +2181,22 @@ 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 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); + if (!quota) return null; return report(provider, "google-antigravity:fetchAvailableModels", { - customWindows, + ...quota, updatedAt: Date.now(), }); } diff --git a/tests/provider-account-quota.test.ts b/tests/provider-account-quota.test.ts index 634c663ca2..0373da6ebb 100644 --- a/tests/provider-account-quota.test.ts +++ b/tests/provider-account-quota.test.ts @@ -418,4 +418,66 @@ 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 seenProjects: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + expect(String(input)).toContain("/v1internal:fetchAvailableModels"); + const body = JSON.parse(String(init?.body)) as { project?: string }; + seenProjects.push(body.project ?? ""); + 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(seenProjects.sort()).toEqual(["project-1", "project-2"]); + + const byProject = Object.fromEntries(rows.map(r => [ + r.quota?.customWindows?.find(w => w.label === "Gem")?.percent, + r.quota?.customWindows?.find(w => w.label === "Cla")?.percent, + ])); + // 1 - 0.64 = 36% used, 1 - 0.21 = 79% used for project 1 + // 1 - 0.90 = 10% used, 1 - 0.85 = 15% used for project 2 + expect(byProject[36]).toBe(79); + expect(byProject[10]).toBe(15); + }); }); From c186d28e589282105353870300ab5273f8c9f1cd Mon Sep 17 00:00:00 2001 From: chilung Date: Wed, 19 Aug 2026 22:48:41 +0800 Subject: [PATCH 2/9] fix(review): thread provider baseUrl to account probe and verify per-account token isolation --- src/providers/quota.ts | 5 +++-- tests/provider-account-quota.test.ts | 12 +++++++++--- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/src/providers/quota.ts b/src/providers/quota.ts index e4bbfeda78..f5eab49387 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -9,7 +9,7 @@ import type { StoredAccountQuota } from "../codex/quota"; import { isMainAccountIdentityGenerationLive } from "../codex/main-account-cache"; import { MAIN_CODEX_ACCOUNT_ID } from "../codex/main-account"; import { codexPlanKey } from "../codex/plan"; -import { resolveEnvValue } from "../config"; +import { loadConfig, resolveEnvValue } from "../config"; import { getValidAccessToken, getValidAccessTokenForAccount } from "../oauth"; import { getAccountCredential, getAccountSet, getCredential } from "../oauth/store"; import { antigravityUserAgent } from "../adapters/client-fingerprint"; @@ -1579,7 +1579,8 @@ async function fetchAccountQuota( const stored = getAccountCredential(provider, accountId); if (!stored?.projectId) throw new Error("google-antigravity account projectId missing"); const token = await getTokenForAccountQuotaProbe(provider, accountId); - quota = await fetchAntigravityUsageQuota(token, stored.projectId); + const provBaseUrl = loadConfig().providers?.["google-antigravity"]?.baseUrl; + quota = await fetchAntigravityUsageQuota(token, stored.projectId, provBaseUrl); } if (!quota) { // Preserve last-good bars and mark unavailable; advance TTL so failures diff --git a/tests/provider-account-quota.test.ts b/tests/provider-account-quota.test.ts index 0373da6ebb..4e15130f60 100644 --- a/tests/provider-account-quota.test.ts +++ b/tests/provider-account-quota.test.ts @@ -438,11 +438,14 @@ describe("fetchProviderAccountQuotas", () => { email: "agy2@example.com", }); - const seenProjects: string[] = []; + const seenRequests: Array<{ project: string; authorization: string | null }> = []; globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { expect(String(input)).toContain("/v1internal:fetchAvailableModels"); const body = JSON.parse(String(init?.body)) as { project?: string }; - seenProjects.push(body.project ?? ""); + seenRequests.push({ + project: body.project ?? "", + authorization: new Headers(init?.headers).get("Authorization"), + }); if (body.project === "project-1") { return new Response(JSON.stringify({ models: { @@ -469,7 +472,10 @@ describe("fetchProviderAccountQuotas", () => { const rows = await fetchProviderAccountQuotas("google-antigravity"); expect(rows.length).toBe(2); - expect(seenProjects.sort()).toEqual(["project-1", "project-2"]); + expect(seenRequests).toEqual(expect.arrayContaining([ + { project: "project-1", authorization: "Bearer token-agy-1" }, + { project: "project-2", authorization: "Bearer token-agy-2" }, + ])); const byProject = Object.fromEntries(rows.map(r => [ r.quota?.customWindows?.find(w => w.label === "Gem")?.percent, From 4b940a6509ef6b3745e059a9a1831662a2716217 Mon Sep 17 00:00:00 2001 From: chilung Date: Wed, 19 Aug 2026 16:09:39 +0000 Subject: [PATCH 3/9] fix(review): thread provider baseUrl from caller and skip projectId-less accounts silently (#1082) --- src/providers/quota.ts | 21 +++-- src/server/management/oauth-account-routes.ts | 2 +- tests/provider-account-quota.test.ts | 90 +++++++++++++++++++ 3 files changed, 107 insertions(+), 6 deletions(-) diff --git a/src/providers/quota.ts b/src/providers/quota.ts index f5eab49387..f10df4b3b0 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -9,7 +9,7 @@ import type { StoredAccountQuota } from "../codex/quota"; import { isMainAccountIdentityGenerationLive } from "../codex/main-account-cache"; import { MAIN_CODEX_ACCOUNT_ID } from "../codex/main-account"; import { codexPlanKey } from "../codex/plan"; -import { loadConfig, resolveEnvValue } from "../config"; +import { resolveEnvValue } from "../config"; import { getValidAccessToken, getValidAccessTokenForAccount } from "../oauth"; import { getAccountCredential, getAccountSet, getCredential } from "../oauth/store"; import { antigravityUserAgent } from "../adapters/client-fingerprint"; @@ -1561,6 +1561,7 @@ async function fetchAccountQuota( provider: string, accountId: string, forceRefresh: boolean, + baseUrl?: string, ): Promise { const key = accountCacheKey(provider, accountId); const writerGeneration = captureConfigGeneration(); @@ -1577,10 +1578,19 @@ async function fetchAccountQuota( quota = await fetchAnthropicUsageQuota(token); } else if (provider === "google-antigravity") { const stored = getAccountCredential(provider, accountId); - if (!stored?.projectId) throw new Error("google-antigravity account projectId missing"); + 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; + } const token = await getTokenForAccountQuotaProbe(provider, accountId); - const provBaseUrl = loadConfig().providers?.["google-antigravity"]?.baseUrl; - quota = await fetchAntigravityUsageQuota(token, stored.projectId, provBaseUrl); + quota = await fetchAntigravityUsageQuota(token, stored.projectId, baseUrl); } if (!quota) { // Preserve last-good bars and mark unavailable; advance TTL so failures @@ -1628,12 +1638,13 @@ async function fetchAccountQuota( export async function fetchProviderAccountQuotas( provider: string, forceRefresh = false, + baseUrl?: string, ): 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); return { accountId: account.id, quota: entry.quota, diff --git a/src/server/management/oauth-account-routes.ts b/src/server/management/oauth-account-routes.ts index ee57515b41..2c47797505 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); const byId = new Map(rows.map(row => [row.accountId, row])); const projected = projectAccounts(); return jsonResponse({ diff --git a/tests/provider-account-quota.test.ts b/tests/provider-account-quota.test.ts index 4e15130f60..5c3ecc7adc 100644 --- a/tests/provider-account-quota.test.ts +++ b/tests/provider-account-quota.test.ts @@ -486,4 +486,94 @@ describe("fetchProviderAccountQuotas", () => { expect(byProject[36]).toBe(79); expect(byProject[10]).toBe(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); + }); }); From b2e7f8be675c90964ed1bc2d4440fd6c6e80bd73 Mon Sep 17 00:00:00 2001 From: chilung Date: Thu, 20 Aug 2026 13:18:34 +0000 Subject: [PATCH 4/9] fix(quota): seed active account cache from provider report and clamp percentage bounds (#1082) --- src/providers/quota.ts | 10 ++++++++- tests/provider-account-quota.test.ts | 32 ++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/src/providers/quota.ts b/src/providers/quota.ts index f10df4b3b0..4dceebd17b 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -2146,7 +2146,7 @@ function antigravityUsedPercent(quotaInfo: Record): number | un ? toFiniteNumber(quotaInfo.remainingPercentage)! * 100 : undefined); if (remaining === undefined) return undefined; - return normalizePercent(100 - remaining); + return Math.min(100, Math.max(0, normalizePercent(100 - remaining) ?? 0)); } export async function fetchAntigravityUsageQuota( @@ -2207,6 +2207,14 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig } const quota = await fetchAntigravityUsageQuota(accessToken, credential.projectId, config.baseUrl); if (!quota) return null; + const probedAccountId = getAccountSet("google-antigravity")?.activeAccountId; + if (probedAccountId) { + const probedAccountKey = accountCacheKey("google-antigravity", probedAccountId); + const stillOwnsToken = getAccountCredential("google-antigravity", probedAccountId)?.access === accessToken; + if (stillOwnsToken && mayCommitAccountQuotaKey(probedAccountKey, captureConfigGeneration())) { + accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota }); + } + } return report(provider, "google-antigravity:fetchAvailableModels", { ...quota, updatedAt: Date.now(), diff --git a/tests/provider-account-quota.test.ts b/tests/provider-account-quota.test.ts index 5c3ecc7adc..3a3ba76742 100644 --- a/tests/provider-account-quota.test.ts +++ b/tests/provider-account-quota.test.ts @@ -576,4 +576,36 @@ describe("fetchProviderAccountQuotas", () => { 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(); + }); }); From 2222a77526ce564e399feae4ba4efe000b75caed Mon Sep 17 00:00:00 2001 From: chilung Date: Fri, 21 Aug 2026 12:02:09 +0800 Subject: [PATCH 5/9] fix(quota): destination-bound account cache, pre-token policy check, and fail-closed redirects (#1082) --- src/providers/quota.ts | 74 +++++++++++---- tests/provider-account-quota.test.ts | 129 +++++++++++++++++++++++++++ 2 files changed, 185 insertions(+), 18 deletions(-) diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 4dceebd17b..3a0764ff77 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -1,3 +1,4 @@ +import { providerDestinationConfigError } from "../lib/destination-policy"; import { createHash } from "node:crypto"; import { effectiveCodexAuthAccountId, @@ -1455,8 +1456,18 @@ export function supportsPerAccountQuota(provider: string): boolean { return provider === "anthropic" || provider === "google-antigravity"; } -function accountCacheKey(provider: string, accountId: string): string { - return `${provider}\u0000${accountId}`; +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; } /** @@ -1563,7 +1574,8 @@ async function fetchAccountQuota( forceRefresh: boolean, baseUrl?: string, ): 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; @@ -1589,6 +1601,18 @@ async function fetchAccountQuota( } 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", + }); + 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); } @@ -2154,18 +2178,28 @@ export async function fetchAntigravityUsageQuota( projectId: string, baseUrl = "https://daily-cloudcode-pa.googleapis.com", ): Promise { - const normalizedUrl = (baseUrl || "https://daily-cloudcode-pa.googleapis.com").replace(/\/+$/, ""); - const response = await fetch(`${normalizedUrl}/v1internal:fetchAvailableModels`, { - method: "POST", - 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), + const destError = providerDestinationConfigError("google-antigravity", { + baseUrl: baseUrl || "https://daily-cloudcode-pa.googleapis.com", }); + if (destError) return null; + const normalizedUrl = (baseUrl || "https://daily-cloudcode-pa.googleapis.com").replace(/\/+$/, ""); + let response: Response; + try { + response = await fetch(normalizedUrl + "/v1internal:fetchAvailableModels", { + method: "POST", + headers: { + Accept: "application/json", + "Content-Type": "application/json", + "User-Agent": antigravityUserAgent(), + Authorization: "Bearer " + accessToken, + }, + body: JSON.stringify({ project: projectId }), + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + } catch { + return null; + } if (!response.ok) return null; const body = asRecord(await readQuotaJson(response)); const models = asRecord(body?.models); @@ -2197,6 +2231,12 @@ export async function fetchAntigravityUsageQuota( } 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; @@ -2207,11 +2247,9 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig } const quota = await fetchAntigravityUsageQuota(accessToken, credential.projectId, config.baseUrl); if (!quota) return null; - const probedAccountId = getAccountSet("google-antigravity")?.activeAccountId; - if (probedAccountId) { - const probedAccountKey = accountCacheKey("google-antigravity", probedAccountId); + if (probedAccountId && probedAccountKey) { const stillOwnsToken = getAccountCredential("google-antigravity", probedAccountId)?.access === accessToken; - if (stillOwnsToken && mayCommitAccountQuotaKey(probedAccountKey, captureConfigGeneration())) { + if (stillOwnsToken && mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) { accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota }); } } diff --git a/tests/provider-account-quota.test.ts b/tests/provider-account-quota.test.ts index 3a3ba76742..81b1479017 100644 --- a/tests/provider-account-quota.test.ts +++ b/tests/provider-account-quota.test.ts @@ -608,4 +608,133 @@ describe("fetchProviderAccountQuotas", () => { 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("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); + }); + }); From ae19382be1497cde5543b128cd24c23aa49a2e62 Mon Sep 17 00:00:00 2001 From: chilung Date: Sat, 22 Aug 2026 01:48:13 +0000 Subject: [PATCH 6/9] fix(quota): destination-bound key liveness reconciliation and pinned outbound transport (#1082) --- src/providers/quota.ts | 60 +++++++++++++++++++--------- tests/provider-account-quota.test.ts | 44 ++++++++++++++++++++ 2 files changed, 86 insertions(+), 18 deletions(-) diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 3a0764ff77..2b6331a557 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -1,4 +1,7 @@ import { providerDestinationConfigError } from "../lib/destination-policy"; +import { providerOutboundPost, providerRedirectError } from "../lib/provider-outbound"; + +const originalFetch = globalThis.fetch; import { createHash } from "node:crypto"; import { effectiveCodexAuthAccountId, @@ -1436,8 +1439,16 @@ let lastReconciledGeneration = 0; let liveAccountQuotaKeys = new Set(); let liveProviderQuotaKeys = new Set(); +function canonicalAccountQuotaKey(key: string): string { + const firstNull = key.indexOf("\0"); + if (firstNull === -1) return key; + const secondNull = key.indexOf("\0", firstNull + 1); + if (secondNull === -1) return key; + return key.slice(0, secondNull); +} + function mayCommitAccountQuotaKey(key: string, writerGeneration: number): boolean { - return writerGeneration >= lastReconciledGeneration || liveAccountQuotaKeys.has(key); + return writerGeneration >= lastReconciledGeneration || liveAccountQuotaKeys.has(canonicalAccountQuotaKey(key)); } function mayCommitProviderQuotaKey(key: string, writerGeneration: number): boolean { @@ -1507,7 +1518,8 @@ export function reconcileProviderAccountQuotaRows(context: GenerationContext): n if (context.generation <= lastReconciledGeneration) return 0; let removed = 0; for (const key of accountQuotaCache.keys()) { - if (context.oauthAccountKeys.has(key)) continue; + const canonical = canonicalAccountQuotaKey(key); + if (context.oauthAccountKeys.has(canonical)) continue; accountQuotaCache.delete(key); removed += 1; } @@ -2177,30 +2189,42 @@ 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 destError = providerDestinationConfigError("google-antigravity", { - baseUrl: baseUrl || "https://daily-cloudcode-pa.googleapis.com", - }); - if (destError) return null; + const outboundPost = options?.outboundPost ?? providerOutboundPost; const normalizedUrl = (baseUrl || "https://daily-cloudcode-pa.googleapis.com").replace(/\/+$/, ""); + const targetUrl = normalizedUrl + "/v1internal:fetchAvailableModels"; let response: Response; try { - response = await fetch(normalizedUrl + "/v1internal:fetchAvailableModels", { - method: "POST", - headers: { - Accept: "application/json", - "Content-Type": "application/json", - "User-Agent": antigravityUserAgent(), - Authorization: "Bearer " + accessToken, + const activeFetch = options?.fetch ?? (globalThis.fetch !== originalFetch ? globalThis.fetch : undefined); + response = await outboundPost( + "google-antigravity", + { + baseUrl: normalizedUrl, + allowPrivateNetwork: options?.allowPrivateNetwork ?? false, + ...(activeFetch ? { fetch: activeFetch } : {}), }, - body: JSON.stringify({ project: projectId }), - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); + 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; } - 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; diff --git a/tests/provider-account-quota.test.ts b/tests/provider-account-quota.test.ts index 81b1479017..1c1ca84a18 100644 --- a/tests/provider-account-quota.test.ts +++ b/tests/provider-account-quota.test.ts @@ -737,4 +737,48 @@ describe("fetchProviderAccountQuotas", () => { expect(redirectTargetCalled).toBe(false); }); + 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; + reconcileProviderAccountQuotaRows({ + generation: 50_000, + providerNames: new Set(["google-antigravity"]), + comboIds: new Set(), + comboTargets: new Set(), + codexAccountIds: new Set(), + oauthAccountKeys: new Set([`google-antigravity\0${liveAccountId}`]), + 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); + }); }); From aba7a8bdb744e86c5fd3d1b784e11318cd18b8ed Mon Sep 17 00:00:00 2001 From: chilung Date: Mon, 24 Aug 2026 15:14:29 +0000 Subject: [PATCH 7/9] =?UTF-8?q?fix(quota):=20=E7=B6=81=E5=AE=9A=E5=B8=B3?= =?UTF-8?q?=E6=88=B6=E5=BF=AB=E5=8F=96=E7=9A=84=E7=9B=AE=E7=9A=84=E5=9C=B0?= =?UTF-8?q?=E7=94=9F=E5=91=BD=E9=80=B1=E6=9C=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 拒絕目的地切換後才完成的舊 generation writer，並讓 Antigravity 快取讀取與 liveness context 使用同一目的地 key。 補上 quota 到 pinned outbound transport 的 DNS 邊界回歸。 --- src/lib/state-store-registrations.ts | 6 +- src/lib/state-store-sweeper.ts | 1 + src/providers/quota.ts | 54 +++++--- tests/combos.test.ts | 2 + tests/oauth-store-multi.test.ts | 1 + tests/provider-account-quota.test.ts | 193 ++++++++++++++++++++++++++- tests/request-pacing.test.ts | 1 + tests/state-store-sweeper.test.ts | 1 + 8 files changed, 240 insertions(+), 19 deletions(-) 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 2b6331a557..44fdc1576e 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -1436,19 +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 canonicalAccountQuotaKey(key: string): string { - const firstNull = key.indexOf("\0"); - if (firstNull === -1) return key; - const secondNull = key.indexOf("\0", firstNull + 1); - if (secondNull === -1) return key; - return key.slice(0, secondNull); -} - function mayCommitAccountQuotaKey(key: string, writerGeneration: number): boolean { - return writerGeneration >= lastReconciledGeneration || liveAccountQuotaKeys.has(canonicalAccountQuotaKey(key)); + const liveSinceGeneration = liveAccountQuotaKeys.get(key); + return writerGeneration >= lastReconciledGeneration + || (liveSinceGeneration !== undefined && writerGeneration >= liveSinceGeneration); } function mayCommitProviderQuotaKey(key: string, writerGeneration: number): boolean { @@ -1481,12 +1475,34 @@ function accountCacheKey(provider: string, accountId: string, destination = ""): return destination ? provider + "\u0000" + accountId + "\u0000" + destination : 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; } @@ -1516,10 +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()) { - const canonical = canonicalAccountQuotaKey(key); - if (context.oauthAccountKeys.has(canonical)) continue; + if (context.providerAccountQuotaKeys.has(key)) continue; accountQuotaCache.delete(key); removed += 1; } @@ -1528,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; @@ -1537,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(); } 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 1c1ca84a18..95c82afd26 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(); @@ -715,6 +727,114 @@ describe("fetchProviderAccountQuotas", () => { 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"); @@ -737,6 +857,68 @@ describe("fetchProviderAccountQuotas", () => { 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", { @@ -766,13 +948,22 @@ describe("fetchProviderAccountQuotas", () => { 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: new Set([`google-antigravity\0${liveAccountId}`]), + oauthAccountKeys, + providerAccountQuotaKeys: listLiveProviderAccountQuotaKeys(providers, oauthAccountKeys), configRoots: new Set(), }); 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, }; From 511ef0737c4ed4b3fe3f290d3952013b0a0a32f4 Mon Sep 17 00:00:00 2001 From: chilung Date: Mon, 24 Aug 2026 17:40:59 +0000 Subject: [PATCH 8/9] test(quota): assert Antigravity auth per project --- tests/provider-account-quota.test.ts | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/tests/provider-account-quota.test.ts b/tests/provider-account-quota.test.ts index 95c82afd26..9f00ffd166 100644 --- a/tests/provider-account-quota.test.ts +++ b/tests/provider-account-quota.test.ts @@ -450,14 +450,14 @@ describe("fetchProviderAccountQuotas", () => { email: "agy2@example.com", }); - const seenRequests: Array<{ project: string; authorization: string | null }> = []; + 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 }; - seenRequests.push({ - project: body.project ?? "", - authorization: new Headers(init?.headers).get("Authorization"), - }); + seenAuthorizationByProject.set( + body.project ?? "", + new Headers(init?.headers).get("Authorization"), + ); if (body.project === "project-1") { return new Response(JSON.stringify({ models: { @@ -484,10 +484,8 @@ describe("fetchProviderAccountQuotas", () => { const rows = await fetchProviderAccountQuotas("google-antigravity"); expect(rows.length).toBe(2); - expect(seenRequests).toEqual(expect.arrayContaining([ - { project: "project-1", authorization: "Bearer token-agy-1" }, - { project: "project-2", authorization: "Bearer token-agy-2" }, - ])); + expect(seenAuthorizationByProject.get("project-1")).toBe("Bearer token-agy-1"); + expect(seenAuthorizationByProject.get("project-2")).toBe("Bearer token-agy-2"); const byProject = Object.fromEntries(rows.map(r => [ r.quota?.customWindows?.find(w => w.label === "Gem")?.percent, From aa26492ef9eb0af0ac8a7e93437b7e9f04881d96 Mon Sep 17 00:00:00 2001 From: chilung Date: Tue, 25 Aug 2026 17:13:48 +0000 Subject: [PATCH 9/9] fix(quota): propagate allowPrivateNetwork and isolate account quota assertions (#1082) --- src/providers/quota.ts | 13 ++++-- src/server/management/oauth-account-routes.ts | 2 +- tests/provider-account-quota.test.ts | 46 +++++++++++++++---- 3 files changed, 47 insertions(+), 14 deletions(-) diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 44fdc1576e..9e9e118b53 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -1,7 +1,7 @@ +const originalFetch = globalThis.fetch; import { providerDestinationConfigError } from "../lib/destination-policy"; import { providerOutboundPost, providerRedirectError } from "../lib/provider-outbound"; -const originalFetch = globalThis.fetch; import { createHash } from "node:crypto"; import { effectiveCodexAuthAccountId, @@ -1605,8 +1605,9 @@ async function getTokenForAccountQuotaProbe(provider: string, accountId: string) async function fetchAccountQuota( provider: string, accountId: string, - forceRefresh: boolean, + forceRefresh = false, baseUrl?: string, + allowPrivateNetwork?: boolean, ): Promise { const normalizedDest = provider === "google-antigravity" ? normalizeAntigravityDestination(baseUrl) : ""; const key = accountCacheKey(provider, accountId, normalizedDest); @@ -1638,6 +1639,7 @@ async function fetchAccountQuota( // 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 }; @@ -1648,7 +1650,7 @@ async function fetchAccountQuota( return entry; } const token = await getTokenForAccountQuotaProbe(provider, accountId); - quota = await fetchAntigravityUsageQuota(token, stored.projectId, baseUrl); + quota = await fetchAntigravityUsageQuota(token, stored.projectId, baseUrl, { allowPrivateNetwork }); } if (!quota) { // Preserve last-good bars and mark unavailable; advance TTL so failures @@ -1697,12 +1699,13 @@ 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, baseUrl); + const entry = await fetchAccountQuota(provider, account.id, forceRefresh, baseUrl, allowPrivateNetwork); return { accountId: account.id, quota: entry.quota, @@ -2291,7 +2294,7 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig } catch { return null; } - const quota = await fetchAntigravityUsageQuota(accessToken, credential.projectId, config.baseUrl); + 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; diff --git a/src/server/management/oauth-account-routes.ts b/src/server/management/oauth-account-routes.ts index 2c47797505..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, ctx.config.providers[provider]?.baseUrl); + 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/provider-account-quota.test.ts b/tests/provider-account-quota.test.ts index 9f00ffd166..b0f894417f 100644 --- a/tests/provider-account-quota.test.ts +++ b/tests/provider-account-quota.test.ts @@ -487,14 +487,15 @@ describe("fetchProviderAccountQuotas", () => { expect(seenAuthorizationByProject.get("project-1")).toBe("Bearer token-agy-1"); expect(seenAuthorizationByProject.get("project-2")).toBe("Bearer token-agy-2"); - const byProject = Object.fromEntries(rows.map(r => [ - r.quota?.customWindows?.find(w => w.label === "Gem")?.percent, - r.quota?.customWindows?.find(w => w.label === "Cla")?.percent, - ])); - // 1 - 0.64 = 36% used, 1 - 0.21 = 79% used for project 1 - // 1 - 0.90 = 10% used, 1 - 0.85 = 15% used for project 2 - expect(byProject[36]).toBe(79); - expect(byProject[10]).toBe(15); + 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 () => { @@ -970,4 +971,33 @@ describe("fetchProviderAccountQuotas", () => { 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); + }); });