diff --git a/changelog.d/fixes/8965-fix.plan.md b/changelog.d/fixes/8965-fix.plan.md new file mode 100644 index 0000000000..d557cd7027 --- /dev/null +++ b/changelog.d/fixes/8965-fix.plan.md @@ -0,0 +1 @@ +- fix(providers): switch Antigravity quota RPCs to iterate ANTIGRAVITY_RUNTIME_BASE_URLS (#8965) \ No newline at end of file diff --git a/open-sse/services/usage/antigravity.ts b/open-sse/services/usage/antigravity.ts index 7b2f2ce13a..693771d681 100644 --- a/open-sse/services/usage/antigravity.ts +++ b/open-sse/services/usage/antigravity.ts @@ -272,21 +272,24 @@ async function fetchAntigravityUserQuotaCached( const promise = (async () => { try { - const response = await fetch( - "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota", - { - method: "POST", - headers: getAntigravityContentHeaders(clientProfile, accessToken), - body: JSON.stringify({ project: projectId }), - signal: AbortSignal.timeout(10000), - } - ); + for (const baseUrl of ANTIGRAVITY_RUNTIME_BASE_URLS) { + const response = await fetch( + `${baseUrl}/v1internal:retrieveUserQuota`, + { + method: "POST", + headers: getAntigravityContentHeaders(clientProfile, accessToken), + body: JSON.stringify({ project: projectId }), + signal: AbortSignal.timeout(10000), + } + ); - if (!response.ok) return null; + if (!response.ok) continue; - const data = await response.json(); - _antigravityUserQuotaCache.set(cacheKey, { data, fetchedAt: Date.now() }); - return data; + const data = await response.json(); + _antigravityUserQuotaCache.set(cacheKey, { data, fetchedAt: Date.now() }); + return data; + } + return null; } catch { return null; } diff --git a/open-sse/services/usage/antigravityWeeklyQuota.ts b/open-sse/services/usage/antigravityWeeklyQuota.ts index 3aa4e78d18..a806eb4645 100644 --- a/open-sse/services/usage/antigravityWeeklyQuota.ts +++ b/open-sse/services/usage/antigravityWeeklyQuota.ts @@ -17,6 +17,7 @@ * `fetchAntigravityUserQuotaCached` pattern. */ +import { ANTIGRAVITY_RUNTIME_BASE_URLS } from "../../config/antigravityUpstream.ts"; import { toRecord, toNumber } from "./scalars.ts"; import { type UsageQuota, parseResetTime } from "./quota.ts"; import { getAntigravityContentHeaders } from "../antigravityHeaders.ts"; @@ -81,21 +82,24 @@ export async function fetchAntigravityUserQuotaSummaryCached( const promise = (async () => { try { - const response = await fetch( - "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary", - { - method: "POST", - headers: getAntigravityContentHeaders(clientProfile, accessToken), - body: JSON.stringify({ project: projectId }), - signal: AbortSignal.timeout(10000), - } - ); + for (const baseUrl of ANTIGRAVITY_RUNTIME_BASE_URLS) { + const response = await fetch( + `${baseUrl}/v1internal:retrieveUserQuotaSummary`, + { + method: "POST", + headers: getAntigravityContentHeaders(clientProfile, accessToken), + body: JSON.stringify({ project: projectId }), + signal: AbortSignal.timeout(10000), + } + ); - if (!response.ok) return null; + if (!response.ok) continue; - const data = await response.json(); - _weeklyQuotaCache.set(cacheKey, { data, fetchedAt: Date.now() }); - return data; + const data = await response.json(); + _weeklyQuotaCache.set(cacheKey, { data, fetchedAt: Date.now() }); + return data; + } + return null; } catch { return null; } diff --git a/tests/unit/antigravity-quota-host-8965.test.ts b/tests/unit/antigravity-quota-host-8965.test.ts new file mode 100644 index 0000000000..a0e61832bc --- /dev/null +++ b/tests/unit/antigravity-quota-host-8965.test.ts @@ -0,0 +1,263 @@ +/** + * #8965 — Antigravity quota reads must use the runtime host (daily-cloudcode-pa) + * instead of hardcoding cloudcode-pa.googleapis.com. + * + * Antigravity inference, credit probe, OAuth, and the models catalog all use + * ANTIGRAVITY_RUNTIME_BASE_URLS which starts with daily-cloudcode-pa.googleapis.com. + * The two quota RPCs (retrieveUserQuota, retrieveUserQuotaSummary) were hardcoded + * to cloudcode-pa.googleapis.com, so when only the runtime host serves them, the + * live quota signal is lost and falls back to fetchAvailableModels. + * + * This regression test stubs globalThis.fetch so ONLY daily-cloudcode-pa serves + * the RPCs (cloudcode-pa returns 500), then asserts: + * 1. retrieveUserQuota is the quota source (not fetchAvailableModels) + * 2. Weekly bucket data is populated (not lost) + */ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-ag-host-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = "test-ag-host-secret"; + +const core = await import("../../src/lib/db/core.ts"); +const usageModule = await import("../../open-sse/services/usage.ts"); +const { getUsageForProvider } = usageModule; + +const originalFetch = globalThis.fetch; + +test.after(() => { + globalThis.fetch = originalFetch; + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +const RESET_IN_2_HOURS = new Date(Date.now() + 2 * 60 * 60 * 1000).toISOString(); +const RESET_IN_3_DAYS = new Date(Date.now() + 3 * 24 * 60 * 60 * 1000).toISOString(); + +interface UsageResult { + quotas: Record< + string, + { + remainingPercentage?: number; + resetAt: string | null; + unlimited: boolean; + quotaSource?: string; + } + >; +} + +test("#8965: quota reads use the runtime host (daily-cloudcode-pa), not cloudcode-pa", async () => { + core.resetDbInstance(); + + const dailyCount = { value: 0 }; + const cloudcodeCount = { value: 0 }; + + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + + if (url.includes("daily-cloudcode-pa.googleapis.com")) { + dailyCount.value++; + + if (url.includes("retrieveUserQuotaSummary")) { + return { + ok: true, + json: async () => ({ + groups: [ + { + displayName: "Gemini Models", + buckets: [ + { + bucketId: "gemini-weekly", + displayName: "Weekly Quota", + remainingFraction: 0.6, + resetTime: RESET_IN_3_DAYS, + }, + ], + }, + ], + }), + } as Response; + } + + if (url.includes("retrieveUserQuota")) { + return { + ok: true, + json: async () => ({ + buckets: [ + { + modelId: "gemini-3-flash-agent", + remainingFraction: 0.4, + resetTime: RESET_IN_2_HOURS, + }, + ], + }), + } as Response; + } + + if (url.includes("fetchAvailableModels")) { + return { + ok: true, + json: async () => ({ + models: { + "gemini-3-flash-agent": { + quotaInfo: { remainingFraction: 1.0, resetTime: RESET_IN_2_HOURS }, + }, + "gemini-3.5-flash-low": { + quotaInfo: { remainingFraction: 0.8, resetTime: RESET_IN_2_HOURS }, + }, + }, + }), + } as Response; + } + + // subscription info + return { + ok: true, + json: async () => ({ + cloudaicompanionProject: { id: "test-project" }, + tierId: "FREE", + subscriptionType: "free", + }), + } as Response; + } + + if (url.includes("cloudcode-pa.googleapis.com")) { + cloudcodeCount.value++; + return { ok: false, status: 500, json: async () => ({}) } as Response; + } + + // Default: return 500 for anything else + return { ok: false, status: 500, json: async () => ({}) } as Response; + }) as typeof fetch; + + const connection = { + id: "conn-host-8965", + provider: "antigravity", + accessToken: "fake-token-host-test-8965", + providerSpecificData: { clientProfile: "cli" }, + projectId: "test-project", + }; + + const result = await getUsageForProvider(connection, { forceRefresh: true }); + assert.ok(result && "quotas" in result, "should return quotas"); + const quotas = (result as UsageResult).quotas; + + // The per-model quota should come from retrieveUserQuota (the live source), + // NOT fetchAvailableModels (the stale catalog fallback). + assert.ok(quotas["gemini-3-flash-agent"], "gemini-3-flash-agent quota present"); + assert.equal( + quotas["gemini-3-flash-agent"].quotaSource, + "retrieveUserQuota", + "quota source is retrieveUserQuota (live), not fetchAvailableModels" + ); + + // The weekly group quota should also be populated. + assert.ok(quotas.gemini_weekly, "weekly group quota merged in"); + assert.equal(quotas.gemini_weekly.remainingPercentage, 60); + + // The runtime host should have been used for the quota RPCs. + assert.ok(dailyCount.value > 0, "daily-cloudcode-pa was called at least once"); +}); + +test("#8965 behavioral impact: live quota source + weekly bucket unreachable when only runtime host serves", async () => { + core.resetDbInstance(); + + globalThis.fetch = (async (input: RequestInfo | URL) => { + const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + + if (url.includes("daily-cloudcode-pa.googleapis.com")) { + if (url.includes("retrieveUserQuotaSummary")) { + return { + ok: true, + json: async () => ({ + groups: [ + { + displayName: "Gemini Models", + buckets: [ + { + bucketId: "gemini-weekly", + displayName: "Weekly Quota", + remainingFraction: 0.6, + resetTime: RESET_IN_3_DAYS, + }, + ], + }, + ], + }), + } as Response; + } + + if (url.includes("retrieveUserQuota")) { + return { + ok: true, + json: async () => ({ + buckets: [ + { + modelId: "gemini-3-flash-agent", + remainingFraction: 0.4, + resetTime: RESET_IN_2_HOURS, + }, + ], + }), + } as Response; + } + + if (url.includes("fetchAvailableModels")) { + return { + ok: true, + json: async () => ({ + models: { + "gemini-3-flash-agent": { + quotaInfo: { remainingFraction: 1.0, resetTime: RESET_IN_2_HOURS }, + }, + }, + }), + } as Response; + } + + // subscription info + return { + ok: true, + json: async () => ({ + cloudaicompanionProject: { id: "test-project" }, + tierId: "FREE", + subscriptionType: "free", + }), + } as Response; + } + + if (url.includes("cloudcode-pa.googleapis.com")) { + return { ok: false, status: 500, json: async () => ({}) } as Response; + } + + return { ok: false, status: 500, json: async () => ({}) } as Response; + }) as typeof fetch; + + const connection = { + id: "conn-host-8965-impact", + provider: "antigravity", + accessToken: "fake-token-host-impact", + providerSpecificData: { clientProfile: "cli" }, + projectId: "test-project", + }; + + const result = await getUsageForProvider(connection, { forceRefresh: true }); + assert.ok(result && "quotas" in result, "should return quotas"); + const quotas = (result as UsageResult).quotas; + + // The per-model quota MUST come from retrieveUserQuota — the live signal. + assert.ok(quotas["gemini-3-flash-agent"], "gemini-3-flash-agent quota present"); + assert.equal( + quotas["gemini-3-flash-agent"].quotaSource, + "retrieveUserQuota", + "quota source is retrieveUserQuota (live), not fetchAvailableModels" + ); + + // The weekly group quota MUST also be present because retrieveUserQuotaSummary + // was served by the runtime host. + assert.ok(quotas.gemini_weekly, "weekly group quota present"); +}); \ No newline at end of file