From 6d6a4abb8a204e56a56b0438f802bdb12b5e234c Mon Sep 17 00:00:00 2001 From: Raxxoor Date: Wed, 29 Apr 2026 03:01:23 +0100 Subject: [PATCH] fix(antigravity): stabilize streaming and usage refresh (#1748) Integrated into release/v3.7.5 (PR #1748) --- open-sse/config/antigravityModelAliases.ts | 12 ++ open-sse/services/usage.ts | 182 +++++++++++++------ open-sse/utils/streamReadiness.ts | 19 +- src/app/api/providers/[id]/models/route.ts | 94 ++++++---- src/lib/usage/providerLimits.ts | 22 ++- tests/unit/antigravity-model-aliases.test.ts | 10 + tests/unit/stream-readiness.test.ts | 12 ++ tests/unit/usage-fetcher-antigravity.test.ts | 59 ------ tests/unit/usage-service-hardening.test.ts | 82 ++++++++- 9 files changed, 319 insertions(+), 173 deletions(-) delete mode 100644 tests/unit/usage-fetcher-antigravity.test.ts diff --git a/open-sse/config/antigravityModelAliases.ts b/open-sse/config/antigravityModelAliases.ts index 745249af37..e01375bb12 100644 --- a/open-sse/config/antigravityModelAliases.ts +++ b/open-sse/config/antigravityModelAliases.ts @@ -39,6 +39,11 @@ const CLIENT_VISIBLE_MODEL_NAMES = Object.freeze( }, {}) ); +const PUBLIC_MODEL_IDS = new Set(ANTIGRAVITY_PUBLIC_MODELS.map((model) => model.id)); +const UPSTREAM_PUBLIC_MODEL_IDS = new Set( + ANTIGRAVITY_PUBLIC_MODELS.map((model) => resolveAntigravityModelId(model.id)) +); + export function resolveAntigravityModelId(modelId: string): string { if (!modelId) return modelId; return (ANTIGRAVITY_MODEL_ALIASES as AntigravityModelAliasMap)[modelId] || modelId; @@ -55,3 +60,10 @@ export function getClientVisibleAntigravityModelName( ): string { return CLIENT_VISIBLE_MODEL_NAMES[modelId] || fallbackName || modelId; } + +export function isUserCallableAntigravityModelId(modelId: string): boolean { + if (!modelId) return false; + const clientId = toClientAntigravityModelId(modelId); + const upstreamId = resolveAntigravityModelId(modelId); + return PUBLIC_MODEL_IDS.has(clientId) || UPSTREAM_PUBLIC_MODEL_IDS.has(upstreamId); +} diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index 3122403277..527774b13c 100644 --- a/open-sse/services/usage.ts +++ b/open-sse/services/usage.ts @@ -8,6 +8,7 @@ import { getAntigravityFetchAvailableModelsUrls, ANTIGRAVITY_BASE_URLS, } from "../config/antigravityUpstream.ts"; +import { isUserCallableAntigravityModelId } from "../config/antigravityModelAliases.ts"; import { getGlmQuotaUrl } from "../config/glmProvider.ts"; import { CURSOR_REGISTRY_VERSION, @@ -599,7 +600,7 @@ async function getBailianCodingPlanUsage( * @param {Object} connection - Provider connection with accessToken * @returns {Promise} Usage data with quotas */ -export async function getUsageForProvider(connection) { +export async function getUsageForProvider(connection, options: { forceRefresh?: boolean } = {}) { const { id, provider, accessToken, apiKey, providerSpecificData, projectId, email } = connection; switch (provider) { @@ -608,7 +609,7 @@ export async function getUsageForProvider(connection) { case "gemini-cli": return await getGeminiUsage(accessToken, providerSpecificData, projectId); case "antigravity": - return await getAntigravityUsage(accessToken, providerSpecificData, projectId, id); + return await getAntigravityUsage(accessToken, providerSpecificData, projectId, id, options); case "claude": return await getClaudeUsage(accessToken); case "codex": @@ -1205,6 +1206,80 @@ function getGeminiCliPlanLabel(subscriptionInfo) { // Key: truncated accessToken → { data, fetchedAt } const _antigravitySubCache = new Map(); const ANTIGRAVITY_CACHE_TTL_MS = 5 * 60 * 1000; // 5 minutes +const ANTIGRAVITY_MODELS_CACHE_TTL_MS = 60 * 1000; +const ANTIGRAVITY_CREDIT_PROBE_TTL_MS = 5 * 60 * 1000; +const _antigravityAvailableModelsCache = new Map(); +const _antigravityAvailableModelsInflight = new Map>(); +const _antigravityCreditProbeCache = new Map(); +const _antigravityCreditProbeInflight = new Map>(); + +interface AntigravityUsageOptions { + forceRefresh?: boolean; +} + +function buildAntigravityUsageCacheKey(accessToken: string, projectId?: string | null): string { + return `${accessToken.substring(0, 16)}:${projectId || "default"}`; +} + +async function fetchAntigravityAvailableModelsCached( + accessToken: string, + projectId?: string | null, + options: AntigravityUsageOptions = {} +): Promise { + if (!accessToken) throw new Error("Access token is required"); + + const cacheKey = buildAntigravityUsageCacheKey(accessToken, projectId); + const cached = _antigravityAvailableModelsCache.get(cacheKey); + if (!options.forceRefresh && cached && Date.now() - cached.fetchedAt < ANTIGRAVITY_MODELS_CACHE_TTL_MS) { + return cached.data; + } + + const inflight = _antigravityAvailableModelsInflight.get(cacheKey); + if (inflight) return inflight; + + const promise = (async () => { + let response: Response | null = null; + let lastError: Error | null = null; + + for (const quotaApiUrl of ANTIGRAVITY_CONFIG.quotaApiUrls) { + try { + response = await fetch(quotaApiUrl, { + method: "POST", + headers: getAntigravityHeaders("fetchAvailableModels", accessToken), + body: JSON.stringify(projectId ? { project: projectId } : {}), + signal: AbortSignal.timeout(10000), + }); + + if (response.ok || response.status === 401 || response.status === 403) { + break; + } + } catch (error) { + lastError = error as Error; + } + } + + if (!response) { + throw lastError || new Error("Antigravity API unavailable"); + } + + if (response.status === 403) { + return { __antigravityForbidden: true }; + } + + if (!response.ok) { + throw new Error(`Antigravity API error: ${response.status}`); + } + + const data = await response.json(); + _antigravityAvailableModelsCache.set(cacheKey, { data, fetchedAt: Date.now() }); + return data; + })().finally(() => { + _antigravityAvailableModelsInflight.delete(cacheKey); + }); + + _antigravityAvailableModelsInflight.set(cacheKey, promise); + return promise; +} /** * Map raw loadCodeAssist tier data to short display labels. @@ -1277,6 +1352,40 @@ function getAntigravityPlanLabel(subscriptionInfo) { * Returns the credit balance, or null if the probe failed. */ async function probeAntigravityCreditBalance( + accessToken: string, + accountId: string, + projectId?: string | null, + options: AntigravityUsageOptions = {} +): Promise { + if (!accessToken) return null; + + const cacheKey = buildAntigravityUsageCacheKey(accessToken, projectId || accountId); + const cached = _antigravityCreditProbeCache.get(cacheKey); + if (!options.forceRefresh && cached && Date.now() - cached.fetchedAt < ANTIGRAVITY_CREDIT_PROBE_TTL_MS) { + return cached.data; + } + + const inflight = _antigravityCreditProbeInflight.get(cacheKey); + if (inflight) return inflight; + + const promise = probeAntigravityCreditBalanceUncached(accessToken, accountId, projectId).then( + (data) => { + _antigravityCreditProbeCache.set(cacheKey, { data, fetchedAt: Date.now() }); + return data; + }, + (error) => { + _antigravityCreditProbeCache.set(cacheKey, { data: null, fetchedAt: Date.now() }); + throw error; + } + ).finally(() => { + _antigravityCreditProbeInflight.delete(cacheKey); + }); + + _antigravityCreditProbeInflight.set(cacheKey, promise); + return promise; +} + +async function probeAntigravityCreditBalanceUncached( accessToken: string, accountId: string, projectId?: string | null @@ -1371,8 +1480,13 @@ async function getAntigravityUsage( accessToken, providerSpecificData, connectionProjectId?, - connectionId? + connectionId?, + options: AntigravityUsageOptions = {} ) { + if (!accessToken) { + return { plan: "Free", message: "Antigravity access token not available." }; + } + try { const subscriptionInfo = await getAntigravitySubscriptionInfoCached(accessToken); const projectId = connectionProjectId || subscriptionInfo?.cloudaicompanionProject || null; @@ -1386,68 +1500,18 @@ async function getAntigravityUsage( // If no cached balance and credits mode is enabled, fire a minimal probe const creditsMode = getCreditsMode(); - if (creditBalance === null && creditsMode !== "off") { - creditBalance = await probeAntigravityCreditBalance(accessToken, accountId, projectId); + if ((options.forceRefresh || creditBalance === null) && creditsMode !== "off") { + creditBalance = await probeAntigravityCreditBalance(accessToken, accountId, projectId, options); } - // Fetch model list with quota info from fetchAvailableModels - let response: Response | null = null; - let lastError: Error | null = null; - - for (const quotaApiUrl of ANTIGRAVITY_CONFIG.quotaApiUrls) { - try { - response = await fetch(quotaApiUrl, { - method: "POST", - headers: getAntigravityHeaders("fetchAvailableModels", accessToken), - body: JSON.stringify(projectId ? { project: projectId } : {}), - signal: AbortSignal.timeout(10000), - }); - - if (response.ok || response.status === 401 || response.status === 403) { - break; - } - } catch (error) { - lastError = error as Error; - } - } - - if (!response) { - throw lastError || new Error("Antigravity API unavailable"); - } - - if (response.status === 403) { + const data = await fetchAntigravityAvailableModelsCached(accessToken, projectId, options); + const dataObj = toRecord(data); + if (dataObj.__antigravityForbidden === true) { return { message: "Antigravity access forbidden. Check subscription." }; } - - if (!response.ok) { - throw new Error(`Antigravity API error: ${response.status}`); - } - - const data = await response.json(); - const dataObj = toRecord(data); const modelEntries = toRecord(dataObj.models); const quotas: Record = {}; - // Models excluded from quota display — internal/special-purpose models that - // the Antigravity API returns quota for but are not user-callable via - // generateContent. Matches CLIProxyAPI's hardcoded exclusion list. - const ANTIGRAVITY_EXCLUDED_MODELS = new Set([ - "chat_20706", - "chat_23310", - "tab_flash_lite_preview", - "tab_jump_flash_lite_preview", - "gemini-2.5-flash-thinking", - "gemini-2.5-pro", // browser subagent model — not user-callable - "gemini-2.5-flash", // internal — quota always exhausted on free tier - "gemini-2.5-flash-lite", // internal — quota always exhausted on free tier - "gemini-2.5-flash-preview-image-generation", // image-gen only, not usable for chat - "gemini-3.1-flash-image-preview", // image-gen preview, not usable for chat - "gemini-3-flash-agent", // internal agent model — not user-callable - "gemini-3.1-flash-lite", // not usable for chat - "gemini-3-pro-low", // not usable for chat - "gemini-3-pro-high", // not usable for chat - ]); - // Parse per-model quota info from fetchAvailableModels response. for (const [modelKey, infoValue] of Object.entries(modelEntries)) { const info = toRecord(infoValue); @@ -1456,7 +1520,7 @@ async function getAntigravityUsage( // Skip internal, excluded, and models without quota info if ( info.isInternal === true || - ANTIGRAVITY_EXCLUDED_MODELS.has(modelKey) || + !isUserCallableAntigravityModelId(modelKey) || Object.keys(quotaInfo).length === 0 ) { continue; diff --git a/open-sse/utils/streamReadiness.ts b/open-sse/utils/streamReadiness.ts index d73a883b08..3104f09e99 100644 --- a/open-sse/utils/streamReadiness.ts +++ b/open-sse/utils/streamReadiness.ts @@ -50,6 +50,9 @@ function hasUsefulValue(value: unknown): boolean { "function_call_output", "output", "content_block", + "response", + "choices", + "candidates", "parts", ]) { if (hasUsefulValue(value[key])) return true; @@ -60,21 +63,7 @@ function hasUsefulValue(value: unknown): boolean { function hasUsefulJsonPayload(payload: unknown): boolean { if (!isRecord(payload)) return false; - - const type = typeof payload.type === "string" ? payload.type : ""; - if ( - type.includes("delta") || - type.includes("tool") || - type.includes("function") || - type.includes("content_block") || - type.includes("output_item") - ) { - if (hasUsefulValue(payload)) return true; - } - - return ( - hasUsefulValue(payload.choices) || hasUsefulValue(payload.candidates) || hasUsefulValue(payload) - ); + return hasUsefulValue(payload); } export function hasUsefulStreamContent(text: string): boolean { diff --git a/src/app/api/providers/[id]/models/route.ts b/src/app/api/providers/[id]/models/route.ts index 5aab30e3b2..124c32ee4b 100755 --- a/src/app/api/providers/[id]/models/route.ts +++ b/src/app/api/providers/[id]/models/route.ts @@ -48,6 +48,7 @@ import { import { ANTIGRAVITY_PUBLIC_MODELS, getClientVisibleAntigravityModelName, + isUserCallableAntigravityModelId, toClientAntigravityModelId, } from "@omniroute/open-sse/config/antigravityModelAliases.ts"; import { getEmbeddingProvider } from "@omniroute/open-sse/config/embeddingRegistry.ts"; @@ -64,6 +65,8 @@ import { type JsonRecord = Record; +const antigravityDiscoveryInflight = new Map>>(); + function asRecord(value: unknown): JsonRecord { return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {}; } @@ -173,6 +176,10 @@ function normalizeAntigravityModelsResponse(data: unknown): Array<{ id: string; .filter((value): value is { id: string; name: string } => Boolean(value)); } +function filterUserCallableAntigravityModels(models: Array<{ id: string; name: string }>) { + return models.filter((model) => isUserCallableAntigravityModelId(model.id)); +} + function mapAntigravityModelForClient(model: { id: string; name: string }): { id: string; name: string; @@ -184,6 +191,58 @@ function mapAntigravityModelForClient(model: { id: string; name: string }): { }; } +async function fetchAntigravityDiscoveryModelsCached( + accessToken: string, + connectionId: string, + proxy: unknown +): Promise> { + const cacheKey = `${connectionId}:${accessToken.substring(0, 16)}`; + const inflight = antigravityDiscoveryInflight.get(cacheKey); + if (inflight) return inflight; + + const promise = (async () => { + await resolveAntigravityVersion(); + + for (const discoveryUrl of getAntigravityModelsDiscoveryUrls()) { + try { + const response = await safeOutboundFetch(discoveryUrl, { + ...SAFE_OUTBOUND_FETCH_PRESETS.modelsDiscovery, + guard: getProviderOutboundGuard(), + proxyConfig: proxy, + method: "POST", + headers: getAntigravityHeaders("models", accessToken), + body: JSON.stringify({}), + }); + + if (!response.ok) { + const errorText = await response.text(); + console.warn( + `[models] antigravity discovery failed at ${discoveryUrl} (${response.status}): ${errorText}` + ); + continue; + } + + const models = filterUserCallableAntigravityModels( + normalizeAntigravityModelsResponse(await response.json()) + ).map(mapAntigravityModelForClient); + if (models.length > 0) { + return models; + } + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + console.warn(`[models] antigravity discovery threw for ${discoveryUrl}: ${message}`); + } + } + + return []; + })().finally(() => { + antigravityDiscoveryInflight.delete(cacheKey); + }); + + antigravityDiscoveryInflight.set(cacheKey, promise); + return promise; +} + function normalizeDataRobotCatalogResponse(data: unknown): Array<{ id: string; name: string }> { const items = Array.isArray(asRecord(data).data) ? (asRecord(data).data as unknown[]) : []; @@ -1561,7 +1620,6 @@ export async function GET( if (autoFetchDisabledResponse) return autoFetchDisabledResponse; const staticModels = STATIC_MODEL_PROVIDERS.antigravity(); - const discoveryUrls = getAntigravityModelsDiscoveryUrls(); if (!accessToken) { const fallback = buildDiscoveryFallbackResponse({ @@ -1578,37 +1636,9 @@ export async function GET( }); } - await resolveAntigravityVersion(); - - for (const discoveryUrl of discoveryUrls) { - try { - const response = await safeOutboundFetch(discoveryUrl, { - ...SAFE_OUTBOUND_FETCH_PRESETS.modelsDiscovery, - guard: getProviderOutboundGuard(), - proxyConfig: proxy, - method: "POST", - headers: getAntigravityHeaders("models", accessToken), - body: JSON.stringify({}), - }); - - if (!response.ok) { - const errorText = await response.text(); - console.warn( - `[models] antigravity discovery failed at ${discoveryUrl} (${response.status}): ${errorText}` - ); - continue; - } - - const remoteModels = normalizeAntigravityModelsResponse(await response.json()).map( - mapAntigravityModelForClient - ); - if (remoteModels.length > 0) { - return buildApiDiscoveryResponse(remoteModels); - } - } catch (error) { - const message = error instanceof Error ? error.message : String(error); - console.warn(`[models] antigravity discovery threw for ${discoveryUrl}: ${message}`); - } + const remoteModels = await fetchAntigravityDiscoveryModelsCached(accessToken, connectionId, proxy); + if (remoteModels.length > 0) { + return buildApiDiscoveryResponse(remoteModels); } const fallback = buildDiscoveryFallbackResponse(); diff --git a/src/lib/usage/providerLimits.ts b/src/lib/usage/providerLimits.ts index 5ca5b3523e..d64b851d6d 100644 --- a/src/lib/usage/providerLimits.ts +++ b/src/lib/usage/providerLimits.ts @@ -239,6 +239,16 @@ export function getCachedProviderLimitsMap(): Record { + return fetchLiveProviderLimitsWithOptions(connectionId, { forceRefresh: false }); +} + +async function fetchLiveProviderLimitsWithOptions( + connectionId: string, + options: { forceRefresh?: boolean } = {} +): Promise<{ + connection: ProviderConnectionLike; + usage: JsonRecord; }> { let connection = (await getProviderConnectionById(connectionId)) as ProviderConnectionLike | null; if (!connection) { @@ -250,7 +260,7 @@ export async function fetchLiveProviderLimits(connectionId: string): Promise<{ } if (connection.authType !== "oauth") { - const usage = (await getUsageForProvider(connection)) as JsonRecord; + const usage = (await getUsageForProvider(connection, options)) as JsonRecord; if (isRecord(usage.quotas)) { setQuotaCache(connectionId, connection.provider, usage.quotas); } @@ -274,7 +284,7 @@ export async function fetchLiveProviderLimits(connectionId: string): Promise<{ await syncToCloudIfEnabled(); } - const usageData = (await getUsageForProvider(conn)) as JsonRecord; + const usageData = (await getUsageForProvider(conn, options)) as JsonRecord; connection = conn; return { usage: usageData }; }); @@ -330,7 +340,9 @@ export async function fetchAndPersistProviderLimits( usage: JsonRecord; cache: ProviderLimitsCacheEntry; }> { - const { connection, usage } = await fetchLiveProviderLimits(connectionId); + const { connection, usage } = await fetchLiveProviderLimitsWithOptions(connectionId, { + forceRefresh: source === "manual", + }); const cache = toProviderLimitsCacheEntry(usage, source); setProviderLimitsCache(connectionId, cache); return { connection, usage, cache }; @@ -360,7 +372,9 @@ export async function syncAllProviderLimits( const chunk = connections.slice(i, i + concurrency); const results = await Promise.allSettled( chunk.map(async (connection) => { - const { usage } = await fetchLiveProviderLimits(connection.id); + const { usage } = await fetchLiveProviderLimitsWithOptions(connection.id, { + forceRefresh: source === "manual", + }); const cache = toProviderLimitsCacheEntry(usage, source); return { connectionId: connection.id, cache }; }) diff --git a/tests/unit/antigravity-model-aliases.test.ts b/tests/unit/antigravity-model-aliases.test.ts index 543f9d9c2c..313fa4b620 100644 --- a/tests/unit/antigravity-model-aliases.test.ts +++ b/tests/unit/antigravity-model-aliases.test.ts @@ -2,6 +2,7 @@ import test from "node:test"; import assert from "node:assert/strict"; import { + isUserCallableAntigravityModelId, resolveAntigravityModelId, toClientAntigravityModelId, } from "../../open-sse/config/antigravityModelAliases.ts"; @@ -32,6 +33,15 @@ test("toClientAntigravityModelId exposes client-visible aliases for known upstre assert.equal(toClientAntigravityModelId("claude-opus-4-6-thinking"), "claude-opus-4-6-thinking"); }); +test("isUserCallableAntigravityModelId only allows public chat-capable model IDs", () => { + assert.equal(isUserCallableAntigravityModelId("gemini-3-pro-preview"), true); + assert.equal(isUserCallableAntigravityModelId("gemini-3.1-pro-high"), true); + assert.equal(isUserCallableAntigravityModelId("claude-sonnet-4-6"), true); + assert.equal(isUserCallableAntigravityModelId("gemini-3-flash-agent"), false); + assert.equal(isUserCallableAntigravityModelId("tab_flash_lite_preview"), false); + assert.equal(isUserCallableAntigravityModelId("unknown-model"), false); +}); + test("AntigravityExecutor.transformRequest resolves alias models before dispatching upstream", async () => { const executor = new AntigravityExecutor(); const result = await executor.transformRequest( diff --git a/tests/unit/stream-readiness.test.ts b/tests/unit/stream-readiness.test.ts index 1917dd3827..2e429921d9 100644 --- a/tests/unit/stream-readiness.test.ts +++ b/tests/unit/stream-readiness.test.ts @@ -66,6 +66,18 @@ test("hasUsefulStreamContent detects text, reasoning, and tool deltas", () => { ), true ); + assert.equal( + hasUsefulStreamContent( + `data: ${JSON.stringify({ type: "response.output_text.delta", delta: "hello" })}\n\n` + ), + true + ); + assert.equal( + hasUsefulStreamContent( + `data: ${JSON.stringify({ response: { candidates: [{ content: { parts: [{ text: "hello" }] } }] } })}\n\n` + ), + true + ); }); test("ensureStreamReadiness preserves buffered chunks when stream starts", async () => { diff --git a/tests/unit/usage-fetcher-antigravity.test.ts b/tests/unit/usage-fetcher-antigravity.test.ts deleted file mode 100644 index 5bb58104a1..0000000000 --- a/tests/unit/usage-fetcher-antigravity.test.ts +++ /dev/null @@ -1,59 +0,0 @@ -import test from "node:test"; -import assert from "node:assert/strict"; - -const usageFetcher = await import("../../src/lib/usage/fetcher.ts"); - -const originalFetch = globalThis.fetch; - -test.afterEach(() => { - globalThis.fetch = originalFetch; -}); - -test("usage fetcher retries Antigravity quota discovery across shared fallback URLs", async () => { - const calls = []; - - globalThis.fetch = async (url, init = {}) => { - calls.push({ url: String(url), init }); - - const urlStr = String(url); - if ( - urlStr === "https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels" || - urlStr === "https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels" - ) { - return new Response("unavailable", { status: 503 }); - } - - return new Response( - JSON.stringify({ - models: { - "claude-sonnet-4-6": { - quotaInfo: { - remainingFraction: 0.4, - resetTime: new Date(Date.now() + 60_000).toISOString(), - }, - }, - }, - }), - { status: 200, headers: { "Content-Type": "application/json" } } - ); - }; - - const usage: any = await usageFetcher.getUsageForProvider({ - provider: "antigravity", - accessToken: "ag-token", - providerSpecificData: { email: "coder@example.com" }, - }); - - assert.deepEqual( - calls.map((call) => call.url), - [ - "https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels", - "https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels", - "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:fetchAvailableModels", - ] - ); - assert.match(calls[1].init.headers["User-Agent"], /^antigravity\//); - assert.equal(usage.plan, "Antigravity"); - assert.equal(usage.quotas.models.total, 1); - assert.equal(usage.modelQuotas["claude-sonnet-4-6"].remaining, 40); -}); diff --git a/tests/unit/usage-service-hardening.test.ts b/tests/unit/usage-service-hardening.test.ts index 82b185d8b3..bf17bb6546 100644 --- a/tests/unit/usage-service-hardening.test.ts +++ b/tests/unit/usage-service-hardening.test.ts @@ -5,9 +5,15 @@ const usageService = await import("../../open-sse/services/usage.ts"); const { __testing } = usageService; const originalFetch = globalThis.fetch; +const originalCreditsMode = process.env.ANTIGRAVITY_CREDITS; test.afterEach(() => { globalThis.fetch = originalFetch; + if (originalCreditsMode === undefined) { + delete process.env.ANTIGRAVITY_CREDITS; + } else { + process.env.ANTIGRAVITY_CREDITS = originalCreditsMode; + } }); test("usage service covers GitHub free-plan parsing, auth denial and unsupported providers", async () => { @@ -290,7 +296,7 @@ test("usage service covers Antigravity quota parsing, exclusions and forbidden a "gemini-unlimited": { quotaInfo: {}, }, - "gemini-open": { + "gemini-3.1-pro-high": { quotaInfo: { remainingFraction: 1 }, }, "internal-model": { @@ -312,10 +318,10 @@ test("usage service covers Antigravity quota parsing, exclusions and forbidden a }); assert.equal(usage.plan, "Ultra"); - assert.deepEqual(Object.keys(usage.quotas).sort(), ["claude-sonnet-4-6", "gemini-open"]); + assert.deepEqual(Object.keys(usage.quotas).sort(), ["claude-sonnet-4-6", "gemini-3.1-pro-high"]); assert.equal(usage.quotas["claude-sonnet-4-6"].used, 600); - assert.equal(usage.quotas["gemini-open"].total, 0); - assert.equal(usage.quotas["gemini-open"].remainingPercentage, 100); + assert.equal(usage.quotas["gemini-3.1-pro-high"].total, 0); + assert.equal(usage.quotas["gemini-3.1-pro-high"].remainingPercentage, 100); const loadCodeAssistCall = calls.find((call) => call.url.includes("loadCodeAssist")); assert.equal(loadCodeAssistCall?.init.headers["User-Agent"], "google-api-nodejs-client/10.3.0"); assert.equal( @@ -407,6 +413,74 @@ test("usage service retries Antigravity fetchAvailableModels across the shared f assert.equal(usage.quotas["claude-sonnet-4-6"].used, 500); }); +test("usage service manual Antigravity refresh bypasses usage TTL caches", async () => { + process.env.ANTIGRAVITY_CREDITS = "retry"; + let probeCalls = 0; + let modelCalls = 0; + + globalThis.fetch = async (url) => { + const urlStr = String(url); + if (urlStr.includes("loadCodeAssist")) { + return new Response(JSON.stringify({ cloudaicompanionProject: "ag-project" }), { status: 200 }); + } + + if (urlStr.includes("streamGenerateContent")) { + probeCalls++; + return new Response( + `data: ${JSON.stringify({ remainingCredits: [{ creditType: "GOOGLE_ONE_AI", creditAmount: String(100 - probeCalls) }] })}\n\n`, + { status: 200, headers: { "Content-Type": "text/event-stream" } } + ); + } + + if (urlStr.includes("fetchAvailableModels")) { + modelCalls++; + return new Response( + JSON.stringify({ + models: { + "claude-sonnet-4-6": { + quotaInfo: { remainingFraction: 1 }, + }, + }, + }), + { status: 200, headers: { "Content-Type": "application/json" } } + ); + } + + throw new Error(`unexpected fetch: ${url}`); + }; + + const connection = { + id: "ag-manual-refresh-service-test", + provider: "antigravity", + accessToken: "ag-manual-service-token", + projectId: "ag-project", + }; + + await usageService.getUsageForProvider(connection, { forceRefresh: true }); + await usageService.getUsageForProvider(connection, { forceRefresh: true }); + + assert.equal(probeCalls, 2); + assert.equal(modelCalls, 2); +}); + +test("usage service handles missing Antigravity access tokens without probing upstream", async () => { + let fetchCalls = 0; + + globalThis.fetch = async () => { + fetchCalls++; + return new Response("unexpected", { status: 500 }); + }; + + const usage: any = await usageService.getUsageForProvider({ + provider: "antigravity", + accessToken: undefined, + }); + + assert.equal(fetchCalls, 0); + assert.equal(usage.plan, "Free"); + assert.match(usage.message, /Antigravity access token not available/i); +}); + test("usage service covers Antigravity tier fallbacks and non-403 upstream failures", async () => { globalThis.fetch = async (url) => { if (String(url).includes("loadCodeAssist")) {