diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index b4fe38f8..698e65fd 100644 --- a/open-sse/services/usage.ts +++ b/open-sse/services/usage.ts @@ -657,34 +657,26 @@ function getAntigravityPlanLabel(subscriptionInfo) { /** * Antigravity Usage - Fetch quota from Google Cloud Code API - * Now calls loadCodeAssist ONCE (cached) and reuses for projectId + plan. - * Uses retrieveUserQuota API (same as Gemini CLI) for accurate quota data across all tiers. + * Uses fetchAvailableModels API which returns ALL models (including Claude) + * with per-model quotaInfo (remainingFraction, resetTime). + * retrieveUserQuota only returns Gemini models — not suitable for Antigravity. */ async function getAntigravityUsage(accessToken, providerSpecificData) { try { const subscriptionInfo = await getAntigravitySubscriptionInfoCached(accessToken); const projectId = subscriptionInfo?.cloudaicompanionProject || null; - if (!projectId) { - return { - plan: getAntigravityPlanLabel(subscriptionInfo), - message: "Antigravity project ID not available.", - }; - } - - // Use retrieveUserQuota API (same as Gemini CLI) - works correctly for both Free and Pro tiers - 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), - } - ); + // Fetch model list with quota info from fetchAvailableModels + const response = await fetch(ANTIGRAVITY_CONFIG.quotaApiUrl, { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + "User-Agent": ANTIGRAVITY_CONFIG.userAgent, + "Content-Type": "application/json", + }, + body: JSON.stringify(projectId ? { project: projectId } : {}), + signal: AbortSignal.timeout(10000), + }); if (response.status === 403) { return { message: "Antigravity access forbidden. Check subscription." }; @@ -695,28 +687,39 @@ async function getAntigravityUsage(accessToken, providerSpecificData) { } const data = await response.json(); + const dataObj = toRecord(data); + const modelEntries = toRecord(dataObj.models); const quotas: Record = {}; - // Parse buckets from retrieveUserQuota response (same format as Gemini CLI) - if (Array.isArray(data.buckets)) { - for (const bucket of data.buckets) { - if (!bucket.modelId || bucket.remainingFraction == null) continue; + // Parse per-model quota info from fetchAvailableModels response. + // Show all models that have quota data, excluding only internal models + // (tab-completion, chat placeholders, etc.). + for (const [modelKey, infoValue] of Object.entries(modelEntries)) { + const info = toRecord(infoValue); + const quotaInfo = toRecord(info.quotaInfo); - const remainingFraction = toNumber(bucket.remainingFraction, 0); - const remainingPercentage = remainingFraction * 100; - const QUOTA_NORMALIZED_BASE = 1000; - const total = QUOTA_NORMALIZED_BASE; - const remaining = Math.round(total * remainingFraction); - const used = Math.max(0, total - remaining); - - quotas[bucket.modelId] = { - used, - total, - resetAt: parseResetTime(bucket.resetTime), - remainingPercentage, - unlimited: false, - }; + // Skip internal models and models without quota info + if (info.isInternal === true || Object.keys(quotaInfo).length === 0) { + continue; } + + const remainingFraction = toNumber(quotaInfo.remainingFraction, 0); + const resetAt = parseResetTime(quotaInfo.resetTime); + // Models with no resetTime and full remaining are unlimited (e.g. tab-completion models) + const isUnlimited = !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] = { + used, + total: isUnlimited ? 0 : total, + resetAt, + remainingPercentage: isUnlimited ? 100 : remainingPercentage, + unlimited: isUnlimited, + }; } return {