/** * Builder functions for constructing quota data structures. */ import type { AntigravityQuotaGroup, AntigravityQuotaGroupDefinition, AntigravityQuotaInfo, AntigravityModelsPayload, GeminiCliParsedBucket, GeminiCliQuotaBucketState } from '@/types'; import { ANTIGRAVITY_QUOTA_GROUPS, GEMINI_CLI_GROUP_LOOKUP } from './constants'; import { normalizeQuotaFraction } from './parsers'; import { isIgnoredGeminiCliModel } from './validators'; export function pickEarlierResetTime(current?: string, next?: string): string | undefined { 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; } export 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); } export function buildGeminiCliQuotaBuckets( buckets: GeminiCliParsedBucket[] ): GeminiCliQuotaBucketState[] { if (buckets.length === 0) return []; const grouped = new Map(); buckets.forEach((bucket) => { if (isIgnoredGeminiCliModel(bucket.modelId)) return; 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 mapKey = `${groupId}::${tokenKey}`; const existing = grouped.get(mapKey); if (!existing) { grouped.set(mapKey, { id: `${groupId}${tokenKey ? `-${tokenKey}` : ''}`, label, remainingFraction: bucket.remainingFraction, remainingAmount: bucket.remainingAmount, resetTime: bucket.resetTime, tokenType: bucket.tokenType, modelIds: [bucket.modelId] }); return; } existing.remainingFraction = minNullableNumber( existing.remainingFraction, bucket.remainingFraction ); existing.remainingAmount = minNullableNumber(existing.remainingAmount, bucket.remainingAmount); existing.resetTime = pickEarlierResetTime(existing.resetTime, bucket.resetTime); existing.modelIds.push(bucket.modelId); }); return Array.from(grouped.values()).map((bucket) => { const uniqueModelIds = Array.from(new Set(bucket.modelIds)); return { id: bucket.id, label: bucket.label, remainingFraction: bucket.remainingFraction, remainingAmount: bucket.remainingAmount, resetTime: bucket.resetTime, tokenType: bucket.tokenType, modelIds: uniqueModelIds }; }); } export function getAntigravityQuotaInfo(entry?: AntigravityQuotaInfo): { remainingFraction: number | null; resetTime?: string; displayName?: string; } { if (!entry) { return { remainingFraction: null }; } const quotaInfo = entry.quotaInfo ?? entry.quota_info ?? {}; const remainingValue = quotaInfo.remainingFraction ?? quotaInfo.remaining_fraction ?? quotaInfo.remaining; const remainingFraction = normalizeQuotaFraction(remainingValue); const resetValue = quotaInfo.resetTime ?? quotaInfo.reset_time; const resetTime = typeof resetValue === 'string' ? resetValue : undefined; const displayName = typeof entry.displayName === 'string' ? entry.displayName : undefined; return { remainingFraction, resetTime, displayName }; } export function findAntigravityModel( models: AntigravityModelsPayload, identifier: string ): { id: string; entry: AntigravityQuotaInfo } | null { const direct = models[identifier]; if (direct) { return { id: identifier, entry: direct }; } const match = Object.entries(models).find(([, entry]) => { const name = typeof entry?.displayName === 'string' ? entry.displayName : ''; return name.toLowerCase() === identifier.toLowerCase(); }); if (match) { return { id: match[0], entry: match[1] }; } return null; } export function buildAntigravityQuotaGroups( models: AntigravityModelsPayload ): AntigravityQuotaGroup[] { const groups: AntigravityQuotaGroup[] = []; let geminiProResetTime: string | undefined; const [claudeDef, geminiProDef, flashDef, flashLiteDef, cuDef, geminiFlashDef, imageDef] = ANTIGRAVITY_QUOTA_GROUPS; const buildGroup = ( def: AntigravityQuotaGroupDefinition, overrideResetTime?: string ): AntigravityQuotaGroup | null => { const matches = def.identifiers .map((identifier) => findAntigravityModel(models, identifier)) .filter((entry): entry is { id: string; entry: AntigravityQuotaInfo } => Boolean(entry)); const quotaEntries = matches .map(({ id, entry }) => { const info = getAntigravityQuotaInfo(entry); const remainingFraction = info.remainingFraction ?? (info.resetTime ? 0 : null); if (remainingFraction === null) return null; return { id, remainingFraction, resetTime: info.resetTime, displayName: info.displayName }; }) .filter((entry): entry is NonNullable => entry !== null); if (quotaEntries.length === 0) return null; const remainingFraction = Math.min(...quotaEntries.map((entry) => entry.remainingFraction)); const resetTime = overrideResetTime ?? quotaEntries.map((entry) => entry.resetTime).find(Boolean); const displayName = quotaEntries.map((entry) => entry.displayName).find(Boolean); const label = def.labelFromModel && displayName ? displayName : def.label; return { id: def.id, label, models: quotaEntries.map((entry) => entry.id), remainingFraction, resetTime }; }; const claudeGroup = buildGroup(claudeDef); if (claudeGroup) { groups.push(claudeGroup); } const geminiProGroup = buildGroup(geminiProDef); if (geminiProGroup) { geminiProResetTime = geminiProGroup.resetTime; groups.push(geminiProGroup); } const flashGroup = buildGroup(flashDef); if (flashGroup) { groups.push(flashGroup); } const flashLiteGroup = buildGroup(flashLiteDef); if (flashLiteGroup) { groups.push(flashLiteGroup); } const cuGroup = buildGroup(cuDef); if (cuGroup) { groups.push(cuGroup); } const geminiFlashGroup = buildGroup(geminiFlashDef); if (geminiFlashGroup) { groups.push(geminiFlashGroup); } const imageGroup = buildGroup(imageDef, geminiProResetTime); if (imageGroup) { groups.push(imageGroup); } return groups; }