From a762d8dda0fa1c9c3d0d5f8bfa5eae56a0e2bb69 Mon Sep 17 00:00:00 2001 From: Tam Nhu Tran Date: Fri, 3 Apr 2026 19:43:17 -0400 Subject: [PATCH] feat(cliproxy): enrich gemini quota metadata --- src/cliproxy/gemini-cli-quota-normalizer.ts | 169 ++++++++++ src/cliproxy/quota-fetcher-gemini-cli.ts | 289 +++++++++++------- src/cliproxy/quota-types.ts | 8 + src/commands/cliproxy/quota-subcommand.ts | 12 +- .../cliproxy/quota-fetcher-gemini-cli.test.ts | 151 ++++++++- 5 files changed, 511 insertions(+), 118 deletions(-) create mode 100644 src/cliproxy/gemini-cli-quota-normalizer.ts diff --git a/src/cliproxy/gemini-cli-quota-normalizer.ts b/src/cliproxy/gemini-cli-quota-normalizer.ts new file mode 100644 index 00000000..1d5425b9 --- /dev/null +++ b/src/cliproxy/gemini-cli-quota-normalizer.ts @@ -0,0 +1,169 @@ +import type { GeminiCliBucket } from './quota-types'; + +export interface GeminiCliParsedBucket { + modelId: string; + tokenType: string | null; + remainingFraction: number | null; + remainingAmount: number | null; + resetTime: string | null; +} + +interface GeminiCliQuotaGroupDefinition { + id: string; + label: string; + preferredModelId?: string; + modelIds: string[]; +} + +const GEMINI_CLI_QUOTA_GROUPS: GeminiCliQuotaGroupDefinition[] = [ + { + id: 'gemini-flash-lite-series', + label: 'Gemini Flash Lite Series', + preferredModelId: 'gemini-2.5-flash-lite', + modelIds: ['gemini-2.5-flash-lite'], + }, + { + id: 'gemini-flash-series', + label: 'Gemini Flash Series', + preferredModelId: 'gemini-3-flash-preview', + modelIds: ['gemini-3-flash-preview', 'gemini-3.1-flash-preview', 'gemini-2.5-flash'], + }, + { + id: 'gemini-pro-series', + label: 'Gemini Pro Series', + preferredModelId: 'gemini-3.1-pro-preview', + modelIds: ['gemini-3.1-pro-preview', 'gemini-3-pro-preview', 'gemini-2.5-pro'], + }, +]; + +const GEMINI_CLI_GROUP_ORDER = new Map( + GEMINI_CLI_QUOTA_GROUPS.map((group, index) => [group.id, index] as const) +); + +const GEMINI_CLI_GROUP_LOOKUP = new Map( + GEMINI_CLI_QUOTA_GROUPS.flatMap((group) => + group.modelIds.map((modelId) => [modelId, group] as const) + ) +); + +const GEMINI_CLI_IGNORED_MODEL_PREFIXES = ['gemini-2.0-flash']; + +type GeminiCliQuotaBucketGroup = { + id: string; + label: string; + tokenType: string | null; + modelIds: string[]; + preferredModelId?: string; + preferredBucket?: GeminiCliParsedBucket; + fallbackRemainingFraction: number | null; + fallbackRemainingAmount: number | null; + fallbackResetTime: string | null; +}; + +function isIgnoredGeminiCliModel(modelId: string): boolean { + return GEMINI_CLI_IGNORED_MODEL_PREFIXES.some((prefix) => modelId.startsWith(prefix)); +} + +function pickEarlierResetTime(current: string | null, next: string | null): string | null { + if (!current) return next; + if (!next) return current; + const currentTime = new Date(current).getTime(); + const nextTime = new Date(next).getTime(); + if (Number.isNaN(currentTime)) return next; + if (Number.isNaN(nextTime)) return current; + return currentTime <= nextTime ? current : next; +} + +function minNullableNumber(current: number | null, next: number | null): number | null { + if (current === null) return next; + if (next === null) return current; + return Math.min(current, next); +} + +function clampQuotaFraction(value: number | null): number { + if (value === null) return 0; + return Math.max(0, Math.min(1, value)); +} + +function getGroupOrder(bucket: GeminiCliQuotaBucketGroup): number { + return GEMINI_CLI_GROUP_ORDER.get(bucket.id) ?? Number.MAX_SAFE_INTEGER; +} + +export function buildGeminiCliBucketsFromParsedBuckets( + buckets: GeminiCliParsedBucket[] +): GeminiCliBucket[] { + if (buckets.length === 0) return []; + + const grouped = new Map(); + + for (const bucket of buckets) { + if (isIgnoredGeminiCliModel(bucket.modelId)) continue; + + const group = GEMINI_CLI_GROUP_LOOKUP.get(bucket.modelId); + const groupId = group?.id ?? bucket.modelId; + const label = group?.label ?? bucket.modelId; + const tokenKey = bucket.tokenType ?? ''; + const cacheKey = `${groupId}::${tokenKey || 'combined'}`; + const existing = grouped.get(cacheKey); + + if (!existing) { + const preferredModelId = group?.preferredModelId; + grouped.set(cacheKey, { + id: groupId, + label, + tokenType: bucket.tokenType, + modelIds: [bucket.modelId], + preferredModelId, + preferredBucket: + preferredModelId && bucket.modelId === preferredModelId ? bucket : undefined, + fallbackRemainingFraction: bucket.remainingFraction, + fallbackRemainingAmount: bucket.remainingAmount, + fallbackResetTime: bucket.resetTime, + }); + continue; + } + + existing.modelIds.push(bucket.modelId); + existing.fallbackRemainingFraction = minNullableNumber( + existing.fallbackRemainingFraction, + bucket.remainingFraction + ); + existing.fallbackRemainingAmount = minNullableNumber( + existing.fallbackRemainingAmount, + bucket.remainingAmount + ); + existing.fallbackResetTime = pickEarlierResetTime(existing.fallbackResetTime, bucket.resetTime); + + if (existing.preferredModelId && bucket.modelId === existing.preferredModelId) { + existing.preferredBucket = bucket; + } + } + + return Array.from(grouped.values()) + .sort((a, b) => { + const orderDiff = getGroupOrder(a) - getGroupOrder(b); + if (orderDiff !== 0) return orderDiff; + const tokenDiff = (a.tokenType ?? '').localeCompare(b.tokenType ?? ''); + if (tokenDiff !== 0) return tokenDiff; + return a.label.localeCompare(b.label); + }) + .map((bucket) => { + const preferred = bucket.preferredBucket; + const remainingFraction = clampQuotaFraction( + preferred?.remainingFraction ?? bucket.fallbackRemainingFraction + ); + const remainingAmount = preferred?.remainingAmount ?? bucket.fallbackRemainingAmount; + const resetTime = preferred?.resetTime ?? bucket.fallbackResetTime; + + return { + id: `${bucket.id}::${bucket.tokenType || 'combined'}`, + label: bucket.label, + tokenType: bucket.tokenType, + remainingFraction, + remainingPercent: Math.round(remainingFraction * 100), + remainingAmount, + resetTime, + modelIds: Array.from(new Set(bucket.modelIds)), + }; + }); +} diff --git a/src/cliproxy/quota-fetcher-gemini-cli.ts b/src/cliproxy/quota-fetcher-gemini-cli.ts index c8034ebc..33aaec41 100644 --- a/src/cliproxy/quota-fetcher-gemini-cli.ts +++ b/src/cliproxy/quota-fetcher-gemini-cli.ts @@ -11,43 +11,28 @@ import { getAuthDir } from './config-generator'; import { getProviderAccounts, getPausedDir } from './account-manager'; import { sanitizeEmail, isTokenExpired } from './auth-utils'; import { refreshGeminiToken } from './auth/gemini-token-refresh'; +import { + buildGeminiCliBucketsFromParsedBuckets, + type GeminiCliParsedBucket, +} from './gemini-cli-quota-normalizer'; import type { GeminiCliQuotaResult, GeminiCliBucket } from './quota-types'; /** Google Cloud Code API endpoints */ const GEMINI_CLI_API_BASE = 'https://cloudcode-pa.googleapis.com'; const GEMINI_CLI_API_VERSION = 'v1internal'; +const GEMINI_CLI_QUOTA_URL = `${GEMINI_CLI_API_BASE}/${GEMINI_CLI_API_VERSION}:retrieveUserQuota`; +const GEMINI_CLI_CODE_ASSIST_URL = `${GEMINI_CLI_API_BASE}/${GEMINI_CLI_API_VERSION}:loadCodeAssist`; const GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH = 320; const GEMINI_CLI_ERROR_DETAIL_TRUNCATION_SUFFIX = '...[truncated]'; - -/** - * Model groups for quota consolidation. - * Update when Google releases new Gemini models to include them in quota display. - */ -const GEMINI_CLI_GROUPS: Record< - string, - { - label: string; - models: string[]; - } -> = { - 'gemini-flash-series': { - label: 'Gemini Flash Series', - models: [ - 'gemini-3-flash-preview', - 'gemini-3.1-flash-preview', - 'gemini-2.5-flash', - 'gemini-2.5-flash-lite', - ], - }, - 'gemini-pro-series': { - label: 'Gemini Pro Series', - models: ['gemini-3-pro-preview', 'gemini-3.1-pro-preview', 'gemini-2.5-pro'], - }, +const GEMINI_CLI_G1_CREDIT_TYPE = 'GOOGLE_ONE_AI'; +const GEMINI_CLI_TIER_LABELS: Record = { + 'free-tier': 'Free', + 'legacy-tier': 'Legacy', + 'standard-tier': 'Standard', + 'g1-pro-tier': 'Pro', + 'g1-ultra-tier': 'Ultra', }; -/** Models to ignore in quota display (deprecated) */ -const IGNORED_MODEL_PREFIXES = ['gemini-2.0-flash']; - /** Auth data extracted from Gemini CLI auth file */ interface GeminiCliAuthData { accessToken: string; @@ -75,12 +60,38 @@ interface GeminiCliQuotaResponse { buckets?: RawGeminiCliBucket[]; } +interface GeminiCliCredits { + creditType?: string; + credit_type?: string; + creditAmount?: string | number; + credit_amount?: string | number; +} + +interface GeminiCliUserTier { + id?: string; + availableCredits?: GeminiCliCredits[]; + available_credits?: GeminiCliCredits[]; +} + +interface GeminiCliCodeAssistResponse { + currentTier?: GeminiCliUserTier | null; + current_tier?: GeminiCliUserTier | null; + paidTier?: GeminiCliUserTier | null; + paid_tier?: GeminiCliUserTier | null; +} + interface ParsedGeminiCliErrorBody { errorCode?: string; errorDetail?: string; message?: string; } +interface GeminiCliSupplementaryInfo { + tierLabel: string | null; + tierId: string | null; + creditBalance: number | null; +} + /** * Extract project ID from account field * Input: "user@example.com (cloudaicompanion-abc-123)" @@ -234,23 +245,112 @@ function readGeminiCliAuthData(accountId: string): GeminiCliAuthData | null { return null; } -/** - * Find which group a model belongs to - */ -function findModelGroup(modelId: string): { groupId: string; label: string } | null { - for (const [groupId, group] of Object.entries(GEMINI_CLI_GROUPS)) { - if (group.models.includes(modelId)) { - return { groupId, label: group.label }; +function normalizeStringValue(value: unknown): string | null { + return typeof value === 'string' && value.trim().length > 0 ? value.trim() : null; +} + +function normalizeNumberValue(value: unknown): number | null { + if (typeof value === 'number' && Number.isFinite(value)) { + return value; + } + if (typeof value === 'string' && value.trim().length > 0) { + const parsed = Number(value); + if (Number.isFinite(parsed)) { + return parsed; } } return null; } -/** - * Check if model should be ignored - */ -function shouldIgnoreModel(modelId: string): boolean { - return IGNORED_MODEL_PREFIXES.some((prefix) => modelId.startsWith(prefix)); +function resolveGeminiCliTierId(payload: GeminiCliCodeAssistResponse | null): string | null { + if (!payload) return null; + const currentTier = payload.currentTier ?? payload.current_tier; + const paidTier = payload.paidTier ?? payload.paid_tier; + const rawId = normalizeStringValue(paidTier?.id) ?? normalizeStringValue(currentTier?.id); + return rawId ? rawId.toLowerCase() : null; +} + +function resolveGeminiCliTierLabel(payload: GeminiCliCodeAssistResponse | null): string | null { + const tierId = resolveGeminiCliTierId(payload); + if (!tierId) return null; + return GEMINI_CLI_TIER_LABELS[tierId] ?? tierId; +} + +function resolveGeminiCliCreditBalance(payload: GeminiCliCodeAssistResponse | null): number | null { + if (!payload) return null; + + const paidTier = payload.paidTier ?? payload.paid_tier; + const currentTier = payload.currentTier ?? payload.current_tier; + const tier = paidTier ?? currentTier; + if (!tier) return null; + + const credits = tier.availableCredits ?? tier.available_credits ?? []; + let total = 0; + let found = false; + for (const credit of credits) { + const creditType = normalizeStringValue(credit.creditType ?? credit.credit_type); + if (creditType !== GEMINI_CLI_G1_CREDIT_TYPE) continue; + + const amount = normalizeNumberValue(credit.creditAmount ?? credit.credit_amount); + if (amount !== null) { + total += amount; + found = true; + } + } + + return found ? total : null; +} + +async function fetchGeminiCliSupplementary( + accessToken: string, + projectId: string, + verbose: boolean +): Promise { + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), 5000); + + try { + const response = await fetch(GEMINI_CLI_CODE_ASSIST_URL, { + method: 'POST', + signal: controller.signal, + headers: { + Authorization: `Bearer ${accessToken}`, + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + cloudaicompanionProject: projectId, + metadata: { + ideType: 'IDE_UNSPECIFIED', + platform: 'PLATFORM_UNSPECIFIED', + pluginType: 'GEMINI', + duetProject: projectId, + }, + }), + }); + + clearTimeout(timeoutId); + + if (!response.ok) { + if (verbose) { + console.error(`[i] Gemini CLI supplementary metadata unavailable: HTTP ${response.status}`); + } + return { tierLabel: null, tierId: null, creditBalance: null }; + } + + const payload = (await response.json()) as GeminiCliCodeAssistResponse; + return { + tierLabel: resolveGeminiCliTierLabel(payload), + tierId: resolveGeminiCliTierId(payload), + creditBalance: resolveGeminiCliCreditBalance(payload), + }; + } catch (error) { + clearTimeout(timeoutId); + if (verbose) { + const message = error instanceof Error ? error.message : 'Unknown error'; + console.error(`[i] Gemini CLI supplementary metadata skipped: ${message}`); + } + return { tierLabel: null, tierId: null, creditBalance: null }; + } } function buildGeminiCliFailureResult( @@ -271,6 +371,9 @@ function buildGeminiCliFailureResult( success: false, buckets: [], projectId, + tierLabel: null, + tierId: null, + creditBalance: null, lastUpdated: Date.now(), accountId, error: options.error, @@ -480,81 +583,38 @@ function buildGeminiCliHttpFailureResult( * Groups buckets by model series and token type */ function buildGeminiCliBuckets(rawBuckets: RawGeminiCliBucket[]): GeminiCliBucket[] { - // Group buckets by groupId::tokenType - const grouped = new Map< - string, - { - label: string; - tokenType: string | null; - remainingFraction: number; - resetTime: string | null; - modelIds: string[]; - } - >(); + const parsedBuckets = rawBuckets + .map((bucket): GeminiCliParsedBucket | null => { + const modelId = normalizeStringValue(bucket.model_id ?? bucket.modelId); + if (!modelId) return null; - for (const bucket of rawBuckets) { - const modelId = bucket.model_id || bucket.modelId || ''; - if (!modelId) continue; + const tokenType = normalizeStringValue(bucket.token_type ?? bucket.tokenType); + const remainingFractionRaw = normalizeNumberValue( + bucket.remaining_fraction ?? bucket.remainingFraction + ); + const remainingAmount = normalizeNumberValue( + bucket.remaining_amount ?? bucket.remainingAmount + ); + const resetTime = normalizeStringValue(bucket.reset_time ?? bucket.resetTime); - // Skip ignored models - if (shouldIgnoreModel(modelId)) continue; - - const tokenType = bucket.token_type ?? bucket.tokenType ?? null; - // Clamp remainingFraction to [0, 1] range - const rawRemainingFraction = bucket.remaining_fraction ?? bucket.remainingFraction ?? 1; - const remainingFraction = Math.max(0, Math.min(1, rawRemainingFraction)); - const resetTime = bucket.reset_time ?? bucket.resetTime ?? null; - - // Find group for this model - const group = findModelGroup(modelId); - const groupId = group?.groupId || 'other'; - const label = group?.label || 'Other Models'; - - // Create compound key for grouping - const key = `${groupId}::${tokenType || 'combined'}`; - - const existing = grouped.get(key); - if (existing) { - // Merge: take the minimum remaining fraction (most limiting) - existing.remainingFraction = Math.min(existing.remainingFraction, remainingFraction); - // Keep earliest reset time if available - if (resetTime && (!existing.resetTime || resetTime < existing.resetTime)) { - existing.resetTime = resetTime; + let fallbackFraction: number | null = null; + if (remainingAmount !== null) { + fallbackFraction = remainingAmount <= 0 ? 0 : null; + } else if (resetTime) { + fallbackFraction = 0; } - existing.modelIds.push(modelId); - } else { - grouped.set(key, { - label, + + return { + modelId, tokenType, - remainingFraction, + remainingFraction: remainingFractionRaw ?? fallbackFraction ?? 1, + remainingAmount, resetTime, - modelIds: [modelId], - }); - } - } + }; + }) + .filter((bucket): bucket is GeminiCliParsedBucket => bucket !== null); - // Convert to array - const buckets: GeminiCliBucket[] = []; - for (const [key, data] of grouped.entries()) { - buckets.push({ - id: key, - label: data.label, - tokenType: data.tokenType, - remainingFraction: data.remainingFraction, - remainingPercent: Math.round(data.remainingFraction * 100), - resetTime: data.resetTime, - modelIds: data.modelIds, - }); - } - - // Sort by label then token type - buckets.sort((a, b) => { - const labelCompare = a.label.localeCompare(b.label); - if (labelCompare !== 0) return labelCompare; - return (a.tokenType || '').localeCompare(b.tokenType || ''); - }); - - return buckets; + return buildGeminiCliBucketsFromParsedBuckets(parsedBuckets); } /** @@ -577,12 +637,11 @@ async function fetchWithAuthData( }); } - const url = `${GEMINI_CLI_API_BASE}/${GEMINI_CLI_API_VERSION}:retrieveUserQuota`; const controller = new AbortController(); const timeoutId = setTimeout(() => controller.abort(), 5000); try { - const response = await fetch(url, { + const response = await fetch(GEMINI_CLI_QUOTA_URL, { method: 'POST', signal: controller.signal, headers: { @@ -609,6 +668,11 @@ async function fetchWithAuthData( const data = (await response.json()) as GeminiCliQuotaResponse; const rawBuckets = data.buckets || []; const buckets = buildGeminiCliBuckets(rawBuckets); + const supplementary = await fetchGeminiCliSupplementary( + authData.accessToken, + authData.projectId, + verbose + ); if (verbose) console.error(`[i] Gemini CLI buckets found: ${buckets.length}`); @@ -616,6 +680,9 @@ async function fetchWithAuthData( success: true, buckets, projectId: authData.projectId, + tierLabel: supplementary.tierLabel, + tierId: supplementary.tierId, + creditBalance: supplementary.creditBalance, lastUpdated: Date.now(), accountId, }; diff --git a/src/cliproxy/quota-types.ts b/src/cliproxy/quota-types.ts index fa54ed69..29981f15 100644 --- a/src/cliproxy/quota-types.ts +++ b/src/cliproxy/quota-types.ts @@ -176,6 +176,8 @@ export interface GeminiCliBucket { remainingFraction: number; /** Remaining quota as percentage (0-100) */ remainingPercent: number; + /** Remaining quota count when the upstream API provides it */ + remainingAmount?: number | null; /** ISO timestamp when quota resets, null if unknown */ resetTime: string | null; /** Model IDs in this bucket */ @@ -192,6 +194,12 @@ export interface GeminiCliQuotaResult extends QuotaErrorMetadata { buckets: GeminiCliBucket[]; /** GCP project ID for this account */ projectId: string | null; + /** Human-readable Gemini tier label when available */ + tierLabel?: string | null; + /** Stable Gemini tier identifier when available */ + tierId?: string | null; + /** Available Google One AI credits when reported by the API */ + creditBalance?: number | null; /** Timestamp of fetch */ lastUpdated: number; /** Error message if fetch failed */ diff --git a/src/commands/cliproxy/quota-subcommand.ts b/src/commands/cliproxy/quota-subcommand.ts index 2ff268ed..eac0690b 100644 --- a/src/commands/cliproxy/quota-subcommand.ts +++ b/src/commands/cliproxy/quota-subcommand.ts @@ -638,15 +638,25 @@ function displayGeminiCliQuotaSection( if (quota.projectId) { console.log(` Project: ${dim(quota.projectId)}`); } + if (quota.tierLabel) { + console.log(` Tier: ${dim(quota.tierLabel)}`); + } + if (quota.creditBalance !== null && quota.creditBalance !== undefined) { + console.log(` Credits: ${dim(quota.creditBalance.toLocaleString())}`); + } for (const bucket of quota.buckets) { const bar = formatQuotaBar(bucket.remainingPercent); const tokenLabel = bucket.tokenType ? dim(` (${bucket.tokenType})`) : ''; + const amountLabel = + bucket.remainingAmount !== null && bucket.remainingAmount !== undefined + ? dim(` ${bucket.remainingAmount.toLocaleString()} left`) + : ''; const resetLabel = bucket.resetTime ? dim(` Resets ${formatResetTimeISO(bucket.resetTime)}`) : ''; console.log( - ` ${bucket.label.padEnd(24)} ${bar} ${bucket.remainingPercent.toFixed(0)}%${tokenLabel}${resetLabel}` + ` ${bucket.label.padEnd(24)} ${bar} ${bucket.remainingPercent.toFixed(0)}%${tokenLabel}${amountLabel}${resetLabel}` ); } console.log(''); diff --git a/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts b/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts index 66262b3c..10574df0 100644 --- a/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts +++ b/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts @@ -12,6 +12,7 @@ import { getCapturedFetchRequests, mockFetch, restoreFetch } from '../../mocks'; describe('Gemini CLI Quota Fetcher', () => { const GEMINI_QUOTA_URL = 'https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota'; + const GEMINI_CODE_ASSIST_URL = 'https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist'; const GOOGLE_TOKEN_URL = 'https://oauth2.googleapis.com/token'; let tempHome: string; let originalCcsHome: string | undefined; @@ -159,15 +160,29 @@ describe('Gemini CLI Quota Fetcher', () => { const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series'); expect(flashBucket).toBeDefined(); - // Takes minimum remaining fraction (0.6) - expect(flashBucket!.remainingFraction).toBe(0.6); - expect(flashBucket!.remainingPercent).toBe(60); + // Uses the preferred representative model when it exists + expect(flashBucket!.remainingFraction).toBe(0.8); + expect(flashBucket!.remainingPercent).toBe(80); const proBucket = buckets.find((b) => b.label === 'Gemini Pro Series'); expect(proBucket).toBeDefined(); expect(proBucket!.remainingFraction).toBe(0.9); }); + it('should split Flash Lite into its own bucket', () => { + const rawBuckets = [ + { model_id: 'gemini-2.5-flash-lite', remaining_fraction: 1 }, + { model_id: 'gemini-2.5-flash', remaining_fraction: 0.7 }, + ]; + + const buckets = buildGeminiCliBuckets(rawBuckets); + + expect(buckets.map((bucket) => bucket.label)).toEqual([ + 'Gemini Flash Lite Series', + 'Gemini Flash Series', + ]); + }); + it('should recognize Gemini 3.1 preview IDs during the rollout', () => { const rawBuckets = [ { model_id: 'gemini-3.1-flash-preview', remaining_fraction: 0.7 }, @@ -242,13 +257,13 @@ describe('Gemini CLI Quota Fetcher', () => { expect(buckets[0].remainingFraction).toBe(0.9); }); - it('should categorize unknown models as "other"', () => { + it('should preserve unknown model IDs instead of collapsing them', () => { const rawBuckets = [{ model_id: 'unknown-model-xyz', remaining_fraction: 0.7 }]; const buckets = buildGeminiCliBuckets(rawBuckets); expect(buckets).toHaveLength(1); - expect(buckets[0].label).toBe('Other Models'); + expect(buckets[0].label).toBe('unknown-model-xyz'); }); it('should handle empty buckets array', () => { @@ -268,7 +283,7 @@ describe('Gemini CLI Quota Fetcher', () => { expect(buckets[0].remainingFraction).toBe(0.8); }); - it('should keep earliest reset time when merging', () => { + it('should keep the representative model reset time when it exists', () => { const rawBuckets = [ { model_id: 'gemini-3-flash-preview', @@ -284,6 +299,26 @@ describe('Gemini CLI Quota Fetcher', () => { const buckets = buildGeminiCliBuckets(rawBuckets); + const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series'); + expect(flashBucket!.resetTime).toBe('2026-01-30T12:00:00Z'); + }); + + it('should keep earliest reset time when the representative model is missing', () => { + const rawBuckets = [ + { + model_id: 'gemini-3.1-flash-preview', + remaining_fraction: 0.8, + reset_time: '2026-01-30T12:00:00Z', + }, + { + model_id: 'gemini-2.5-flash', + remaining_fraction: 0.6, + reset_time: '2026-01-30T10:00:00Z', + }, + ]; + + const buckets = buildGeminiCliBuckets(rawBuckets); + const flashBucket = buckets.find((b) => b.label === 'Gemini Flash Series'); expect(flashBucket!.resetTime).toBe('2026-01-30T10:00:00Z'); }); @@ -311,6 +346,110 @@ describe('Gemini CLI Quota Fetcher', () => { }); }); + describe('fetchGeminiCliQuota success metadata', () => { + it('merges tier and credit metadata into successful quota responses', async () => { + writeActiveGeminiAccount('success@example.com'); + + mockFetch([ + { + url: GEMINI_QUOTA_URL, + method: 'POST', + status: 200, + response: { + buckets: [ + { + model_id: 'gemini-2.5-flash-lite', + remaining_fraction: 1, + remaining_amount: 100, + reset_time: '2026-01-30T09:00:00Z', + }, + { + model_id: 'gemini-3-flash-preview', + remaining_fraction: 0.82, + remaining_amount: 82, + reset_time: '2026-01-30T14:00:00Z', + }, + { + model_id: 'gemini-2.5-flash', + remaining_fraction: 0.4, + remaining_amount: 40, + reset_time: '2026-01-30T10:00:00Z', + }, + { + model_id: 'gemini-3.1-pro-preview', + remaining_fraction: 0.91, + remaining_amount: 91, + reset_time: '2026-01-30T15:00:00Z', + }, + ], + }, + }, + { + url: GEMINI_CODE_ASSIST_URL, + method: 'POST', + status: 200, + response: { + paidTier: { + id: 'g1-pro-tier', + availableCredits: [{ creditType: 'GOOGLE_ONE_AI', creditAmount: 12 }], + }, + }, + }, + ]); + + const result = await fetchGeminiCliQuota('success@example.com'); + + expect(result.success).toBe(true); + expect(result.tierLabel).toBe('Pro'); + expect(result.tierId).toBe('g1-pro-tier'); + expect(result.creditBalance).toBe(12); + expect(result.buckets.map((bucket) => bucket.label)).toEqual([ + 'Gemini Flash Lite Series', + 'Gemini Flash Series', + 'Gemini Pro Series', + ]); + expect(result.buckets[0].remainingAmount).toBe(100); + + const flashBucket = result.buckets.find((bucket) => bucket.label === 'Gemini Flash Series'); + expect(flashBucket?.remainingPercent).toBe(82); + expect(flashBucket?.remainingAmount).toBe(82); + expect(flashBucket?.resetTime).toBe('2026-01-30T14:00:00Z'); + + const requestUrls = getCapturedFetchRequests().map((request) => request.url); + expect(requestUrls).toContain(GEMINI_QUOTA_URL); + expect(requestUrls).toContain(GEMINI_CODE_ASSIST_URL); + }); + + it('keeps base quota success when supplementary metadata fails', async () => { + writeActiveGeminiAccount('supplementary-failure@example.com'); + + mockFetch([ + { + url: GEMINI_QUOTA_URL, + method: 'POST', + status: 200, + response: { + buckets: [{ model_id: 'gemini-3-flash-preview', remaining_fraction: 0.75 }], + }, + }, + { + url: GEMINI_CODE_ASSIST_URL, + method: 'POST', + status: 503, + response: { error: { message: 'Service unavailable' } }, + }, + ]); + + const result = await fetchGeminiCliQuota('supplementary-failure@example.com'); + + expect(result.success).toBe(true); + expect(result.tierLabel).toBeNull(); + expect(result.tierId).toBeNull(); + expect(result.creditBalance).toBeNull(); + expect(result.buckets[0].remainingPercent).toBe(75); + }); + }); + describe('fetchGeminiCliQuota failure metadata', () => { it('maps 401 responses to reauth-required metadata', async () => { writeActiveGeminiAccount('reauth@example.com');