diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index f18ba970bc..e96eeaa123 100644 --- a/open-sse/services/usage.ts +++ b/open-sse/services/usage.ts @@ -3,14 +3,6 @@ */ import { PROVIDERS } from "../config/constants.ts"; - -// Quota / usage upstream URLs (overridable for testing or relays). -const CROF_USAGE_URL = process.env.OMNIROUTE_CROF_USAGE_URL ?? "https://crof.ai/usage_api/"; -const GEMINI_CLI_USAGE_URL = - process.env.OMNIROUTE_GEMINI_CLI_USAGE_URL ?? - "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"; -const CODEWHISPERER_BASE_URL = - process.env.OMNIROUTE_CODEWHISPERER_BASE_URL ?? "https://codewhisperer.us-east-1.amazonaws.com"; import { getAntigravityFetchAvailableModelsUrls, ANTIGRAVITY_BASE_URLS, @@ -43,6 +35,14 @@ import { extractCodeAssistSubscriptionTier, } from "./codeAssistSubscription.ts"; +// Quota / usage upstream URLs (overridable for testing or relays). +const CROF_USAGE_URL = process.env.OMNIROUTE_CROF_USAGE_URL ?? "https://crof.ai/usage_api/"; +const GEMINI_CLI_USAGE_URL = + process.env.OMNIROUTE_GEMINI_CLI_USAGE_URL ?? + "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"; +const CODEWHISPERER_BASE_URL = + process.env.OMNIROUTE_CODEWHISPERER_BASE_URL ?? "https://codewhisperer.us-east-1.amazonaws.com"; + // Antigravity API config (credentials from PROVIDERS via credential loader) const ANTIGRAVITY_CONFIG = { quotaApiUrls: getAntigravityFetchAvailableModelsUrls(), @@ -1940,8 +1940,9 @@ async function getAntigravityUsage( const rawFraction = toNumber(quotaInfo.remainingFraction, -1); const resetAt = parseResetTime(quotaInfo.resetTime); - // Default to 100% when the API doesn't report a fraction - const remainingFraction = rawFraction < 0 ? 1 : rawFraction; + // When remainingFraction is undefined/NaN (exhausted quota), default to 0%. Clamp to valid range [0, 1] + // Unlimited models have remainingFraction=1 AND no resetTime + const remainingFraction = rawFraction < 0 ? 0 : Math.max(0, Math.min(1, rawFraction)); // Models with no resetTime and full remaining are unlimited (e.g. tab-completion models) const isUnlimited = !resetAt && remainingFraction >= 1; const remainingPercentage = remainingFraction * 100; diff --git a/src/lib/usage/fetcher.ts b/src/lib/usage/fetcher.ts index 58f4d3d086..092de749d7 100644 --- a/src/lib/usage/fetcher.ts +++ b/src/lib/usage/fetcher.ts @@ -348,9 +348,13 @@ async function getAntigravityUsage( if (info.isInternal) continue; const quotaInfo = (info.quotaInfo as Record) ?? {}; - if ("remainingFraction" in quotaInfo) { + // Process models with any quota metadata (exhausted models may have only resetTime, active models have remainingFraction, credit-based models have empty quotaInfo) + if (Object.keys(quotaInfo).length > 0) { + // Default to 0 when remainingFraction is undefined/null/invalid, clamp to valid range [0, 1] const fraction = - typeof quotaInfo.remainingFraction === "number" ? quotaInfo.remainingFraction : 1; + typeof quotaInfo.remainingFraction === "number" + ? Math.max(0, Math.min(1, quotaInfo.remainingFraction)) + : 0; const resetTime = typeof quotaInfo.resetTime === "string" ? quotaInfo.resetTime : null; modelQuotas[modelId] = { remaining: Math.round(fraction * 100), @@ -360,7 +364,7 @@ async function getAntigravityUsage( quotaModelsTotal++; if (fraction > 0) quotaModelsAvailable++; } - // Credit-based models have no remainingFraction — their availability is + // Credit-based models have empty quotaInfo — their availability is // tracked via the GOOGLE_ONE_AI credit balance cached from SSE responses. } diff --git a/tests/unit/antigravity-usage-fetcher.test.ts b/tests/unit/antigravity-usage-fetcher.test.ts new file mode 100644 index 0000000000..b53e482d53 --- /dev/null +++ b/tests/unit/antigravity-usage-fetcher.test.ts @@ -0,0 +1,285 @@ +/** + * Tests for src/lib/usage/fetcher.ts — Antigravity quota parsing. + * + * Verifies that remainingFraction is correctly parsed: + * - undefined → 0% remaining (exhausted quota) + * - 0 → 0% remaining (exhausted quota, explicit) + * - 1.0 → 100% remaining (full quota) + * - 0.5 → 50% remaining (partial quota) + */ + +import { describe, it, mock } from "node:test"; +import assert from "node:assert/strict"; + +describe("getUsageForProvider (antigravity in fetcher.ts)", () => { + const connectionBase = { + provider: "antigravity", + accessToken: "fake-token", + providerSpecificData: {}, + projectId: undefined, + id: "test-conn", + }; + + it("defaults to 0% remaining when remainingFraction is undefined", async () => { + const fetcherModule = await import("../../src/lib/usage/fetcher.ts"); + const { getUsageForProvider } = fetcherModule; + + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-2.5-pro": { + quotaInfo: { + remainingFraction: undefined, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const result = await getUsageForProvider(connectionBase); + assert.ok(result, "should return a result"); + + if ("modelQuotas" in result) { + const quota = result.modelQuotas["gemini-2.5-pro"]; + assert.ok(quota, "should have quota for gemini-2.5-pro"); + assert.equal(quota.remaining, 0, "remaining should be 0%"); + assert.equal(quota.limited, true, "should be marked as limited"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("defaults to 0% remaining when remainingFraction key is absent (exhausted quota)", async () => { + const fetcherModule = await import("../../src/lib/usage/fetcher.ts"); + const { getUsageForProvider } = fetcherModule; + + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-2.5-pro": { + quotaInfo: { + // remainingFraction key is omitted (exhausted quota) + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const result = await getUsageForProvider(connectionBase); + assert.ok(result, "should return a result"); + + if ("modelQuotas" in result) { + const quota = result.modelQuotas["gemini-2.5-pro"]; + assert.ok(quota, "should have quota for gemini-2.5-pro"); + assert.equal(quota.remaining, 0, "remaining should be 0% when key is absent"); + assert.equal(quota.limited, true, "should be marked as limited"); + assert.equal(quota.resetAt, "2026-05-26T00:00:00Z", "should preserve resetTime"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("parses remainingFraction=0 as exhausted quota", async () => { + const fetcherModule = await import("../../src/lib/usage/fetcher.ts"); + const { getUsageForProvider } = fetcherModule; + + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-2.5-pro": { + quotaInfo: { + remainingFraction: 0, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const result = await getUsageForProvider(connectionBase); + assert.ok(result, "should return a result"); + + if ("modelQuotas" in result) { + const quota = result.modelQuotas["gemini-2.5-pro"]; + assert.ok(quota, "should have quota for gemini-2.5-pro"); + assert.equal(quota.remaining, 0, "remaining should be 0%"); + assert.equal(quota.limited, true, "should be marked as limited"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("parses remainingFraction=1.0 with resetTime as full quota", async () => { + const fetcherModule = await import("../../src/lib/usage/fetcher.ts"); + const { getUsageForProvider } = fetcherModule; + + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-2.5-pro": { + quotaInfo: { + remainingFraction: 1.0, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const result = await getUsageForProvider(connectionBase); + assert.ok(result, "should return a result"); + + if ("modelQuotas" in result) { + const quota = result.modelQuotas["gemini-2.5-pro"]; + assert.ok(quota, "should have quota for gemini-2.5-pro"); + assert.equal(quota.remaining, 100, "remaining should be 100%"); + assert.equal(quota.limited, false, "should not be marked as limited"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("parses remainingFraction=0.5 as partial quota", async () => { + const fetcherModule = await import("../../src/lib/usage/fetcher.ts"); + const { getUsageForProvider } = fetcherModule; + + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-2.5-pro": { + quotaInfo: { + remainingFraction: 0.5, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const result = await getUsageForProvider(connectionBase); + assert.ok(result, "should return a result"); + + if ("modelQuotas" in result) { + const quota = result.modelQuotas["gemini-2.5-pro"]; + assert.ok(quota, "should have quota for gemini-2.5-pro"); + assert.equal(quota.remaining, 50, "remaining should be 50%"); + assert.equal(quota.limited, false, "should not be marked as limited"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("skips credit-based models without remainingFraction", async () => { + const fetcherModule = await import("../../src/lib/usage/fetcher.ts"); + const { getUsageForProvider } = fetcherModule; + + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-2.5-pro": { + quotaInfo: {}, + }, + }, + }), + })); + + try { + const result = await getUsageForProvider(connectionBase); + assert.ok(result, "should return a result"); + + if ("modelQuotas" in result) { + assert.equal( + Object.keys(result.modelQuotas).length, + 0, + "should not include credit-based models" + ); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("clamps remainingFraction > 1 to 100%", async () => { + const fetcherModule = await import("../../src/lib/usage/fetcher.ts"); + const { getUsageForProvider } = fetcherModule; + + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-2.5-pro": { + quotaInfo: { + remainingFraction: 1.5, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const result = await getUsageForProvider(connectionBase); + assert.ok(result, "should return a result"); + + if ("modelQuotas" in result) { + const quota = result.modelQuotas["gemini-2.5-pro"]; + assert.ok(quota, "should have quota for gemini-2.5-pro"); + assert.equal(quota.remaining, 100, "remaining should be clamped to 100%"); + assert.equal(quota.limited, false, "should not be marked as limited"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("clamps negative remainingFraction to 0%", async () => { + const fetcherModule = await import("../../src/lib/usage/fetcher.ts"); + const { getUsageForProvider } = fetcherModule; + + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-2.5-pro": { + quotaInfo: { + remainingFraction: -0.5, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const result = await getUsageForProvider(connectionBase); + assert.ok(result, "should return a result"); + + if ("modelQuotas" in result) { + const quota = result.modelQuotas["gemini-2.5-pro"]; + assert.ok(quota, "should have quota for gemini-2.5-pro"); + assert.equal(quota.remaining, 0, "remaining should be clamped to 0%"); + assert.equal(quota.limited, true, "should be marked as limited"); + } + } finally { + mockFetch.mock.restore(); + } + }); +}); diff --git a/tests/unit/antigravity-usage-service.test.ts b/tests/unit/antigravity-usage-service.test.ts new file mode 100644 index 0000000000..9a357f6a01 --- /dev/null +++ b/tests/unit/antigravity-usage-service.test.ts @@ -0,0 +1,261 @@ +/** + * Tests for open-sse/services/usage.ts — Antigravity quota parsing. + * + * Verifies that remainingFraction is correctly parsed: + * - undefined → 0% remaining (exhausted quota) + * - 0 → 0% remaining (exhausted quota, explicit) + * - 1.0 → 100% remaining (full quota) + * - 1.0 without resetTime → unlimited (e.g. tab-completion) + * - 0.5 → 50% remaining (partial quota) + */ + +import { describe, it, mock } from "node:test"; +import assert from "node:assert/strict"; + +describe("getUsageForProvider (antigravity in usage.ts)", () => { + const connectionBase = { + id: "test-conn", + provider: "antigravity", + accessToken: "fake-token", + providerSpecificData: {}, + projectId: undefined, + }; + + it("defaults to 0% remaining when remainingFraction is undefined", async () => { + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-3.5-flash-preview": { + quotaInfo: { + remainingFraction: undefined, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const usageModule = await import("../../open-sse/services/usage.ts"); + const { getUsageForProvider } = usageModule; + + const result = await getUsageForProvider(connectionBase, { forceRefresh: true }); + assert.ok(result, "should return a result"); + assert.ok("quotas" in result, "should have quotas"); + + if ("quotas" in result) { + const quota = result.quotas["gemini-3.5-flash-preview"]; + assert.ok(quota, "should have quota for gemini-3.5-flash-preview"); + assert.equal(quota.remainingPercentage, 0, "remaining should be 0%"); + assert.equal(quota.unlimited, false, "should not be unlimited"); + assert.equal(quota.used > 0, true, "used should be > 0 when quota is exhausted"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("parses remainingFraction=0 as exhausted quota", async () => { + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-3.5-flash-preview": { + quotaInfo: { + remainingFraction: 0, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const usageModule = await import("../../open-sse/services/usage.ts"); + const { getUsageForProvider } = usageModule; + + const result = await getUsageForProvider(connectionBase, { forceRefresh: true }); + assert.ok(result, "should return a result"); + assert.ok("quotas" in result, "should have quotas"); + + if ("quotas" in result) { + const quota = result.quotas["gemini-3.5-flash-preview"]; + assert.ok(quota, "should have quota for gemini-3.5-flash-preview"); + assert.equal(quota.remainingPercentage, 0, "remaining should be 0%"); + assert.equal(quota.unlimited, false, "should not be unlimited"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("parses remainingFraction=1.0 with resetTime as full quota", async () => { + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-3.5-flash-preview": { + quotaInfo: { + remainingFraction: 1.0, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const usageModule = await import("../../open-sse/services/usage.ts"); + const { getUsageForProvider } = usageModule; + + const result = await getUsageForProvider(connectionBase, { forceRefresh: true }); + assert.ok(result, "should return a result"); + assert.ok("quotas" in result, "should have quotas"); + + if ("quotas" in result) { + const quota = result.quotas["gemini-3.5-flash-preview"]; + assert.ok(quota, "should have quota for gemini-3.5-flash-preview"); + assert.equal(quota.remainingPercentage, 100, "remaining should be 100%"); + assert.equal(quota.unlimited, false, "should not be unlimited (has resetTime)"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("parses remainingFraction=1.0 without resetTime as unlimited (e.g. tab-completion)", async () => { + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-3.1-flash-lite": { + quotaInfo: { + remainingFraction: 1.0, + }, + }, + }, + }), + })); + + try { + const usageModule = await import("../../open-sse/services/usage.ts"); + const { getUsageForProvider } = usageModule; + + const result = await getUsageForProvider(connectionBase, { forceRefresh: true }); + assert.ok(result, "should return a result"); + assert.ok("quotas" in result, "should have quotas"); + + if ("quotas" in result) { + const quota = result.quotas["gemini-3.1-flash-lite"]; + assert.ok(quota, "should have quota for gemini-3.1-flash-lite"); + assert.equal(quota.remainingPercentage, 100, "remaining should be 100%"); + assert.equal(quota.unlimited, true, "should be unlimited (no resetTime)"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("parses remainingFraction=0.5 as partial quota", async () => { + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-3.5-flash-preview": { + quotaInfo: { + remainingFraction: 0.5, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const usageModule = await import("../../open-sse/services/usage.ts"); + const { getUsageForProvider } = usageModule; + + const result = await getUsageForProvider(connectionBase, { forceRefresh: true }); + assert.ok(result, "should return a result"); + assert.ok("quotas" in result, "should have quotas"); + + if ("quotas" in result) { + const quota = result.quotas["gemini-3.5-flash-preview"]; + assert.ok(quota, "should have quota for gemini-3.5-flash-preview"); + assert.equal(quota.remainingPercentage, 50, "remaining should be 50%"); + assert.equal(quota.unlimited, false, "should not be unlimited"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("clamps remainingFraction > 1 to 100%", async () => { + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-3.5-flash-preview": { + quotaInfo: { + remainingFraction: 1.5, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const usageModule = await import("../../open-sse/services/usage.ts"); + const { getUsageForProvider } = usageModule; + + const result = await getUsageForProvider(connectionBase, { forceRefresh: true }); + assert.ok(result, "should return a result"); + assert.ok("quotas" in result, "should have quotas"); + + if ("quotas" in result) { + const quota = result.quotas["gemini-3.5-flash-preview"]; + assert.ok(quota, "should have quota for gemini-3.5-flash-preview"); + assert.equal(quota.remainingPercentage, 100, "remaining should be clamped to 100%"); + assert.equal(quota.unlimited, false, "should not be unlimited (has resetTime)"); + } + } finally { + mockFetch.mock.restore(); + } + }); + + it("clamps negative remainingFraction to 0%", async () => { + const mockFetch = mock.method(global, "fetch", async () => ({ + ok: true, + json: async () => ({ + models: { + "gemini-3.5-flash-preview": { + quotaInfo: { + remainingFraction: -0.5, + resetTime: "2026-05-26T00:00:00Z", + }, + }, + }, + }), + })); + + try { + const usageModule = await import("../../open-sse/services/usage.ts"); + const { getUsageForProvider } = usageModule; + + const result = await getUsageForProvider(connectionBase, { forceRefresh: true }); + assert.ok(result, "should return a result"); + assert.ok("quotas" in result, "should have quotas"); + + if ("quotas" in result) { + const quota = result.quotas["gemini-3.5-flash-preview"]; + assert.ok(quota, "should have quota for gemini-3.5-flash-preview"); + assert.equal(quota.remainingPercentage, 0, "remaining should be clamped to 0%"); + assert.equal(quota.unlimited, false, "should not be unlimited"); + } + } finally { + mockFetch.mock.restore(); + } + }); +});