From 766967d6885ffca9bd590df41148a56439ec19db Mon Sep 17 00:00:00 2001 From: Diego Rodrigues de Sa e Souza Date: Wed, 24 Jun 2026 13:01:56 -0300 Subject: [PATCH] =?UTF-8?q?refactor(sse):=20extrai=20a=20fam=C3=ADlia=20An?= =?UTF-8?q?tigravity=20de=20services/usage.ts?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Continua a decomposição do god-file services/usage.ts (empilhado sobre a fatia GLM). Move a família Antigravity (Gemini Code Assist) completa para services/usage/antigravity.ts (794L): fallback de uso local, mapeamento de tier/plan code-assist, probing de saldo de crédito, os fetchers de user-quota + available-models (com seus caches de módulo) e getAntigravityUsage, mais ANTIGRAVITY_CONFIG e a interface AntigravityUsageOptions. Timer compartilhado DIVIDIDO (decisão do operador): o setInterval de purga de TTL em usage.ts varria 5 caches juntos (1 do Gemini + 4 do Antigravity). Agora cada módulo é dono dos seus caches e da sua limpeza — o leaf ganha os 4 caches + setInterval próprio (varre só eles); o host mantém o cache Gemini + um timer só pra ele. Runtime passa de 1 para 2 timers (funcionalmente equivalente, ambos unref'd). Única mudança não-byte-idêntica da fatia; o resto é move verbatim. - Importa só os leaves irmãos scalars/quota + os módulos antigravity-* (config/ headers/identity/credits/clientProfile/codeAssist) e executors/antigravity. - usage.ts reimporta getAntigravityUsage (dispatcher) + getAntigravityPlanLabel/ mapCodeAssist* (__testing). 9 imports antigravity órfãos removidos do host (PROVIDERS, antigravityHeaders, antigravityIdentity, etc — agora só no leaf). usage.ts 2670 → 1915. Validação: typecheck:core, eslint 0/0, check:cycles, check:file-size (leaf 794<800), check:public-creds OK, usage-utils (47/47), usage-service-hardening (24/24), code-assist-subscription (6/6) + caracterização nova (4/4, surface + identidade do re-export __testing). --- open-sse/services/usage.ts | 781 +---------------- open-sse/services/usage/antigravity.ts | 794 ++++++++++++++++++ .../usage-antigravity-family-split.test.ts | 45 + 3 files changed, 853 insertions(+), 767 deletions(-) create mode 100644 open-sse/services/usage/antigravity.ts create mode 100644 tests/unit/usage-antigravity-family-split.test.ts diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index 9ef9fffe0919..76ac708a5061 100644 --- a/open-sse/services/usage.ts +++ b/open-sse/services/usage.ts @@ -2,16 +2,6 @@ * Usage Fetcher - Get usage data from provider APIs */ -import { PROVIDERS } from "../config/constants.ts"; -import { - getAntigravityFetchAvailableModelsUrls, - ANTIGRAVITY_BASE_URLS, -} from "../config/antigravityUpstream.ts"; -import { - isUserCallableAntigravityModelId, - toClientAntigravityQuotaModelId, -} from "../config/antigravityModelAliases.ts"; -import { isUserCallableAgyModelId } from "../config/agyModels.ts"; import { buildCodexUsageQuotas } from "./codexUsageQuotas.ts"; import { getGlmQuotaUrl } from "../config/glmProvider.ts"; import { getGitHubCopilotInternalUserHeaders } from "../config/providerHeaderProfiles.ts"; @@ -22,24 +12,8 @@ import { fetchDeepseekQuota, type DeepseekQuota } from "./deepseekQuotaFetcher.t import { fetchOpencodeQuota, type OpencodeTripleWindowQuota } from "./opencodeQuotaFetcher.ts"; import { getOllamaCloudUsage, getOpenCodeGoUsage } from "./opencodeOllamaUsage.ts"; import { getCodeBuddyCnUsage } from "./usage/codebuddy-cn.ts"; -import { - applyAntigravityClientProfileHeaders, - getAntigravityBootstrapHeaders, - getAntigravityClientProfile, -} from "./antigravityClientProfile.ts"; -import { - antigravityUserAgent, - getAntigravityHeaders, - getAntigravityLoadCodeAssistMetadata, -} from "./antigravityHeaders.ts"; -import { - getAntigravityRemainingCredits, - updateAntigravityRemainingCredits, -} from "../executors/antigravity.ts"; -import { getCreditsMode } from "./antigravityCredits.ts"; import { CLAUDE_CODE_VERSION, fetchClaudeBootstrap } from "../executors/claudeIdentity.ts"; import { isClaudeOauthUsageCoolingDown, markClaudeOauthUsage429 } from "./claudeUsageCooldown.ts"; -import { generateAntigravityRequestId, getAntigravitySessionId } from "./antigravityIdentity.ts"; import { extractCodeAssistOnboardTierId, extractCodeAssistSubscriptionTier, @@ -74,6 +48,13 @@ import { import { getGlmUsage } from "./usage/glm.ts"; // Re-exported para o teste glm-coding-plan-monthly (importa de services/usage). export { glmMonthlyRemainingPercentage } from "./usage/glm.ts"; +import { + getAntigravityUsage, + getAntigravityPlanLabel, + mapCodeAssistSubscriptionToPlanLabel, + mapCodeAssistTierIdToLabel, + mapSubscriptionTierStringToPlanLabel, +} from "./usage/antigravity.ts"; // Quota / usage upstream URLs (overridable for testing or relays). const CROF_USAGE_URL = process.env.OMNIROUTE_CROF_USAGE_URL ?? "https://crof.ai/usage_api/"; @@ -83,22 +64,6 @@ const GEMINI_CLI_USAGE_URL = const CODEWHISPERER_BASE_URL = process.env.OMNIROUTE_CODEWHISPERER_BASE_URL ?? "https://codewhisperer.us-east-1.amazonaws.com"; -// Antigravity API config (credentials from PROVIDERS via credential loader) -const ANTIGRAVITY_CONFIG = { - quotaApiUrls: getAntigravityFetchAvailableModelsUrls(), - loadProjectApiUrl: "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:loadCodeAssist", - tokenUrl: "https://oauth2.googleapis.com/token", - get clientId() { - return PROVIDERS.antigravity.clientId; - }, - get clientSecret() { - return PROVIDERS.antigravity.clientSecret; - }, - get userAgent() { - return antigravityUserAgent(); - }, -}; - // Codex (OpenAI) API config const CODEX_CONFIG = { usageUrl: "https://chatgpt.com/backend-api/wham/usage", @@ -1221,739 +1186,21 @@ function getGeminiCliPlanLabel(subscriptionInfo: unknown): string { } // ── Antigravity subscription info cache ────────────────────────────────────── -// Prevents duplicate loadCodeAssist calls within the same quota cycle. -// Key: truncated accessToken → { data, fetchedAt } -const _antigravitySubCache = new Map(); -const ANTIGRAVITY_CACHE_TTL_MS = 5 * 60 * 1000; // 5 minutes -const ANTIGRAVITY_MODELS_CACHE_TTL_MS = 60 * 1000; -const ANTIGRAVITY_CREDIT_PROBE_TTL_MS = 5 * 60 * 1000; -const _antigravityAvailableModelsCache = new Map(); -const _antigravityAvailableModelsInflight = new Map>(); -const _antigravityUserQuotaCache = new Map(); -const _antigravityUserQuotaInflight = new Map>(); -const _antigravityCreditProbeCache = new Map(); -const _antigravityCreditProbeInflight = new Map>(); - -// ── Proactive TTL purging for module-level caches ────────────────────────── -// All 4 data caches only evict on read (passive TTL). This interval proactively -// 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( +// ── Proactive TTL purging for the Gemini CLI subscription cache ──────────── +// Passive TTL evicts on read; this interval proactively purges stale entries so +// keys accessed once and never again don't leak memory. The Antigravity caches +// + their own purge timer were split out into ./usage/antigravity.ts (god-file +// decomposition), so each module now owns its caches and their cleanup. +const _geminiCacheCleanupTimer = 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; - -// `toClientAntigravityQuotaModelId` was an inline if-ladder here; it is now the single -// source of truth in open-sse/config/antigravityModelAliases.ts (imported above), shared -// with the provider-limits cache sanitizer. (#3821-review LEDGER-5) - -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"}`; -} - -async function fetchAntigravityAvailableModelsCached( - accessToken: string, - projectId?: string | null, - options: AntigravityUsageOptions = {} -): Promise { - if (!accessToken) throw new Error("Access token is required"); - - const cacheKey = buildAntigravityUsageCacheKey(accessToken, projectId); - const cached = _antigravityAvailableModelsCache.get(cacheKey); - if ( - !options.forceRefresh && - cached && - Date.now() - cached.fetchedAt < ANTIGRAVITY_MODELS_CACHE_TTL_MS - ) { - return cached.data; - } - - const inflight = _antigravityAvailableModelsInflight.get(cacheKey); - if (inflight) return inflight; - - const promise = (async () => { - let response: Response | null = null; - let lastError: Error | null = null; - - for (const quotaApiUrl of ANTIGRAVITY_CONFIG.quotaApiUrls) { - try { - response = await fetch(quotaApiUrl, { - method: "POST", - headers: getAntigravityHeaders("fetchAvailableModels", accessToken), - body: JSON.stringify(projectId ? { project: projectId } : {}), - signal: AbortSignal.timeout(10000), - }); - - if (response.ok || response.status === 401 || response.status === 403) { - break; - } - } catch (error) { - lastError = error as Error; - } - } - - if (!response) { - throw lastError || new Error("Antigravity API unavailable"); - } - - if (response.status === 403) { - return { __antigravityForbidden: true }; - } - - if (!response.ok) { - throw new Error(`Antigravity API error: ${response.status}`); - } - - const data = await response.json(); - _antigravityAvailableModelsCache.set(cacheKey, { data, fetchedAt: Date.now() }); - return data; - })().finally(() => { - _antigravityAvailableModelsInflight.delete(cacheKey); - }); - - _antigravityAvailableModelsInflight.set(cacheKey, promise); - return promise; -} - -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 ""; - const upper = tierId.toUpperCase(); - return mapCodeAssistTierIdToLabel(upper) ? upper : ""; -} - -function mapCodeAssistTierIdToLabel(tierId: string): string | null { - const upper = tierId.toUpperCase(); - if (upper.includes("ULTRA")) return "Ultra"; - if ( - upper.includes("PRO") || - upper.includes("PREMIUM") || - upper.includes("GOOGLE_ONE") || - upper.includes("ONE_AI") - ) - return "Pro"; - if (upper.includes("ENTERPRISE")) return "Enterprise"; - if (upper.includes("BUSINESS") || upper.includes("STANDARD")) return "Business"; - if (upper.includes("PLUS")) return "Plus"; - if (upper.includes("LITE") || upper.includes("LIGHT")) return "Lite"; - if (upper.includes("FREE") || upper.includes("INDIVIDUAL") || upper.includes("LEGACY")) - return "Free"; - return null; -} - -function mapSubscriptionTierStringToPlanLabel(tierText: string): string | null { - const upper = tierText.toUpperCase(); - if (upper.includes("ULTRA")) return "Ultra"; - if (upper.includes("PRO") || upper.includes("PREMIUM") || upper.includes("GOOGLE ONE")) - return "Pro"; - if (upper.includes("ENTERPRISE")) return "Enterprise"; - if (upper.includes("STANDARD") || upper.includes("BUSINESS")) return "Business"; - if (upper.includes("PLUS")) return "Plus"; - if (upper.includes("LITE")) return "Lite"; - if (upper.includes("INDIVIDUAL") || upper.includes("FREE")) return "Free"; - // Strip a trailing "(RESTRICTED)" marker. Match the fixed literal anywhere then - // trim, instead of /\s*\(RESTRICTED\)\s*$/ whose overlapping \s* runs backtrack - // polynomially on whitespace-heavy upstream input (js/polynomial-redos). - const normalizedId = upper.replace(/\(RESTRICTED\)/i, "").trim(); - if (normalizedId) { - const mapped = mapCodeAssistTierIdToLabel(normalizedId); - if (mapped) return mapped; - } - return null; -} - -function mapCodeAssistSubscriptionToPlanLabel(subscriptionInfo: unknown): string { - const subscription = toRecord(subscriptionInfo); - if (Object.keys(subscription).length === 0) return "Free"; - - const subscriptionTier = extractCodeAssistSubscriptionTier(subscriptionInfo); - if (subscriptionTier) { - const mapped = mapSubscriptionTierStringToPlanLabel(subscriptionTier); - if (mapped) return mapped; - if (subscriptionTier.toLowerCase() !== "free") { - return subscriptionTier.charAt(0).toUpperCase() + subscriptionTier.slice(1).toLowerCase(); - } - } - - const currentTier = toRecord(subscription.currentTier); - const tierName = String( - getFieldValue(currentTier, "name", "displayName") || - subscription.subscriptionType || - subscription.tier || - "" - ); - const mappedName = tierName ? mapSubscriptionTierStringToPlanLabel(tierName) : null; - if (mappedName) return mappedName; - - const tierId = extractCodeAssistTierId(subscription); - if (tierId) { - const mapped = mapCodeAssistTierIdToLabel(tierId); - if (mapped) return mapped; - } - if (currentTier.upgradeSubscriptionType) return "Free"; - if (tierName) return tierName.charAt(0).toUpperCase() + tierName.slice(1).toLowerCase(); - return "Free"; -} - -const KNOWN_ANTIGRAVITY_PLAN_LABELS = new Set([ - "Ultra", - "Pro", - "Enterprise", - "Business", - "Plus", - "Lite", -]); - -/** - * Map raw loadCodeAssist tier data to short display labels (Antigravity Manager parity). - */ -function getAntigravityPlanLabel(subscriptionInfo: unknown, fallbackInfo?: unknown): string { - const livePlan = mapCodeAssistSubscriptionToPlanLabel(subscriptionInfo); - const fallbackPlan = mapCodeAssistSubscriptionToPlanLabel(fallbackInfo); - - if (KNOWN_ANTIGRAVITY_PLAN_LABELS.has(livePlan)) return livePlan; - if (KNOWN_ANTIGRAVITY_PLAN_LABELS.has(fallbackPlan)) return fallbackPlan; - if (livePlan !== "Free") return livePlan; - return fallbackPlan !== "Free" ? fallbackPlan : livePlan; -} - -/** - * Proactive credit balance probe for Antigravity. - * - * Fires a minimal streamGenerateContent request with GOOGLE_ONE_AI credits enabled - * and maxOutputTokens=1 to extract the `remainingCredits` field from the SSE stream. - * This uses ~1 credit but lets us show the balance on the dashboard without waiting - * for a real user request. - * - * Returns the credit balance, or null if the probe failed. - */ -async function probeAntigravityCreditBalance( - accessToken: string, - accountId: string, - projectId?: string | null, - options: AntigravityUsageOptions = {}, - providerSpecificData: JsonRecord = {} -): Promise { - if (!accessToken) return null; - - const cacheKey = buildAntigravityUsageCacheKey(accessToken, projectId || accountId); - const cached = _antigravityCreditProbeCache.get(cacheKey); - if ( - !options.forceRefresh && - cached && - Date.now() - cached.fetchedAt < ANTIGRAVITY_CREDIT_PROBE_TTL_MS - ) { - return cached.data; - } - - const inflight = _antigravityCreditProbeInflight.get(cacheKey); - if (inflight) return inflight; - - const promise = probeAntigravityCreditBalanceUncached( - accessToken, - accountId, - projectId, - providerSpecificData - ) - .then( - (data) => { - _antigravityCreditProbeCache.set(cacheKey, { data, fetchedAt: Date.now() }); - return data; - }, - (error) => { - _antigravityCreditProbeCache.set(cacheKey, { data: null, fetchedAt: Date.now() }); - throw error; - } - ) - .finally(() => { - _antigravityCreditProbeInflight.delete(cacheKey); - }); - - _antigravityCreditProbeInflight.set(cacheKey, promise); - return promise; -} - -async function probeAntigravityCreditBalanceUncached( - accessToken: string, - accountId: string, - projectId?: string | null, - providerSpecificData: JsonRecord = {} -): Promise { - try { - if (!projectId) return null; - - // Try all base URLs (some accounts only work with specific endpoints) - for (const baseUrl of ANTIGRAVITY_BASE_URLS) { - const url = `${baseUrl}/v1internal:streamGenerateContent?alt=sse`; - - const sessionId = getAntigravitySessionId({ connectionId: accountId, projectId }); - const body = { - project: projectId, - model: "gemini-2-flash", - userAgent: "antigravity", - requestType: "agent", - requestId: generateAntigravityRequestId(), - enabledCreditTypes: ["GOOGLE_ONE_AI"], - request: { - model: "gemini-2-flash", - contents: [{ role: "user", parts: [{ text: "hi" }] }], - generationConfig: { maxOutputTokens: 1 }, - sessionId, - }, - }; - - const headers: Record = { - "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, - Accept: "text/event-stream", - }; - applyAntigravityClientProfileHeaders( - headers, - { connectionId: accountId, projectId, providerSpecificData }, - body - ); - - try { - const res = await fetch(url, { - method: "POST", - headers, - body: JSON.stringify(body), - signal: AbortSignal.timeout(10_000), - }); - - if (!res.ok) continue; - - // Read the full SSE response and scan for remainingCredits - const rawSSE = await res.text(); - const lines = rawSSE.split("\n"); - - for (const line of lines) { - const trimmed = line.trim(); - if (!trimmed.startsWith("data:")) continue; - const payload = trimmed.slice(5).trim(); - if (payload === "[DONE]") break; - try { - const parsed = JSON.parse(payload); - if (Array.isArray(parsed?.remainingCredits)) { - const googleCredit = parsed.remainingCredits.find( - (c: { creditType?: string }) => c?.creditType === "GOOGLE_ONE_AI" - ); - if (googleCredit) { - const balance = parseInt(googleCredit.creditAmount, 10); - if (!isNaN(balance)) { - updateAntigravityRemainingCredits(accountId, balance); - return balance; - } - } - } - } catch { - // Skip malformed SSE lines - } - } - } catch { - // Individual endpoint failure; try next - } - } - - return null; - } catch { - // Probe is best-effort — don't let it break the usage fetch - return null; - } -} - -/** - * 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, - connectionId?: string, - options: AntigravityUsageOptions = {} -) { - if (!accessToken) { - return { plan: "Free", message: "Antigravity access token not available." }; - } - - let subscriptionInfo: unknown = null; - try { - subscriptionInfo = await getAntigravitySubscriptionInfoCached( - accessToken, - providerSpecificData, - options - ); - const savedProjectId = - typeof providerSpecificData?.projectId === "string" && providerSpecificData.projectId.trim() - ? providerSpecificData.projectId.trim() - : null; - const subscriptionProject = toRecord(subscriptionInfo).cloudaicompanionProject; - const projectId = - savedProjectId || - connectionProjectId || - (typeof subscriptionProject === "string" - ? subscriptionProject - : typeof toRecord(subscriptionProject).id === "string" - ? (toRecord(subscriptionProject).id as string) - : null); - - // Derive accountId for credit balance cache. - // Must match executor key: credentials.connectionId - const accountId: string = connectionId || "unknown"; - - // Read cached credit balance (hydrated from DB on first access) - let creditBalance = getAntigravityRemainingCredits(accountId); - - // If no cached balance and credits mode is enabled, fire a minimal probe - const creditsMode = getCreditsMode(); - if ((options.forceRefresh || creditBalance === null) && creditsMode !== "off") { - creditBalance = await probeAntigravityCreditBalance( - accessToken, - accountId, - projectId, - options, - providerSpecificData || {} - ); - } - - 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 [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 || - !(provider === "agy" - ? isUserCallableAgyModelId(modelKey) - : isUserCallableAntigravityModelId(modelKey)) || - Object.keys(quotaInfo).length === 0 - ) { - continue; - } - - 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%". - // 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( - `[Antigravity] model ${modelKey} returned no remainingFraction — quota unknown` - ); - } - const remainingFraction = fractionReported ? Math.max(0, Math.min(1, rawFraction)) : 0; - // Models with no resetTime AND a reported full fraction are unlimited - // (e.g. tab-completion models). Unreported fraction is NEVER unlimited. - const isUnlimited = fractionReported && !resetAt && remainingFraction >= 1; - const remainingPercentage = remainingFraction * 100; - const QUOTA_NORMALIZED_BASE = 1000; - const total = QUOTA_NORMALIZED_BASE; - 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: isUnlimited ? 0 : Math.max(0, total - remaining), - total: isUnlimited ? 0 : total, - resetAt, - remainingPercentage: isUnlimited ? 100 : remainingFraction * 100, - unlimited: isUnlimited, - fractionReported: true, - quotaSource: "retrieveUserQuota", - }; - } - - return { - plan: getAntigravityPlanLabel(subscriptionInfo, providerSpecificData), - quotas: { - ...quotas, - ...(creditBalance !== null && { - credits: { - used: 0, - total: 0, - remaining: creditBalance, - unlimited: false, - resetAt: null, - }, - }), - }, - subscriptionInfo, - }; - } catch (error) { - return { - plan: getAntigravityPlanLabel(subscriptionInfo, providerSpecificData), - subscriptionInfo, - message: `Antigravity error: ${(error as Error).message}`, - }; - } -} - -/** - * Get Antigravity subscription info (cached, 5 min TTL) - * Prevents duplicate loadCodeAssist calls within the same quota cycle. - */ -async function getAntigravitySubscriptionInfoCached( - accessToken: string, - providerSpecificData?: JsonRecord, - options: AntigravityUsageOptions = {} -): Promise { - const profile = getAntigravityClientProfile({ providerSpecificData }); - const cacheKey = `${accessToken.substring(0, 16)}:${profile}`; - - if (options.forceRefresh) { - _antigravitySubCache.delete(cacheKey); - } else { - const cached = _antigravitySubCache.get(cacheKey); - if (cached && Date.now() - cached.fetchedAt < ANTIGRAVITY_CACHE_TTL_MS) { - return cached.data; - } - } - - const data = await getAntigravitySubscriptionInfo(accessToken, providerSpecificData); - if (data != null) { - _antigravitySubCache.set(cacheKey, { data, fetchedAt: Date.now() }); - } - return data; -} - -/** - * Get Antigravity subscription info using correct Antigravity headers. - * Must match the headers used in providers.js postExchange (not CLI headers). - */ -async function getAntigravitySubscriptionInfo( - accessToken: string, - providerSpecificData?: JsonRecord -): Promise { - try { - const profile = getAntigravityClientProfile({ providerSpecificData }); - const response = await fetch(ANTIGRAVITY_CONFIG.loadProjectApiUrl, { - method: "POST", - headers: - profile === "harness" - ? getAntigravityBootstrapHeaders(profile, accessToken) - : getAntigravityHeaders("loadCodeAssist", accessToken), - body: JSON.stringify({ metadata: getAntigravityLoadCodeAssistMetadata() }), - }); - - if (!response.ok) return null; - - return await response.json(); - } catch { - return null; - } -} +_geminiCacheCleanupTimer.unref?.(); // Don't prevent process exit /** * Claude Usage - Try to fetch from Anthropic API diff --git a/open-sse/services/usage/antigravity.ts b/open-sse/services/usage/antigravity.ts new file mode 100644 index 000000000000..02f98b010d8c --- /dev/null +++ b/open-sse/services/usage/antigravity.ts @@ -0,0 +1,794 @@ +/** + * usage/antigravity.ts — Antigravity (Gemini Code Assist) usage fetcher + quota helpers. + * + * Extracted from services/usage.ts (god-file decomposition): the full Antigravity family — + * local-usage fallback, code-assist tier/plan mapping, credit-balance probing, the user-quota + * + available-models fetchers (with their module-level caches), and getAntigravityUsage. The + * 4 data caches + their proactive TTL-purge setInterval move here as a self-contained unit + * (previously the purge timer was shared with the Gemini CLI cache in usage.ts; that timer was + * split so each module owns its own caches + cleanup). usage.ts imports getAntigravityUsage + * (dispatcher) + getAntigravityPlanLabel/mapCodeAssist* (__testing). Behavior-preserving move. + */ + +import { PROVIDERS } from "../../config/constants.ts"; +import { + getAntigravityFetchAvailableModelsUrls, + ANTIGRAVITY_BASE_URLS, +} from "../../config/antigravityUpstream.ts"; +import { + isUserCallableAntigravityModelId, + toClientAntigravityQuotaModelId, +} from "../../config/antigravityModelAliases.ts"; +import { isUserCallableAgyModelId } from "../../config/agyModels.ts"; +import { getDbInstance } from "@/lib/db/core"; +import { + applyAntigravityClientProfileHeaders, + getAntigravityBootstrapHeaders, + getAntigravityClientProfile, +} from "../antigravityClientProfile.ts"; +import { + antigravityUserAgent, + getAntigravityHeaders, + getAntigravityLoadCodeAssistMetadata, +} from "../antigravityHeaders.ts"; +import { + getAntigravityRemainingCredits, + updateAntigravityRemainingCredits, +} from "../../executors/antigravity.ts"; +import { getCreditsMode } from "../antigravityCredits.ts"; +import { generateAntigravityRequestId, getAntigravitySessionId } from "../antigravityIdentity.ts"; +import { + extractCodeAssistOnboardTierId, + extractCodeAssistSubscriptionTier, +} from "../codeAssistSubscription.ts"; +import { toRecord, toNumber, getFieldValue } from "./scalars.ts"; +import { type UsageQuota, parseResetTime } from "./quota.ts"; + +type JsonRecord = Record; +type SubscriptionCacheEntry = { + data: unknown; + fetchedAt: number; +}; + +const ANTIGRAVITY_CONFIG = { + quotaApiUrls: getAntigravityFetchAvailableModelsUrls(), + loadProjectApiUrl: "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:loadCodeAssist", + tokenUrl: "https://oauth2.googleapis.com/token", + get clientId() { + return PROVIDERS.antigravity.clientId; + }, + get clientSecret() { + return PROVIDERS.antigravity.clientSecret; + }, + get userAgent() { + return antigravityUserAgent(); + }, +}; + +const _antigravitySubCache = new Map(); +const ANTIGRAVITY_CACHE_TTL_MS = 5 * 60 * 1000; // 5 minutes +const ANTIGRAVITY_MODELS_CACHE_TTL_MS = 60 * 1000; +const ANTIGRAVITY_CREDIT_PROBE_TTL_MS = 5 * 60 * 1000; +const _antigravityAvailableModelsCache = new Map(); +const _antigravityAvailableModelsInflight = new Map>(); +const _antigravityUserQuotaCache = new Map(); +const _antigravityUserQuotaInflight = new Map>(); +const _antigravityCreditProbeCache = new Map(); +const _antigravityCreditProbeInflight = new Map>(); + +// ── Proactive TTL purging for the Antigravity module-level caches ────────── +// Split out of the shared usage.ts cleanup timer (god-file decomposition): this +// leaf owns its 4 data caches, so it owns their purge too. The 2 inflight Maps +// self-clean when their Promise settles, so they are NOT swept here. +const _antigravityCacheCleanupTimer = setInterval( + () => { + const now = Date.now(); + 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 +_antigravityCacheCleanupTimer.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; + +// `toClientAntigravityQuotaModelId` was an inline if-ladder here; it is now the single +// source of truth in open-sse/config/antigravityModelAliases.ts (imported above), shared +// with the provider-limits cache sanitizer. (#3821-review LEDGER-5) + +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"}`; +} + +async function fetchAntigravityAvailableModelsCached( + accessToken: string, + projectId?: string | null, + options: AntigravityUsageOptions = {} +): Promise { + if (!accessToken) throw new Error("Access token is required"); + + const cacheKey = buildAntigravityUsageCacheKey(accessToken, projectId); + const cached = _antigravityAvailableModelsCache.get(cacheKey); + if ( + !options.forceRefresh && + cached && + Date.now() - cached.fetchedAt < ANTIGRAVITY_MODELS_CACHE_TTL_MS + ) { + return cached.data; + } + + const inflight = _antigravityAvailableModelsInflight.get(cacheKey); + if (inflight) return inflight; + + const promise = (async () => { + let response: Response | null = null; + let lastError: Error | null = null; + + for (const quotaApiUrl of ANTIGRAVITY_CONFIG.quotaApiUrls) { + try { + response = await fetch(quotaApiUrl, { + method: "POST", + headers: getAntigravityHeaders("fetchAvailableModels", accessToken), + body: JSON.stringify(projectId ? { project: projectId } : {}), + signal: AbortSignal.timeout(10000), + }); + + if (response.ok || response.status === 401 || response.status === 403) { + break; + } + } catch (error) { + lastError = error as Error; + } + } + + if (!response) { + throw lastError || new Error("Antigravity API unavailable"); + } + + if (response.status === 403) { + return { __antigravityForbidden: true }; + } + + if (!response.ok) { + throw new Error(`Antigravity API error: ${response.status}`); + } + + const data = await response.json(); + _antigravityAvailableModelsCache.set(cacheKey, { data, fetchedAt: Date.now() }); + return data; + })().finally(() => { + _antigravityAvailableModelsInflight.delete(cacheKey); + }); + + _antigravityAvailableModelsInflight.set(cacheKey, promise); + return promise; +} + +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 ""; + const upper = tierId.toUpperCase(); + return mapCodeAssistTierIdToLabel(upper) ? upper : ""; +} + +export function mapCodeAssistTierIdToLabel(tierId: string): string | null { + const upper = tierId.toUpperCase(); + if (upper.includes("ULTRA")) return "Ultra"; + if ( + upper.includes("PRO") || + upper.includes("PREMIUM") || + upper.includes("GOOGLE_ONE") || + upper.includes("ONE_AI") + ) + return "Pro"; + if (upper.includes("ENTERPRISE")) return "Enterprise"; + if (upper.includes("BUSINESS") || upper.includes("STANDARD")) return "Business"; + if (upper.includes("PLUS")) return "Plus"; + if (upper.includes("LITE") || upper.includes("LIGHT")) return "Lite"; + if (upper.includes("FREE") || upper.includes("INDIVIDUAL") || upper.includes("LEGACY")) + return "Free"; + return null; +} + +export function mapSubscriptionTierStringToPlanLabel(tierText: string): string | null { + const upper = tierText.toUpperCase(); + if (upper.includes("ULTRA")) return "Ultra"; + if (upper.includes("PRO") || upper.includes("PREMIUM") || upper.includes("GOOGLE ONE")) + return "Pro"; + if (upper.includes("ENTERPRISE")) return "Enterprise"; + if (upper.includes("STANDARD") || upper.includes("BUSINESS")) return "Business"; + if (upper.includes("PLUS")) return "Plus"; + if (upper.includes("LITE")) return "Lite"; + if (upper.includes("INDIVIDUAL") || upper.includes("FREE")) return "Free"; + // Strip a trailing "(RESTRICTED)" marker. Match the fixed literal anywhere then + // trim, instead of /\s*\(RESTRICTED\)\s*$/ whose overlapping \s* runs backtrack + // polynomially on whitespace-heavy upstream input (js/polynomial-redos). + const normalizedId = upper.replace(/\(RESTRICTED\)/i, "").trim(); + if (normalizedId) { + const mapped = mapCodeAssistTierIdToLabel(normalizedId); + if (mapped) return mapped; + } + return null; +} + +export function mapCodeAssistSubscriptionToPlanLabel(subscriptionInfo: unknown): string { + const subscription = toRecord(subscriptionInfo); + if (Object.keys(subscription).length === 0) return "Free"; + + const subscriptionTier = extractCodeAssistSubscriptionTier(subscriptionInfo); + if (subscriptionTier) { + const mapped = mapSubscriptionTierStringToPlanLabel(subscriptionTier); + if (mapped) return mapped; + if (subscriptionTier.toLowerCase() !== "free") { + return subscriptionTier.charAt(0).toUpperCase() + subscriptionTier.slice(1).toLowerCase(); + } + } + + const currentTier = toRecord(subscription.currentTier); + const tierName = String( + getFieldValue(currentTier, "name", "displayName") || + subscription.subscriptionType || + subscription.tier || + "" + ); + const mappedName = tierName ? mapSubscriptionTierStringToPlanLabel(tierName) : null; + if (mappedName) return mappedName; + + const tierId = extractCodeAssistTierId(subscription); + if (tierId) { + const mapped = mapCodeAssistTierIdToLabel(tierId); + if (mapped) return mapped; + } + if (currentTier.upgradeSubscriptionType) return "Free"; + if (tierName) return tierName.charAt(0).toUpperCase() + tierName.slice(1).toLowerCase(); + return "Free"; +} + +const KNOWN_ANTIGRAVITY_PLAN_LABELS = new Set([ + "Ultra", + "Pro", + "Enterprise", + "Business", + "Plus", + "Lite", +]); + +/** + * Map raw loadCodeAssist tier data to short display labels (Antigravity Manager parity). + */ +export function getAntigravityPlanLabel(subscriptionInfo: unknown, fallbackInfo?: unknown): string { + const livePlan = mapCodeAssistSubscriptionToPlanLabel(subscriptionInfo); + const fallbackPlan = mapCodeAssistSubscriptionToPlanLabel(fallbackInfo); + + if (KNOWN_ANTIGRAVITY_PLAN_LABELS.has(livePlan)) return livePlan; + if (KNOWN_ANTIGRAVITY_PLAN_LABELS.has(fallbackPlan)) return fallbackPlan; + if (livePlan !== "Free") return livePlan; + return fallbackPlan !== "Free" ? fallbackPlan : livePlan; +} + +/** + * Proactive credit balance probe for Antigravity. + * + * Fires a minimal streamGenerateContent request with GOOGLE_ONE_AI credits enabled + * and maxOutputTokens=1 to extract the `remainingCredits` field from the SSE stream. + * This uses ~1 credit but lets us show the balance on the dashboard without waiting + * for a real user request. + * + * Returns the credit balance, or null if the probe failed. + */ +async function probeAntigravityCreditBalance( + accessToken: string, + accountId: string, + projectId?: string | null, + options: AntigravityUsageOptions = {}, + providerSpecificData: JsonRecord = {} +): Promise { + if (!accessToken) return null; + + const cacheKey = buildAntigravityUsageCacheKey(accessToken, projectId || accountId); + const cached = _antigravityCreditProbeCache.get(cacheKey); + if ( + !options.forceRefresh && + cached && + Date.now() - cached.fetchedAt < ANTIGRAVITY_CREDIT_PROBE_TTL_MS + ) { + return cached.data; + } + + const inflight = _antigravityCreditProbeInflight.get(cacheKey); + if (inflight) return inflight; + + const promise = probeAntigravityCreditBalanceUncached( + accessToken, + accountId, + projectId, + providerSpecificData + ) + .then( + (data) => { + _antigravityCreditProbeCache.set(cacheKey, { data, fetchedAt: Date.now() }); + return data; + }, + (error) => { + _antigravityCreditProbeCache.set(cacheKey, { data: null, fetchedAt: Date.now() }); + throw error; + } + ) + .finally(() => { + _antigravityCreditProbeInflight.delete(cacheKey); + }); + + _antigravityCreditProbeInflight.set(cacheKey, promise); + return promise; +} + +async function probeAntigravityCreditBalanceUncached( + accessToken: string, + accountId: string, + projectId?: string | null, + providerSpecificData: JsonRecord = {} +): Promise { + try { + if (!projectId) return null; + + // Try all base URLs (some accounts only work with specific endpoints) + for (const baseUrl of ANTIGRAVITY_BASE_URLS) { + const url = `${baseUrl}/v1internal:streamGenerateContent?alt=sse`; + + const sessionId = getAntigravitySessionId({ connectionId: accountId, projectId }); + const body = { + project: projectId, + model: "gemini-2-flash", + userAgent: "antigravity", + requestType: "agent", + requestId: generateAntigravityRequestId(), + enabledCreditTypes: ["GOOGLE_ONE_AI"], + request: { + model: "gemini-2-flash", + contents: [{ role: "user", parts: [{ text: "hi" }] }], + generationConfig: { maxOutputTokens: 1 }, + sessionId, + }, + }; + + const headers: Record = { + "Content-Type": "application/json", + Authorization: `Bearer ${accessToken}`, + Accept: "text/event-stream", + }; + applyAntigravityClientProfileHeaders( + headers, + { connectionId: accountId, projectId, providerSpecificData }, + body + ); + + try { + const res = await fetch(url, { + method: "POST", + headers, + body: JSON.stringify(body), + signal: AbortSignal.timeout(10_000), + }); + + if (!res.ok) continue; + + // Read the full SSE response and scan for remainingCredits + const rawSSE = await res.text(); + const lines = rawSSE.split("\n"); + + for (const line of lines) { + const trimmed = line.trim(); + if (!trimmed.startsWith("data:")) continue; + const payload = trimmed.slice(5).trim(); + if (payload === "[DONE]") break; + try { + const parsed = JSON.parse(payload); + if (Array.isArray(parsed?.remainingCredits)) { + const googleCredit = parsed.remainingCredits.find( + (c: { creditType?: string }) => c?.creditType === "GOOGLE_ONE_AI" + ); + if (googleCredit) { + const balance = parseInt(googleCredit.creditAmount, 10); + if (!isNaN(balance)) { + updateAntigravityRemainingCredits(accountId, balance); + return balance; + } + } + } + } catch { + // Skip malformed SSE lines + } + } + } catch { + // Individual endpoint failure; try next + } + } + + return null; + } catch { + // Probe is best-effort — don't let it break the usage fetch + return null; + } +} + +/** + * 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). + */ +export async function getAntigravityUsage( + provider: "antigravity" | "agy", + accessToken?: string, + providerSpecificData?: JsonRecord, + connectionProjectId?: string, + connectionId?: string, + options: AntigravityUsageOptions = {} +) { + if (!accessToken) { + return { plan: "Free", message: "Antigravity access token not available." }; + } + + let subscriptionInfo: unknown = null; + try { + subscriptionInfo = await getAntigravitySubscriptionInfoCached( + accessToken, + providerSpecificData, + options + ); + const savedProjectId = + typeof providerSpecificData?.projectId === "string" && providerSpecificData.projectId.trim() + ? providerSpecificData.projectId.trim() + : null; + const subscriptionProject = toRecord(subscriptionInfo).cloudaicompanionProject; + const projectId = + savedProjectId || + connectionProjectId || + (typeof subscriptionProject === "string" + ? subscriptionProject + : typeof toRecord(subscriptionProject).id === "string" + ? (toRecord(subscriptionProject).id as string) + : null); + + // Derive accountId for credit balance cache. + // Must match executor key: credentials.connectionId + const accountId: string = connectionId || "unknown"; + + // Read cached credit balance (hydrated from DB on first access) + let creditBalance = getAntigravityRemainingCredits(accountId); + + // If no cached balance and credits mode is enabled, fire a minimal probe + const creditsMode = getCreditsMode(); + if ((options.forceRefresh || creditBalance === null) && creditsMode !== "off") { + creditBalance = await probeAntigravityCreditBalance( + accessToken, + accountId, + projectId, + options, + providerSpecificData || {} + ); + } + + 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 [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 || + !(provider === "agy" + ? isUserCallableAgyModelId(modelKey) + : isUserCallableAntigravityModelId(modelKey)) || + Object.keys(quotaInfo).length === 0 + ) { + continue; + } + + 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%". + // 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( + `[Antigravity] model ${modelKey} returned no remainingFraction — quota unknown` + ); + } + const remainingFraction = fractionReported ? Math.max(0, Math.min(1, rawFraction)) : 0; + // Models with no resetTime AND a reported full fraction are unlimited + // (e.g. tab-completion models). Unreported fraction is NEVER unlimited. + const isUnlimited = fractionReported && !resetAt && remainingFraction >= 1; + const remainingPercentage = remainingFraction * 100; + const QUOTA_NORMALIZED_BASE = 1000; + const total = QUOTA_NORMALIZED_BASE; + 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: isUnlimited ? 0 : Math.max(0, total - remaining), + total: isUnlimited ? 0 : total, + resetAt, + remainingPercentage: isUnlimited ? 100 : remainingFraction * 100, + unlimited: isUnlimited, + fractionReported: true, + quotaSource: "retrieveUserQuota", + }; + } + + return { + plan: getAntigravityPlanLabel(subscriptionInfo, providerSpecificData), + quotas: { + ...quotas, + ...(creditBalance !== null && { + credits: { + used: 0, + total: 0, + remaining: creditBalance, + unlimited: false, + resetAt: null, + }, + }), + }, + subscriptionInfo, + }; + } catch (error) { + return { + plan: getAntigravityPlanLabel(subscriptionInfo, providerSpecificData), + subscriptionInfo, + message: `Antigravity error: ${(error as Error).message}`, + }; + } +} + +/** + * Get Antigravity subscription info (cached, 5 min TTL) + * Prevents duplicate loadCodeAssist calls within the same quota cycle. + */ +async function getAntigravitySubscriptionInfoCached( + accessToken: string, + providerSpecificData?: JsonRecord, + options: AntigravityUsageOptions = {} +): Promise { + const profile = getAntigravityClientProfile({ providerSpecificData }); + const cacheKey = `${accessToken.substring(0, 16)}:${profile}`; + + if (options.forceRefresh) { + _antigravitySubCache.delete(cacheKey); + } else { + const cached = _antigravitySubCache.get(cacheKey); + if (cached && Date.now() - cached.fetchedAt < ANTIGRAVITY_CACHE_TTL_MS) { + return cached.data; + } + } + + const data = await getAntigravitySubscriptionInfo(accessToken, providerSpecificData); + if (data != null) { + _antigravitySubCache.set(cacheKey, { data, fetchedAt: Date.now() }); + } + return data; +} + +/** + * Get Antigravity subscription info using correct Antigravity headers. + * Must match the headers used in providers.js postExchange (not CLI headers). + */ +async function getAntigravitySubscriptionInfo( + accessToken: string, + providerSpecificData?: JsonRecord +): Promise { + try { + const profile = getAntigravityClientProfile({ providerSpecificData }); + const response = await fetch(ANTIGRAVITY_CONFIG.loadProjectApiUrl, { + method: "POST", + headers: + profile === "harness" + ? getAntigravityBootstrapHeaders(profile, accessToken) + : getAntigravityHeaders("loadCodeAssist", accessToken), + body: JSON.stringify({ metadata: getAntigravityLoadCodeAssistMetadata() }), + }); + + if (!response.ok) return null; + + return await response.json(); + } catch { + return null; + } +} diff --git a/tests/unit/usage-antigravity-family-split.test.ts b/tests/unit/usage-antigravity-family-split.test.ts new file mode 100644 index 000000000000..f1d3a5887a6f --- /dev/null +++ b/tests/unit/usage-antigravity-family-split.test.ts @@ -0,0 +1,45 @@ +// Characterization of the services/usage.ts Antigravity-family split (god-file decomposition): the +// full Antigravity (Gemini Code Assist) usage family — code-assist tier/plan mapping, credit probing, +// the user-quota/available-models fetchers + their module-level caches, and getAntigravityUsage — +// moved into services/usage/antigravity.ts. The 4 data caches + their proactive TTL-purge setInterval +// moved as a unit (the previously-shared purge timer was split so each module owns its own cleanup). +// Behavior-preserving move; the fetcher + quota math stay covered by usage-utils / usage-service- +// hardening / code-assist-subscription (which exercise them via usage.ts __testing). +import { test } from "node:test"; +import assert from "node:assert/strict"; + +const A = await import("../../open-sse/services/usage/antigravity.ts"); +const HOST = await import("../../open-sse/services/usage.ts"); + +test("leaf exports the host-facing Antigravity helpers", () => { + for (const name of [ + "getAntigravityUsage", + "getAntigravityPlanLabel", + "mapCodeAssistSubscriptionToPlanLabel", + "mapCodeAssistTierIdToLabel", + "mapSubscriptionTierStringToPlanLabel", + ]) { + assert.equal(typeof (A as Record)[name], "function", `missing ${name}`); + } +}); + +test("host __testing re-exports the same Antigravity function identities", () => { + const t = (HOST as Record>).__testing; + assert.equal(t.getAntigravityPlanLabel, A.getAntigravityPlanLabel); + assert.equal(t.mapCodeAssistSubscriptionToPlanLabel, A.mapCodeAssistSubscriptionToPlanLabel); + assert.equal(t.mapCodeAssistTierIdToLabel, A.mapCodeAssistTierIdToLabel); + assert.equal(t.mapSubscriptionTierStringToPlanLabel, A.mapSubscriptionTierStringToPlanLabel); +}); + +test("mapCodeAssistTierIdToLabel maps known tier ids, null otherwise", () => { + // unknown tier id → null (the label mapping only knows specific ids) + assert.equal(A.mapCodeAssistTierIdToLabel("totally-unknown-tier"), null); + // shape: always string|null + const out = A.mapCodeAssistTierIdToLabel("free-tier"); + assert.ok(out === null || typeof out === "string"); +}); + +test("getAntigravityPlanLabel returns a string for arbitrary input (never throws)", () => { + assert.equal(typeof A.getAntigravityPlanLabel(null), "string"); + assert.equal(typeof A.getAntigravityPlanLabel({ tier: "x" }, { other: 1 }), "string"); +});