pjandro/OmniRoute
1
0
Fork 0
mirror of https://github.com/diegosouzapw/OmniRoute.git synced 2026-07-26 09:52:11 +03:00
Code Issues Packages Projects Releases Wiki Activity

fix(usage): normalize Antigravity and agy provider quotas (#3604)

Browse Source 
Split-out PR B from #3584. Normalizes Antigravity/agy provider quotas: prefers retrieveUserQuota for live consumption, falls back to fetchAvailableModels and local usage_history, sanitizes cached Provider Limits so retired upstream IDs are not re-exposed, and schedules a deduplicated post-usage refresh. Maintainer follow-up: decoupled the post-usage refresh via a lightweight usageEvents bus (usageHistory no longer dynamic-imports providerLimits) so it does not pull the executors/translator graph into the typecheck-core surface — typecheck:core stays at 0. Integrated into release/v3.8.21. Thanks @dhaern!
This commit is contained in:
Raxxoor
2026-06-11 04:32:51 +01:00
committed by GitHub
parent 5c11d1c92e
commit 68b71bc826
12 changed files with 853 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files

306
open-sse/services/usage.ts
View File

@@ -7,16 +7,18 @@ import {
getAntigravityFetchAvailableModelsUrls,
ANTIGRAVITY_BASE_URLS,
} from "../config/antigravityUpstream.ts";
import { isUserCallableAntigravityModelId } from "../config/antigravityModelAliases.ts";
import {
isUserCallableAntigravityModelId,
toClientAntigravityModelId,
} from "../config/antigravityModelAliases.ts";
import { isUserCallableAgyModelId } from "../config/agyModels.ts";
import { getGlmQuotaUrl } from "../config/glmProvider.ts";
import { getGitHubCopilotInternalUserHeaders } from "../config/providerHeaderProfiles.ts";
import { safePercentage } from "@/shared/utils/formatting";
import { getDbInstance } from "@/lib/db/core";
import { fetchBailianQuota, type BailianTripleWindowQuota } from "./bailianQuotaFetcher.ts";
import { fetchDeepseekQuota, type DeepseekQuota } from "./deepseekQuotaFetcher.ts";
import {
fetchOpencodeQuota,
type OpencodeTripleWindowQuota,
} from "./opencodeQuotaFetcher.ts";
import { fetchOpencodeQuota, type OpencodeTripleWindowQuota } from "./opencodeQuotaFetcher.ts";
import {
applyAntigravityClientProfileHeaders,
getAntigravityBootstrapHeaders,
@@ -139,6 +141,7 @@ type UsageQuota = {
* NOT a confirmed-exhausted state. Antigravity-specific.
*/
fractionReported?: boolean;
quotaSource?: "retrieveUserQuota" | "fetchAvailableModels" | "localUsageHistory";
displayName?: string;
details?: Array<{
name: string;
@@ -942,7 +945,9 @@ async function getOpenCodeGoUsage(apiKey: string) {
if (!res.ok) {
if (res.status === 401 || res.status === 403) {
return { message: "OpenCode Go quota endpoint rejected this API key. Chat requests still work." };
return {
message: "OpenCode Go quota endpoint rejected this API key. Chat requests still work.",
};
}
return { message: `OpenCode Go quota API error (${res.status})` };
}
@@ -956,7 +961,9 @@ async function getOpenCodeGoUsage(apiKey: string) {
const code = toNumber((json as Record<string, unknown>).code, 200);
if (code === 401 || code === 403 || (json as Record<string, unknown>).success === false) {
return { message: "OpenCode Go quota endpoint rejected this API key. Chat requests still work." };
return {
message: "OpenCode Go quota endpoint rejected this API key. Chat requests still work.",
};
}
const data = toRecord((json as Record<string, unknown>).data);
@@ -1161,7 +1168,9 @@ async function getOpencodeUsage(connectionId: string, apiKey: string) {
}
try {
const quota = (await fetchOpencodeQuota(connectionId, { apiKey })) as OpencodeTripleWindowQuota | null;
const quota = (await fetchOpencodeQuota(connectionId, {
apiKey,
})) as OpencodeTripleWindowQuota | null;
if (!quota) {
return { message: "OpenCode connected. Unable to fetch quota data." };
@@ -1490,7 +1499,14 @@ export async function getUsageForProvider(
return await getGeminiUsage(accessToken, providerSpecificData, projectId);
case "antigravity":
case "agy":
return await getAntigravityUsage(accessToken, providerSpecificData, projectId, id, options);
return await getAntigravityUsage(
provider,
accessToken,
providerSpecificData,
projectId,
id,
options
);
case "claude":
return await getClaudeUsage(accessToken);
case "codex":
@@ -1643,10 +1659,7 @@ async function getGitHubUsage(accessToken?: string, providerSpecificData?: JsonR
const addLimitedQuota = (name: string) => {
const total = toNumber(getFieldValue(monthlyQuotas, name, name), 0);
if (total <= 0) return null;
const remainingRaw = Math.max(
0,
toNumber(getFieldValue(remainingQuotas, name, name), 0)
);
const remainingRaw = Math.max(0, toNumber(getFieldValue(remainingQuotas, name, name), 0));
const remaining = Math.min(remainingRaw, total);
const used = Math.max(total - remaining, 0);
quotas[name] = {
@@ -1918,6 +1931,8 @@ const ANTIGRAVITY_MODELS_CACHE_TTL_MS = 60 * 1000;
const ANTIGRAVITY_CREDIT_PROBE_TTL_MS = 5 * 60 * 1000;
const _antigravityAvailableModelsCache = new Map<string, { data: unknown; fetchedAt: number }>();
const _antigravityAvailableModelsInflight = new Map<string, Promise<unknown>>();
const _antigravityUserQuotaCache = new Map<string, { data: unknown; fetchedAt: number }>();
const _antigravityUserQuotaInflight = new Map<string, Promise<unknown>>();
const _antigravityCreditProbeCache = new Map<string, { data: number | null; fetchedAt: number }>();
const _antigravityCreditProbeInflight = new Map<string, Promise<number | null>>();
@@ -1926,27 +1941,125 @@ const _antigravityCreditProbeInflight = new Map<string, Promise<number | null>>(
// purges stale entries so keys accessed once and never again don't leak memory.
// The 2 inflight Maps (availableModelsInflight, creditProbeInflight) self-clean
// when the Promise resolves/rejects, so they are NOT touched here.
const _usageCacheCleanupTimer = setInterval(() => {
const now = Date.now();
for (const [key, entry] of _geminiCliSubCache) {
if (now - entry.fetchedAt > GEMINI_CLI_CACHE_TTL_MS) _geminiCliSubCache.delete(key);
}
for (const [key, entry] of _antigravitySubCache) {
if (now - entry.fetchedAt > ANTIGRAVITY_CACHE_TTL_MS) _antigravitySubCache.delete(key);
}
for (const [key, entry] of _antigravityAvailableModelsCache) {
if (now - entry.fetchedAt > ANTIGRAVITY_MODELS_CACHE_TTL_MS) _antigravityAvailableModelsCache.delete(key);
}
for (const [key, entry] of _antigravityCreditProbeCache) {
if (now - entry.fetchedAt > ANTIGRAVITY_CREDIT_PROBE_TTL_MS) _antigravityCreditProbeCache.delete(key);
}
}, 5 * 60 * 1000); // every 5 minutes
const _usageCacheCleanupTimer = setInterval(
() => {
const now = Date.now();
for (const [key, entry] of _geminiCliSubCache) {
if (now - entry.fetchedAt > GEMINI_CLI_CACHE_TTL_MS) _geminiCliSubCache.delete(key);
}
for (const [key, entry] of _antigravitySubCache) {
if (now - entry.fetchedAt > ANTIGRAVITY_CACHE_TTL_MS) _antigravitySubCache.delete(key);
}
for (const [key, entry] of _antigravityAvailableModelsCache) {
if (now - entry.fetchedAt > ANTIGRAVITY_MODELS_CACHE_TTL_MS)
_antigravityAvailableModelsCache.delete(key);
}
for (const [key, entry] of _antigravityUserQuotaCache) {
if (now - entry.fetchedAt > ANTIGRAVITY_MODELS_CACHE_TTL_MS)
_antigravityUserQuotaCache.delete(key);
}
for (const [key, entry] of _antigravityCreditProbeCache) {
if (now - entry.fetchedAt > ANTIGRAVITY_CREDIT_PROBE_TTL_MS)
_antigravityCreditProbeCache.delete(key);
}
},
5 * 60 * 1000
); // every 5 minutes
_usageCacheCleanupTimer.unref?.(); // Don't prevent process exit
interface AntigravityUsageOptions {
forceRefresh?: boolean;
}
const ANTIGRAVITY_LOCAL_USAGE_WINDOW_MS = 5 * 60 * 60 * 1000;
const ANTIGRAVITY_LOCAL_USAGE_TOKENS_PER_UNIT = 1000;
const ANTIGRAVITY_QUOTA_MODEL_ALIASES: Record<string, string | null> = {
"gemini-3.5-flash-preview": null,
"gemini-3-flash-preview": null,
};
function normalizeAntigravityQuotaModelId(modelId: string): string | null {
if (!modelId) return null;
return Object.prototype.hasOwnProperty.call(ANTIGRAVITY_QUOTA_MODEL_ALIASES, modelId)
? ANTIGRAVITY_QUOTA_MODEL_ALIASES[modelId]
: modelId;
}
function toClientAntigravityQuotaModelId(modelId: string): string | null {
if (!modelId) return null;
if (normalizeAntigravityQuotaModelId(modelId) === null) return null;
if (modelId === "gemini-3.5-flash-extra-low") return "gemini-3.5-flash-low";
if (modelId === "gemini-3.5-flash-low") return "gemini-3.5-flash-medium";
if (modelId === "gemini-3-flash-agent") return "gemini-3.5-flash-high";
return toClientAntigravityModelId(modelId);
}
function getAntigravityLocalUsageUnits(
provider: "antigravity" | "agy",
connectionId: string | undefined,
modelId: string,
resetAt: string | null
): number {
if (!connectionId || !modelId || !resetAt) return 0;
const resetMs = Date.parse(resetAt);
if (!Number.isFinite(resetMs)) return 0;
const windowStart = new Date(resetMs - ANTIGRAVITY_LOCAL_USAGE_WINDOW_MS).toISOString();
const windowEnd = new Date(resetMs).toISOString();
try {
const db = getDbInstance() as unknown as {
prepare: (sql: string) => { get: (...params: unknown[]) => unknown };
};
const row = db
.prepare(
`SELECT COALESCE(SUM(
COALESCE(tokens_input, 0) + COALESCE(tokens_output, 0) + COALESCE(tokens_reasoning, 0)
), 0) AS tokens
FROM usage_history
WHERE provider = ?
AND connection_id = ?
AND model = ?
AND success = 1
AND timestamp >= ?
AND timestamp < ?`
)
.get(provider, connectionId, modelId, windowStart, windowEnd) as
| { tokens?: unknown }
| undefined;
const tokens = Number(row?.tokens || 0);
if (!Number.isFinite(tokens) || tokens <= 0) return 0;
return Math.max(1, Math.ceil(tokens / ANTIGRAVITY_LOCAL_USAGE_TOKENS_PER_UNIT));
} catch {
return 0;
}
}
function applyLocalUsageFallback(
quota: UsageQuota,
provider: "antigravity" | "agy",
connectionId: string | undefined,
modelId: string
): UsageQuota {
if (quota.quotaSource !== "fetchAvailableModels" || quota.used > 0 || quota.unlimited) {
return quota;
}
const localUsed = getAntigravityLocalUsageUnits(provider, connectionId, modelId, quota.resetAt);
if (localUsed <= 0 || quota.total <= 0) return quota;
const used = Math.min(quota.total, localUsed);
return {
...quota,
used,
remainingPercentage: Math.max(0, ((quota.total - used) / quota.total) * 100),
quotaSource: "localUsageHistory",
};
}
function buildAntigravityUsageCacheKey(accessToken: string, projectId?: string | null): string {
return `${accessToken.substring(0, 16)}:${projectId || "default"}`;
}
@@ -2015,6 +2128,57 @@ async function fetchAntigravityAvailableModelsCached(
return promise;
}
async function fetchAntigravityUserQuotaCached(
accessToken: string,
projectId?: string | null,
options: AntigravityUsageOptions = {}
): Promise<unknown | null> {
if (!accessToken || !projectId) return null;
const cacheKey = buildAntigravityUsageCacheKey(accessToken, projectId);
const cached = _antigravityUserQuotaCache.get(cacheKey);
if (
!options.forceRefresh &&
cached &&
Date.now() - cached.fetchedAt < ANTIGRAVITY_MODELS_CACHE_TTL_MS
) {
return cached.data;
}
const inflight = _antigravityUserQuotaInflight.get(cacheKey);
if (inflight) return inflight;
const promise = (async () => {
try {
const response = await fetch(
"https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota",
{
method: "POST",
headers: {
Authorization: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
body: JSON.stringify({ project: projectId }),
signal: AbortSignal.timeout(10000),
}
);
if (!response.ok) return null;
const data = await response.json();
_antigravityUserQuotaCache.set(cacheKey, { data, fetchedAt: Date.now() });
return data;
} catch {
return null;
}
})().finally(() => {
_antigravityUserQuotaInflight.delete(cacheKey);
});
_antigravityUserQuotaInflight.set(cacheKey, promise);
return promise;
}
function extractCodeAssistTierId(subscription: JsonRecord): string {
const tierId = extractCodeAssistOnboardTierId(subscription);
if (tierId === "legacy-tier") return "";
@@ -2263,12 +2427,14 @@ async function probeAntigravityCreditBalanceUncached(
}
/**
* Antigravity Usage - Fetch quota from Google Cloud Code API
* Uses fetchAvailableModels API which returns ALL models (including Claude)
* with per-model quotaInfo (remainingFraction, resetTime).
* retrieveUserQuota only returns Gemini models — not suitable for Antigravity.
* Antigravity Usage - Fetch quota from Google Cloud Code API.
* fetchAvailableModels is catalog/eligibility data and may keep reporting full buckets
* after real usage. retrieveUserQuota is the consumption signal for Gemini-family
* buckets, so prefer it when present and fall back to fetchAvailableModels only for
* models that have no retrieveUserQuota entry (for example Claude/GPT OSS buckets).
*/
async function getAntigravityUsage(
provider: "antigravity" | "agy",
accessToken?: string,
providerSpecificData?: JsonRecord,
connectionProjectId?: string,
@@ -2319,33 +2485,52 @@ async function getAntigravityUsage(
);
}
const data = await fetchAntigravityAvailableModelsCached(accessToken, projectId, options);
const [data, userQuotaData] = await Promise.all([
fetchAntigravityAvailableModelsCached(accessToken, projectId, options),
fetchAntigravityUserQuotaCached(accessToken, projectId, options),
]);
const dataObj = toRecord(data);
if (dataObj.__antigravityForbidden === true) {
return { message: "Antigravity access forbidden. Check subscription." };
}
const modelEntries = toRecord(dataObj.models);
const userQuotaEntries = new Map<string, JsonRecord>();
const userQuotaObj = toRecord(userQuotaData);
if (Array.isArray(userQuotaObj.buckets)) {
for (const bucketValue of userQuotaObj.buckets) {
const bucket = toRecord(bucketValue);
const modelId = toClientAntigravityQuotaModelId(String(bucket.modelId || "").trim());
if (!modelId) continue;
userQuotaEntries.set(modelId, bucket);
}
}
const quotas: Record<string, UsageQuota> = {};
// Parse per-model quota info from fetchAvailableModels response.
for (const [modelKey, infoValue] of Object.entries(modelEntries)) {
for (const [rawModelKey, infoValue] of Object.entries(modelEntries)) {
const info = toRecord(infoValue);
const quotaInfo = toRecord(info.quotaInfo);
const modelKey = toClientAntigravityQuotaModelId(rawModelKey);
// Skip internal, excluded, and models without quota info
if (
!modelKey ||
info.isInternal === true ||
!isUserCallableAntigravityModelId(modelKey) ||
!(provider === "agy"
? isUserCallableAgyModelId(modelKey)
: isUserCallableAntigravityModelId(modelKey)) ||
Object.keys(quotaInfo).length === 0
) {
continue;
}
const rawFraction = toNumber(quotaInfo.remainingFraction, -1);
const resetAt = parseResetTime(quotaInfo.resetTime);
const liveQuota = userQuotaEntries.get(modelKey);
const quotaSource = liveQuota || quotaInfo;
const rawFraction = toNumber(quotaSource.remainingFraction, -1);
const resetAt = parseResetTime(quotaSource.resetTime);
// Distinguish "upstream did not report remainingFraction" from "remaining is 0%".
// A schema drift in Antigravity's quota API (very plausible — internal Google product)
// would otherwise silently mark every model as exhausted across the dashboard.
// fetchAvailableModels is a catalog view and can be stale/full; retrieveUserQuota is
// the source of truth for actual Gemini consumption when it includes the model.
const fractionReported = rawFraction >= 0;
if (!fractionReported) {
console.warn(
@@ -2362,13 +2547,50 @@ async function getAntigravityUsage(
const remaining = Math.round(total * remainingFraction);
const used = isUnlimited ? 0 : Math.max(0, total - remaining);
quotas[modelKey] = applyLocalUsageFallback(
{
used,
total: isUnlimited ? 0 : total,
resetAt,
remainingPercentage: isUnlimited ? 100 : remainingPercentage,
unlimited: isUnlimited,
fractionReported,
quotaSource: liveQuota ? "retrieveUserQuota" : "fetchAvailableModels",
},
provider,
connectionId,
modelKey
);
}
// Include retrieveUserQuota buckets not listed in the static/public Antigravity catalog yet.
// This keeps Provider Limits honest when Google adds a new Gemini tier before our catalog is
// updated. Hidden/internal catalog entries above are still filtered by the public pass.
for (const [modelKey, bucket] of userQuotaEntries) {
if (
quotas[modelKey] ||
!(provider === "agy"
? isUserCallableAgyModelId(modelKey)
: isUserCallableAntigravityModelId(modelKey))
) {
continue;
}
const rawFraction = toNumber(bucket.remainingFraction, -1);
if (rawFraction < 0) continue;
const remainingFraction = Math.max(0, Math.min(1, rawFraction));
const resetAt = parseResetTime(bucket.resetTime);
const isUnlimited = !resetAt && remainingFraction >= 1;
const QUOTA_NORMALIZED_BASE = 1000;
const total = QUOTA_NORMALIZED_BASE;
const remaining = Math.round(total * remainingFraction);
quotas[modelKey] = {
used,
used: isUnlimited ? 0 : Math.max(0, total - remaining),
total: isUnlimited ? 0 : total,
resetAt,
remainingPercentage: isUnlimited ? 100 : remainingPercentage,
remainingPercentage: isUnlimited ? 100 : remainingFraction * 100,
unlimited: isUnlimited,
fractionReported,
fractionReported: true,
quotaSource: "retrieveUserQuota",
};
}

27
src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx
View File

@@ -124,6 +124,21 @@ function getSoonestResetMs(quotas: any[] | undefined): number {
return soonest;
}
function shouldAutoRefreshQuota(provider: string, cached: any): boolean {
const quotas = cached?.quotas;
if (!Array.isArray(quotas) || quotas.length === 0) return true;
if (provider !== "antigravity" && provider !== "agy") return false;
return quotas.some(
(q: any) =>
q &&
typeof q.modelKey === "string" &&
q.modelKey.startsWith("gemini-") &&
!q.isCredits &&
q.quotaSource !== "retrieveUserQuota"
);
}
const getQuotaBarWidthClass = (pct: number) => {
if (pct <= 10) return "w-[10%]";
if (pct <= 20) return "w-1/5";
@@ -670,8 +685,7 @@ export default function ProviderLimits({
autoLiveFetchedRef.current = true;
for (const conn of visibleConnections) {
const cached = quotaData[conn.id];
const hasQuota = Array.isArray(cached?.quotas) && cached.quotas.length > 0;
if (!hasQuota) {
if (shouldAutoRefreshQuota(conn.provider, cached)) {
void fetchQuota(conn.id, conn.provider, { force: true }).catch(() => {});
}
}
@@ -779,7 +793,9 @@ export default function ProviderLimits({
onClick={refreshAll}
disabled={refreshingAll}
className="flex items-center gap-1.5 px-3.5 py-1.5 rounded-lg bg-bg-subtle border border-border text-text-main text-[13px] disabled:opacity-50 disabled:cursor-not-allowed cursor-pointer"
title={autoRefreshIntervalMs > 0 ? tr("autoRefreshing", "Auto-refreshing") : t("refreshAll")}
title={
autoRefreshIntervalMs > 0 ? tr("autoRefreshing", "Auto-refreshing") : t("refreshAll")
}
>
<span
className={`material-symbols-outlined text-[16px] ${refreshingAll ? "animate-spin" : ""}`}
@@ -790,7 +806,10 @@ export default function ProviderLimits({
? tr("refreshing", "Refreshing")
: autoRefreshIntervalMs > 0
? `${tr("autoRefreshing", "Auto-refreshing")} ${formatAutoRefreshCountdown(
Math.max(0, autoRefreshIntervalMs - (autoRefreshClock - lastRefreshAllAtRef.current))
Math.max(
0,
autoRefreshIntervalMs - (autoRefreshClock - lastRefreshAllAtRef.current)
)
)}`
: t("refreshAll")}
</button>

4
src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx
View File

@@ -290,6 +290,10 @@ export function parseQuotaData(provider, data) {
normalizedQuotas.push(
normalizeQuotaEntry(modelKey, quota, {
modelKey: modelKey,
...(quota?.quotaSource ? { quotaSource: quota.quotaSource } : {}),
...(quota?.fractionReported !== undefined
? { fractionReported: quota.fractionReported }
: {}),
})
);
});

6
src/app/api/usage/provider-limits/route.ts
View File

@@ -1,8 +1,8 @@
import { NextResponse } from "next/server";
import {
getCachedProviderLimitsMap,
getLastProviderLimitsAutoSyncTime,
getProviderLimitsSyncIntervalMinutes,
getSanitizedCachedProviderLimitsMap,
syncAllProviderLimits,
} from "@/lib/usage/providerLimits";
@@ -13,7 +13,7 @@ import {
export async function GET() {
try {
return NextResponse.json({
caches: getCachedProviderLimitsMap(),
caches: await getSanitizedCachedProviderLimitsMap(),
intervalMinutes: getProviderLimitsSyncIntervalMinutes(),
lastAutoSyncAt: await getLastProviderLimitsAutoSyncTime(),
});
@@ -30,7 +30,7 @@ export async function GET() {
export async function POST() {
try {
const result = await syncAllProviderLimits({ source: "manual" });
const caches = getCachedProviderLimitsMap();
const caches = await getSanitizedCachedProviderLimitsMap();
return NextResponse.json({
...result,
caches,

200
src/lib/usage/providerLimits.ts
View File

@@ -18,12 +18,21 @@ import { getMachineId } from "@/shared/utils/machine";
import { USAGE_SUPPORTED_PROVIDERS } from "@/shared/constants/providers";
import { getExecutor } from "@omniroute/open-sse/executors/index.ts";
import { getUsageForProvider } from "@omniroute/open-sse/services/usage.ts";
import { rotationGroupFor, serializeRefresh } from "@omniroute/open-sse/services/refreshSerializer.ts";
import {
rotationGroupFor,
serializeRefresh,
} from "@omniroute/open-sse/services/refreshSerializer.ts";
import {
extractCodeAssistOnboardTierId,
extractCodeAssistSubscriptionTier,
} from "@omniroute/open-sse/services/codeAssistSubscription.ts";
import { runWithProxyContext } from "@omniroute/open-sse/utils/proxyFetch.ts";
import {
isUserCallableAntigravityModelId,
toClientAntigravityModelId,
} from "@omniroute/open-sse/config/antigravityModelAliases.ts";
import { isUserCallableAgyModelId } from "@omniroute/open-sse/config/agyModels.ts";
import { onUsageRecorded } from "./usageEvents";
type JsonRecord = Record<string, unknown>;
@@ -64,6 +73,8 @@ const PROVIDER_LIMITS_APIKEY_PROVIDERS = new Set([
]);
const DEFAULT_PROVIDER_LIMITS_SYNC_INTERVAL_MINUTES = 70;
const PROVIDER_LIMITS_AUTO_SYNC_SETTING_KEY = "provider_limits_auto_sync_last_run";
const DEFAULT_PROVIDER_LIMITS_POST_USAGE_REFRESH_DELAY_MS = 5_000;
const pendingPostUsageRefreshes = new Set<string>();
function isRecord(value: unknown): value is JsonRecord {
return value !== null && typeof value === "object" && !Array.isArray(value);
@@ -83,6 +94,143 @@ function toProviderLimitsCacheEntry(
};
}
function getProviderLimitsPostUsageRefreshDelayMs(): number {
const raw = Number(process.env.PROVIDER_LIMITS_POST_USAGE_REFRESH_DELAY_MS ?? "");
return Number.isFinite(raw) && raw >= 0
? raw
: DEFAULT_PROVIDER_LIMITS_POST_USAGE_REFRESH_DELAY_MS;
}
function scheduleProviderLimitsPostUsageRefresh(connectionId: string): void {
if (!connectionId || pendingPostUsageRefreshes.has(connectionId)) return;
pendingPostUsageRefreshes.add(connectionId);
const timer = setTimeout(() => {
pendingPostUsageRefreshes.delete(connectionId);
void fetchAndPersistProviderLimits(connectionId, "scheduled", {
allowRotatingRefresh: true,
}).catch((error) => {
const message = error instanceof Error ? error.message : String(error);
console.warn(
`[ProviderLimits] Post-usage refresh failed for connection ${connectionId}: ${message}`
);
});
}, getProviderLimitsPostUsageRefreshDelayMs());
timer.unref?.();
}
export function notifyProviderUsageRecorded(
provider: string | null | undefined,
connectionId: string | null | undefined
): void {
if ((provider !== "antigravity" && provider !== "agy") || !connectionId) return;
scheduleProviderLimitsPostUsageRefresh(connectionId);
}
// Subscribe at module load so usageHistory can emit usage events without importing
// this module (and its executors/translator import graph). This module is loaded by
// the provider-limits route and the background auto-sync scheduler at server boot.
onUsageRecorded(notifyProviderUsageRecorded);
function isUsageQuotaKeyAllowed(provider: string, quotaKey: string): boolean {
if (quotaKey === "credits" || quotaKey === "models") return true;
if (provider === "antigravity") return isUserCallableAntigravityModelId(quotaKey);
if (provider === "agy") return isUserCallableAgyModelId(quotaKey);
return true;
}
function normalizeUsageQuotaKey(provider: string, quotaKey: string): string | null {
if (quotaKey === "credits" || quotaKey === "models") return quotaKey;
if (provider === "antigravity" || provider === "agy") {
const clientKey = toClientAntigravityModelId(quotaKey);
return isUsageQuotaKeyAllowed(provider, clientKey) ? clientKey : null;
}
return isUsageQuotaKeyAllowed(provider, quotaKey) ? quotaKey : null;
}
function normalizeUsageQuotasForProvider(
provider: string,
quotas: JsonRecord | null | undefined
): JsonRecord | null {
if (!isRecord(quotas)) return quotas ?? null;
const normalized: JsonRecord = {};
let changed = false;
for (const [quotaKey, quota] of Object.entries(quotas)) {
const normalizedKey = normalizeUsageQuotaKey(provider, quotaKey);
if (!normalizedKey) {
changed = true;
continue;
}
const existing = normalized[normalizedKey];
if (existing && isRecord(existing) && isRecord(quota)) {
const existingSource = String(existing.quotaSource ?? "");
const nextSource = String(quota.quotaSource ?? "");
const sourceRank: Record<string, number> = {
fetchAvailableModels: 0,
localUsageHistory: 1,
retrieveUserQuota: 2,
};
if ((sourceRank[existingSource] ?? 0) > (sourceRank[nextSource] ?? 0)) {
continue;
}
}
normalized[normalizedKey] = quota as JsonRecord;
if (normalizedKey !== quotaKey) changed = true;
}
return changed ? normalized : quotas;
}
function sanitizeUsageQuotasForProvider(provider: string, usage: JsonRecord): JsonRecord {
if (provider !== "antigravity" && provider !== "agy") return usage;
if (!isRecord(usage.quotas)) return usage;
const sanitizedQuotas = normalizeUsageQuotasForProvider(provider, usage.quotas);
return sanitizedQuotas === usage.quotas ? usage : { ...usage, quotas: sanitizedQuotas };
}
function hasRetrieveUserQuotaSource(
provider: string,
cache: ProviderLimitsCacheEntry | undefined
): boolean {
if (provider !== "antigravity" && provider !== "agy") return true;
if (!cache?.quotas) return false;
return Object.values(cache.quotas).some((quota) => {
if (!isRecord(quota)) return false;
return quota.quotaSource === "retrieveUserQuota";
});
}
function sanitizeProviderLimitsCacheForConnection(
connection: ProviderConnectionLike | null | undefined,
entry: ProviderLimitsCacheEntry | null
): ProviderLimitsCacheEntry | null {
if (!connection || !entry || !entry.quotas) return entry;
if (connection.provider !== "antigravity" && connection.provider !== "agy") return entry;
const sanitizedQuotas = normalizeUsageQuotasForProvider(connection.provider, entry.quotas);
return sanitizedQuotas === entry.quotas ? entry : { ...entry, quotas: sanitizedQuotas };
}
function shouldRefreshProviderLimitsCache(
connection: ProviderConnectionLike,
cache: ProviderLimitsCacheEntry | undefined
): boolean {
if (!cache?.quotas) return true;
if (connection.provider !== "antigravity" && connection.provider !== "agy") return false;
return (
!hasRetrieveUserQuotaSource(connection.provider, cache) ||
Object.keys(cache.quotas).some(
(quotaKey) => !isUsageQuotaKeyAllowed(connection.provider, quotaKey)
)
);
}
function isSupportedUsageConnection(connection: ProviderConnectionLike | null): boolean {
if (
!connection ||
@@ -162,9 +310,18 @@ export async function refreshAndUpdateCredentials(
// Serialize the actual token mint per rotation group so two sibling accounts
// never hit Auth0 concurrently (passthrough for non-rotating providers).
const refreshResult = await serializeRefresh(connection.provider, () =>
const refreshResult = (await serializeRefresh(connection.provider, () =>
executor.refreshCredentials(credentials, console)
);
)) as
| (JsonRecord & {
accessToken?: string;
refreshToken?: string;
expiresIn?: number;
expiresAt?: string;
copilotToken?: string;
copilotTokenExpiresAt?: string;
})
| null;
if (!refreshResult) {
if (connection.provider === "github" && connection.accessToken) {
@@ -443,6 +600,22 @@ export function getCachedProviderLimitsMap(): Record<string, ProviderLimitsCache
return getAllProviderLimitsCache();
}
export async function getSanitizedCachedProviderLimitsMap(): Promise<
Record<string, ProviderLimitsCacheEntry>
> {
const caches = getAllProviderLimitsCache();
const connections = (await getProviderConnections({
isActive: true,
})) as unknown as ProviderConnectionLike[];
const byId = new Map(connections.map((conn) => [conn.id, conn]));
const sanitized: Record<string, ProviderLimitsCacheEntry> = {};
for (const [connectionId, entry] of Object.entries(caches)) {
sanitized[connectionId] =
sanitizeProviderLimitsCacheForConnection(byId.get(connectionId), entry) || entry;
}
return sanitized;
}
export async function fetchLiveProviderLimits(connectionId: string): Promise<{
connection: ProviderConnectionLike;
usage: JsonRecord;
@@ -469,7 +642,10 @@ async function fetchLiveProviderLimitsWithOptions(
}
if (connection.authType !== "oauth") {
const usage = (await getUsageForProvider(connection, options)) as JsonRecord;
const usage = sanitizeUsageQuotasForProvider(
connection.provider,
(await getUsageForProvider(connection as unknown as JsonRecord, options)) as JsonRecord
);
if (isRecord(usage.quotas)) {
setQuotaCache(connectionId, connection.provider, usage.quotas);
}
@@ -497,7 +673,10 @@ async function fetchLiveProviderLimitsWithOptions(
await syncToCloudIfEnabled();
}
let usageData = (await getUsageForProvider(conn, options)) as JsonRecord;
let usageData = sanitizeUsageQuotasForProvider(
conn.provider,
(await getUsageForProvider(conn as unknown as JsonRecord, options)) as JsonRecord
);
// Reactive 401 recovery (on-demand/force path only): an unauthorized usage
// response means the access token is actually dead. Force ONE serialized
@@ -512,7 +691,10 @@ async function fetchLiveProviderLimitsWithOptions(
if (forced.refreshed) {
conn = forced.connection;
await syncToCloudIfEnabled();
usageData = (await getUsageForProvider(conn, options)) as JsonRecord;
usageData = sanitizeUsageQuotasForProvider(
conn.provider,
(await getUsageForProvider(conn as unknown as JsonRecord, options)) as JsonRecord
);
}
}
@@ -679,8 +861,12 @@ export async function syncAllProviderLimits(
};
const fetchOne = async (connection: ProviderConnectionLike) => {
const existingCache = getProviderLimitsCache(connection.id);
const forceRefresh =
source === "manual" ||
shouldRefreshProviderLimitsCache(connection, existingCache || undefined);
const { usage } = await fetchLiveProviderLimitsWithOptions(connection.id, {
forceRefresh: source === "manual",
forceRefresh,
});
const cache = toProviderLimitsCacheEntry(usage, source);
return { connectionId: connection.id, cache };

40
src/lib/usage/usageEvents.ts Normal file
View File

@@ -0,0 +1,40 @@
/**
* Lightweight usage-event bus.
*
* Decouples usage recording (usageHistory.ts) from the provider-limits
* subsystem (providerLimits.ts). usageHistory must NOT import providerLimits:
* providerLimits pulls in the executors barrel (and the whole translator graph),
* so a direct or dynamic import from usageHistory expands the type-check surface
* across modules that have nothing to do with usage recording. usageHistory
* emits here; providerLimits subscribes at module load.
*
* @module lib/usage/usageEvents
*/
export type UsageRecordedListener = (provider: string, connectionId: string) => void;
const listeners = new Set<UsageRecordedListener>();
/** Register a listener for usage-recorded events. Returns an unsubscribe fn. */
export function onUsageRecorded(listener: UsageRecordedListener): () => void {
listeners.add(listener);
return () => {
listeners.delete(listener);
};
}
/** Emit a usage-recorded event. No-ops when provider/connectionId is missing. */
export function emitUsageRecorded(
provider: string | null | undefined,
connectionId: string | null | undefined
): void {
if (!provider || !connectionId) return;
for (const listener of listeners) {
try {
listener(provider, connectionId);
} catch (error) {
const message = error instanceof Error ? error.message : String(error);
console.warn(`[usageEvents] usage-recorded listener failed: ${message}`);
}
}
}

21
src/lib/usage/usageHistory.ts
View File

@@ -10,6 +10,7 @@
import { getDbInstance } from "../db/core";
import { protectPayloadForLog } from "../logPayloads";
import { shouldPersistToDisk } from "./migrations";
import { emitUsageRecorded } from "./usageEvents";
import {
getLoggedInputTokens,
getLoggedOutputTokens,
@@ -74,6 +75,7 @@ function toNumber(value: unknown): number {
return 0;
}
function percentile(sortedValues: number[], p: number): number {
if (sortedValues.length === 0) return 0;
if (sortedValues.length === 1) return sortedValues[0];
@@ -340,7 +342,8 @@ export function finalizeMostRecentPendingRequest(
// detail from persisted call_log artifacts (best-effort, non-blocking).
(async () => {
try {
const missingProvider = updated.providerResponse === undefined || updated.providerResponse === null;
const missingProvider =
updated.providerResponse === undefined || updated.providerResponse === null;
const missingClient = updated.clientResponse === undefined || updated.clientResponse === null;
if ((missingProvider || missingClient) && connectionId) {
const db = getDbInstance();
@@ -363,7 +366,10 @@ export function finalizeMostRecentPendingRequest(
if (missingClient && pipeline?.clientResponse) {
updated.clientResponse = pipeline.clientResponse;
}
if ((missingProvider && art.artifact.responseBody) || (missingClient && art.artifact.responseBody)) {
if (
(missingProvider && art.artifact.responseBody) ||
(missingClient && art.artifact.responseBody)
) {
// use responseBody as a fallback for both
if (missingProvider) updated.providerResponse = art.artifact.responseBody;
if (missingClient) updated.clientResponse = art.artifact.responseBody;
@@ -376,7 +382,12 @@ export function finalizeMostRecentPendingRequest(
}
}
} catch (e) {
try { console.warn("[usageHistory] failed to enrich completed detail from artifacts:", (e && (e.message || e))); } catch {}
try {
console.warn(
"[usageHistory] failed to enrich completed detail from artifacts:",
e && (e.message || e)
);
} catch {}
}
})();
@@ -582,6 +593,10 @@ export async function saveRequestUsage(entry: any) {
entry.comboStrategy || entry.combo_strategy || null,
timestamp
);
// Decoupled via the event bus so usageHistory never imports providerLimits
// (which would pull the executors/translator graph into the type-check surface).
emitUsageRecorded(entry.provider, entry.connectionId);
} catch (error) {
console.error("Failed to save usage stats:", error);
}

41
tests/unit/agy-usage-quota.test.ts
View File

@@ -1,4 +1,4 @@
import test, { mock } from "node:test";
import test from "node:test";
import assert from "node:assert/strict";
const usageModule = await import("../../open-sse/services/usage.ts");
@@ -13,19 +13,21 @@ test("agy is registered for the background usage fetcher", () => {
});
test("getUsageForProvider routes agy through the Antigravity usage implementation", async () => {
const fetchMock = mock.method(globalThis, "fetch", async () => ({
ok: true,
json: async () => ({
models: {
"gemini-3.5-flash-preview": {
quotaInfo: {
remainingFraction: 0.75,
resetTime: "2026-06-06T00:00:00Z",
const originalFetch = globalThis.fetch;
globalThis.fetch = async () =>
new Response(
JSON.stringify({
models: {
"gemini-3-flash-agent": {
quotaInfo: {
remainingFraction: 0.75,
resetTime: "2026-06-06T00:00:00Z",
},
},
},
},
}),
}));
}),
{ status: 200, headers: { "Content-Type": "application/json" } }
);
try {
const result = await usageModule.getUsageForProvider(
@@ -46,11 +48,16 @@ test("getUsageForProvider routes agy through the Antigravity usage implementatio
);
assert.ok("quotas" in result, "agy should return quota data when upstream responds");
const quota = (result as { quotas: Record<string, any> }).quotas["gemini-3.5-flash-preview"];
assert.ok(quota, "should expose the Antigravity per-model quota for agy");
const quota = (result as { quotas: Record<string, any> }).quotas["gemini-3.5-flash-high"];
assert.ok(quota, "should expose the clean agy per-model quota");
assert.equal(quota.remainingPercentage, 75);
assert.equal(
(result as { quotas: Record<string, any> }).quotas["gemini-3-flash-agent"],
undefined,
"agy quota should not expose retired upstream IDs"
);
} finally {
fetchMock.mock.restore();
globalThis.fetch = originalFetch;
}
});
@@ -58,7 +65,7 @@ test("parseQuotaData treats agy quota payloads like Antigravity", () => {
const parsed = providerLimitUtils.parseQuotaData("agy", {
quotas: {
credits: { remaining: 42 },
"gemini-3.5-flash-preview": {
"gemini-3.5-flash-high": {
used: 250,
total: 1000,
remainingPercentage: 75,
@@ -76,7 +83,7 @@ test("parseQuotaData treats agy quota payloads like Antigravity", () => {
assert.ok(credits, "credits quota should be rendered");
assert.equal(credits.isCredits, true);
const modelQuota = parsed.find((quota: any) => quota.name === "gemini-3.5-flash-preview");
const modelQuota = parsed.find((quota: any) => quota.name === "gemini-3.5-flash-high");
assert.ok(modelQuota, "model quota should be rendered");
assert.equal(modelQuota.remainingPercentage, 75);
});

206
tests/unit/antigravity-usage-service.test.ts
View File

@@ -9,7 +9,7 @@
* - 0.5 → 50% remaining (partial quota)
*/
import { describe, it, mock } from "node:test";
import { describe, it } from "node:test";
import assert from "node:assert/strict";
// IMPORTANT: load usage.ts up-front so the proxyFetch patch in
@@ -29,19 +29,21 @@ describe("getUsageForProvider (antigravity in usage.ts)", () => {
};
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",
const originalFetch = globalThis.fetch;
globalThis.fetch = async () =>
({
ok: true,
json: async () => ({
models: {
"gemini-3.5-flash-high": {
quotaInfo: {
remainingFraction: undefined,
resetTime: "2026-05-26T00:00:00Z",
},
},
},
},
}),
}));
}),
}) as Response;
try {
const result = await getUsageForProvider(connectionBase, { forceRefresh: true });
@@ -49,31 +51,33 @@ describe("getUsageForProvider (antigravity in usage.ts)", () => {
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");
const quota = result.quotas["gemini-3.5-flash-high"];
assert.ok(quota, "should have quota for gemini-3.5-flash-high");
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();
globalThis.fetch = originalFetch;
}
});
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",
const originalFetch = globalThis.fetch;
globalThis.fetch = async () =>
({
ok: true,
json: async () => ({
models: {
"gemini-3.5-flash-high": {
quotaInfo: {
remainingFraction: 0,
resetTime: "2026-05-26T00:00:00Z",
},
},
},
},
}),
}));
}),
}) as Response;
try {
const usageModule = await import("../../open-sse/services/usage.ts");
@@ -84,30 +88,32 @@ describe("getUsageForProvider (antigravity in usage.ts)", () => {
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");
const quota = result.quotas["gemini-3.5-flash-high"];
assert.ok(quota, "should have quota for gemini-3.5-flash-high");
assert.equal(quota.remainingPercentage, 0, "remaining should be 0%");
assert.equal(quota.unlimited, false, "should not be unlimited");
}
} finally {
mockFetch.mock.restore();
globalThis.fetch = originalFetch;
}
});
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",
const originalFetch = globalThis.fetch;
globalThis.fetch = async () =>
({
ok: true,
json: async () => ({
models: {
"gemini-3.5-flash-high": {
quotaInfo: {
remainingFraction: 1.0,
resetTime: "2026-05-26T00:00:00Z",
},
},
},
},
}),
}));
}),
}) as Response;
try {
const usageModule = await import("../../open-sse/services/usage.ts");
@@ -118,29 +124,31 @@ describe("getUsageForProvider (antigravity in usage.ts)", () => {
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");
const quota = result.quotas["gemini-3.5-flash-high"];
assert.ok(quota, "should have quota for gemini-3.5-flash-high");
assert.equal(quota.remainingPercentage, 100, "remaining should be 100%");
assert.equal(quota.unlimited, false, "should not be unlimited (has resetTime)");
}
} finally {
mockFetch.mock.restore();
globalThis.fetch = originalFetch;
}
});
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,
const originalFetch = globalThis.fetch;
globalThis.fetch = async () =>
({
ok: true,
json: async () => ({
models: {
"gemini-3.1-flash-lite": {
quotaInfo: {
remainingFraction: 1.0,
},
},
},
},
}),
}));
}),
}) as Response;
try {
const usageModule = await import("../../open-sse/services/usage.ts");
@@ -157,24 +165,26 @@ describe("getUsageForProvider (antigravity in usage.ts)", () => {
assert.equal(quota.unlimited, true, "should be unlimited (no resetTime)");
}
} finally {
mockFetch.mock.restore();
globalThis.fetch = originalFetch;
}
});
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",
const originalFetch = globalThis.fetch;
globalThis.fetch = async () =>
({
ok: true,
json: async () => ({
models: {
"gemini-3.5-flash-high": {
quotaInfo: {
remainingFraction: 0.5,
resetTime: "2026-05-26T00:00:00Z",
},
},
},
},
}),
}));
}),
}) as Response;
try {
const usageModule = await import("../../open-sse/services/usage.ts");
@@ -185,30 +195,32 @@ describe("getUsageForProvider (antigravity in usage.ts)", () => {
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");
const quota = result.quotas["gemini-3.5-flash-high"];
assert.ok(quota, "should have quota for gemini-3.5-flash-high");
assert.equal(quota.remainingPercentage, 50, "remaining should be 50%");
assert.equal(quota.unlimited, false, "should not be unlimited");
}
} finally {
mockFetch.mock.restore();
globalThis.fetch = originalFetch;
}
});
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",
const originalFetch = globalThis.fetch;
globalThis.fetch = async () =>
({
ok: true,
json: async () => ({
models: {
"gemini-3.5-flash-high": {
quotaInfo: {
remainingFraction: 1.5,
resetTime: "2026-05-26T00:00:00Z",
},
},
},
},
}),
}));
}),
}) as Response;
try {
const usageModule = await import("../../open-sse/services/usage.ts");
@@ -219,30 +231,32 @@ describe("getUsageForProvider (antigravity in usage.ts)", () => {
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");
const quota = result.quotas["gemini-3.5-flash-high"];
assert.ok(quota, "should have quota for gemini-3.5-flash-high");
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();
globalThis.fetch = originalFetch;
}
});
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",
const originalFetch = globalThis.fetch;
globalThis.fetch = async () =>
({
ok: true,
json: async () => ({
models: {
"gemini-3.5-flash-high": {
quotaInfo: {
remainingFraction: -0.5,
resetTime: "2026-05-26T00:00:00Z",
},
},
},
},
}),
}));
}),
}) as Response;
try {
const usageModule = await import("../../open-sse/services/usage.ts");
@@ -253,13 +267,13 @@ describe("getUsageForProvider (antigravity in usage.ts)", () => {
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");
const quota = result.quotas["gemini-3.5-flash-high"];
assert.ok(quota, "should have quota for gemini-3.5-flash-high");
assert.equal(quota.remainingPercentage, 0, "remaining should be clamped to 0%");
assert.equal(quota.unlimited, false, "should not be unlimited");
}
} finally {
mockFetch.mock.restore();
globalThis.fetch = originalFetch;
}
});
});

6
tests/unit/provider-columns.test.ts
View File

@@ -68,14 +68,15 @@ test("getProviderColumns: Antigravity falls back to dynamic schema (first 3 quot
"claude-opus-4-6-thinking": { used: 0, total: 100, remainingPercentage: 100 },
"claude-sonnet-4-6": { used: 0, total: 100, remainingPercentage: 100 },
"gemini-3.1-pro-low": { used: 0, total: 100, remainingPercentage: 100 },
"gemini-3-flash-agent": { used: 0, total: 100, remainingPercentage: 100 },
"gemini-3.5-flash-low": { used: 0, total: 100, remainingPercentage: 100 },
"gemini-3.5-flash-medium": { used: 0, total: 100, remainingPercentage: 100 },
"gemini-3.5-flash-high": { used: 0, total: 100, remainingPercentage: 100 },
},
});
const schema = providerColumns.getProviderColumns("antigravity", quotas);
assert.equal(schema.columns.length, providerColumns.MAX_DYNAMIC_COLUMNS);
assert.equal(schema.overflowCount, 2, "5 quotas - 3 visible = 2 overflow");
assert.equal(schema.overflowCount, 3, "6 quotas - 3 visible = 3 overflow");
});
test("getProviderColumns: credits never become columns, always counted toward overflow", () => {
@@ -105,7 +106,6 @@ test("getProviderColumns: unknown provider uses dynamic fallback", () => {
});
test("getProviderColumns: tolerates non-array quotas", () => {
// @ts-expect-error — exercise the runtime guard
const schema = providerColumns.getProviderColumns("codex", null);
assert.equal(schema.columns.length, 2);
assert.equal(schema.columns[0].quota, null);

58
tests/unit/usage-events.test.ts Normal file
View File

@@ -0,0 +1,58 @@
import test from "node:test";
import assert from "node:assert/strict";
import { emitUsageRecorded, onUsageRecorded } from "../../src/lib/usage/usageEvents.ts";
test("emitUsageRecorded delivers provider + connectionId to subscribers", () => {
const seen: Array<[string, string]> = [];
const off = onUsageRecorded((provider, connectionId) => seen.push([provider, connectionId]));
try {
emitUsageRecorded("antigravity", "conn-1");
assert.deepEqual(seen, [["antigravity", "conn-1"]]);
} finally {
off();
}
});
test("emitUsageRecorded no-ops when provider or connectionId is missing", () => {
let calls = 0;
const off = onUsageRecorded(() => {
calls += 1;
});
try {
emitUsageRecorded(null, "conn-1");
emitUsageRecorded("agy", "");
emitUsageRecorded(undefined, undefined);
assert.equal(calls, 0);
} finally {
off();
}
});
test("onUsageRecorded unsubscribe stops further delivery", () => {
let calls = 0;
const off = onUsageRecorded(() => {
calls += 1;
});
emitUsageRecorded("agy", "conn-2");
off();
emitUsageRecorded("agy", "conn-2");
assert.equal(calls, 1);
});
test("a throwing listener does not break emit for others", () => {
let reached = false;
const offBad = onUsageRecorded(() => {
throw new Error("boom");
});
const offGood = onUsageRecorded(() => {
reached = true;
});
try {
assert.doesNotThrow(() => emitUsageRecorded("antigravity", "conn-3"));
assert.equal(reached, true);
} finally {
offBad();
offGood();
}
});

113
tests/unit/usage-service-hardening.test.ts
View File

@@ -371,6 +371,119 @@ test("usage service covers Antigravity quota parsing, exclusions and forbidden a
assert.match(forbidden.message, /forbidden/i);
});
test("usage service prefers Antigravity retrieveUserQuota over catalog quotaInfo", async () => {
globalThis.fetch = async (url) => {
const urlString = String(url);
if (urlString.includes("loadCodeAssist")) {
return new Response(
JSON.stringify({
allowedTiers: [{ id: "tier_pro", isDefault: true }],
cloudaicompanionProject: "ag-project",
}),
{ status: 200 }
);
}
if (urlString.includes("fetchAvailableModels")) {
return new Response(
JSON.stringify({
models: {
"gemini-3.5-flash-high": {
quotaInfo: {
remainingFraction: 1,
resetTime: new Date(Date.now() + 60_000).toISOString(),
},
},
},
}),
{ status: 200 }
);
}
if (urlString.includes("retrieveUserQuota")) {
return new Response(
JSON.stringify({
buckets: [
{
modelId: "gemini-3.5-flash-high",
remainingFraction: 0.25,
resetTime: new Date(Date.now() + 60_000).toISOString(),
},
],
}),
{ status: 200 }
);
}
throw new Error(`unexpected fetch: ${url}`);
};
const usage: any = await usageService.getUsageForProvider({
provider: "antigravity",
accessToken: `ag-token-live-quota-${Date.now()}`,
});
assert.equal(usage.quotas["gemini-3.5-flash-high"].remainingPercentage, 25);
assert.equal(usage.quotas["gemini-3.5-flash-high"].used, 750);
assert.equal(usage.quotas["gemini-3.5-flash-high"].quotaSource, "retrieveUserQuota");
});
test("usage service normalizes retired Antigravity quota bucket ids", async () => {
globalThis.fetch = async (url) => {
const urlString = String(url);
if (urlString.includes("loadCodeAssist")) {
return new Response(
JSON.stringify({
allowedTiers: [{ id: "tier_pro", isDefault: true }],
cloudaicompanionProject: "ag-project",
}),
{ status: 200 }
);
}
if (urlString.includes("fetchAvailableModels")) {
return new Response(
JSON.stringify({
models: {
"gemini-3.5-flash-low": { quotaInfo: { remainingFraction: 1 } },
"gemini-3.5-flash-high": { quotaInfo: { remainingFraction: 1 } },
"gemini-3-flash-agent": { quotaInfo: { remainingFraction: 1 } },
"gemini-3.5-flash-extra-low": { quotaInfo: { remainingFraction: 1 } },
},
}),
{ status: 200 }
);
}
if (urlString.includes("retrieveUserQuota")) {
return new Response(
JSON.stringify({
buckets: [
{ modelId: "gemini-3-flash-agent", remainingFraction: 0.5 },
{ modelId: "gemini-3.5-flash-extra-low", remainingFraction: 0.25 },
],
}),
{ status: 200 }
);
}
throw new Error(`unexpected fetch: ${url}`);
};
const usage: any = await usageService.getUsageForProvider({
provider: "antigravity",
accessToken: `ag-token-legacy-buckets-${Date.now()}`,
});
assert.equal(usage.quotas["gemini-3-flash-agent"], undefined);
assert.equal(usage.quotas["gemini-3.5-flash-extra-low"], undefined);
assert.equal(usage.quotas["gemini-3.5-flash-high"].remainingPercentage, 50);
assert.equal(usage.quotas["gemini-3.5-flash-medium"].remainingPercentage, 100);
assert.equal(usage.quotas["gemini-3.5-flash-low"].remainingPercentage, 25);
});
test("usage service retries Antigravity fetchAvailableModels across the shared fallback order", async () => {
const calls: any[] = [];
const expectedQuotaUrls = getAntigravityFetchAvailableModelsUrls();