diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index 1a35c6b41a..1dd48479b7 100644 --- a/open-sse/services/usage.ts +++ b/open-sse/services/usage.ts @@ -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).code, 200); if (code === 401 || code === 403 || (json as Record).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).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(); const _antigravityAvailableModelsInflight = new Map>(); +const _antigravityUserQuotaCache = new Map(); +const _antigravityUserQuotaInflight = new Map>(); const _antigravityCreditProbeCache = new Map(); const _antigravityCreditProbeInflight = new Map>(); @@ -1926,27 +1941,125 @@ const _antigravityCreditProbeInflight = new Map>( // 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 = { + "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 { + 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(); + 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 = {}; // 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", }; } diff --git a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx index 7653c1a16f..968c189fe9 100644 --- a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx +++ b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx @@ -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") + } > 0 ? `${tr("autoRefreshing", "Auto-refreshing")} ${formatAutoRefreshCountdown( - Math.max(0, autoRefreshIntervalMs - (autoRefreshClock - lastRefreshAllAtRef.current)) + Math.max( + 0, + autoRefreshIntervalMs - (autoRefreshClock - lastRefreshAllAtRef.current) + ) )}` : t("refreshAll")} diff --git a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx index 159c8bd46e..3e02c550cc 100644 --- a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx +++ b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.tsx @@ -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 } + : {}), }) ); }); diff --git a/src/app/api/usage/provider-limits/route.ts b/src/app/api/usage/provider-limits/route.ts index 2c1fdd6541..03403ca297 100644 --- a/src/app/api/usage/provider-limits/route.ts +++ b/src/app/api/usage/provider-limits/route.ts @@ -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, diff --git a/src/lib/usage/providerLimits.ts b/src/lib/usage/providerLimits.ts index bf6edd7df5..820025d350 100644 --- a/src/lib/usage/providerLimits.ts +++ b/src/lib/usage/providerLimits.ts @@ -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; @@ -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(); 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 = { + 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 +> { + 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 = {}; + 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 }; diff --git a/src/lib/usage/usageEvents.ts b/src/lib/usage/usageEvents.ts new file mode 100644 index 0000000000..dbd2a507eb --- /dev/null +++ b/src/lib/usage/usageEvents.ts @@ -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(); + +/** 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}`); + } + } +} diff --git a/src/lib/usage/usageHistory.ts b/src/lib/usage/usageHistory.ts index bea7f91db6..2c6359f2ee 100644 --- a/src/lib/usage/usageHistory.ts +++ b/src/lib/usage/usageHistory.ts @@ -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); } diff --git a/tests/unit/agy-usage-quota.test.ts b/tests/unit/agy-usage-quota.test.ts index 61213df837..a99f8e1ac9 100644 --- a/tests/unit/agy-usage-quota.test.ts +++ b/tests/unit/agy-usage-quota.test.ts @@ -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 }).quotas["gemini-3.5-flash-preview"]; - assert.ok(quota, "should expose the Antigravity per-model quota for agy"); + const quota = (result as { quotas: Record }).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 }).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); }); diff --git a/tests/unit/antigravity-usage-service.test.ts b/tests/unit/antigravity-usage-service.test.ts index 69985175c1..30d1f0d6b8 100644 --- a/tests/unit/antigravity-usage-service.test.ts +++ b/tests/unit/antigravity-usage-service.test.ts @@ -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; } }); }); diff --git a/tests/unit/provider-columns.test.ts b/tests/unit/provider-columns.test.ts index eed1c56d09..4034cd3424 100644 --- a/tests/unit/provider-columns.test.ts +++ b/tests/unit/provider-columns.test.ts @@ -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); diff --git a/tests/unit/usage-events.test.ts b/tests/unit/usage-events.test.ts new file mode 100644 index 0000000000..554fa9cfde --- /dev/null +++ b/tests/unit/usage-events.test.ts @@ -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(); + } +}); diff --git a/tests/unit/usage-service-hardening.test.ts b/tests/unit/usage-service-hardening.test.ts index 09701ddeb8..a24fe9c930 100644 --- a/tests/unit/usage-service-hardening.test.ts +++ b/tests/unit/usage-service-hardening.test.ts @@ -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();