From 7af81c2436b9f1b66b90369abdf170a233c42783 Mon Sep 17 00:00:00 2001 From: Nelson Pires Date: Sun, 8 Feb 2026 10:42:15 -0300 Subject: [PATCH] Add per-model quotas with collapsible model groups in header (#355) * feat(ui): show per-model quota groups in header with collapsible families Add per-model quota groups under each provider in the header Introduce collapsible sections for model families to reveal models Show all models when none are explicitly selected or respect explicit selections * feat: add toggle to UsageCard for dropdown visibility Introduce a switch in UsageCard to control inclusion in the dropdown Hide the percent label when the toggle is visible * feat(usage): enhance UsagePage with model grouping and collapsibles Add model lists grouped by family for the selected provider Enable collapsible sections per family to toggle visibility Apply and persist default model selections on provider change * feat(quota): add model family helpers Add utilities to categorize models into families by provider Enable grouping of models by family for header and usage pages Define default models for Gemini 3.x and Claude families * feat(desktop): extend settings with model grouping and selection Track per-provider selected models for usage Allow collapsing and expanding families in the usage page Support per-provider custom model groups with labels and assignments * feat: persist usage preferences in persistence Persist usageSelectedModels per provider Persist usageCollapsedFamilies and usageExpandedFamilies states Support custom usageModelGroups with groups and assignments * feat: track selected quota models and expanded families per provider Initialize selectedModels and expandedFamilies state for quota providers Add actions to set, toggle, and apply default model selections and expanded families Persist selections to desktop settings and apply defaults on load * feat(server): sanitize usage model configuration in settings update Validate and sanitize usage selections per provider Persist only valid usageCollapsedFamilies and usageExpandedFamilies in settings Enforce limits on custom model groups and model assignments * feat: add collapsible model families in header Add collapsible sections for model families within each provider in the header Show per-model usage with percent display and a progress bar Toggle expansion via arrow icons and preserve expanded state per provider * fix(ui): treat explicit per-provider model selections correctly in header Enable showing all models by default when a provider has no explicit selection Recognize an explicit per-provider selection when a provider key exists in selectedModels Filter to selected models only if an explicit selection is present for the provider --- packages/ui/src/components/layout/Header.tsx | 325 +++++++++++++++--- .../components/sections/usage/UsageCard.tsx | 27 +- .../components/sections/usage/UsagePage.tsx | 198 ++++++++++- packages/ui/src/lib/desktop.ts | 8 + packages/ui/src/lib/persistence.ts | 98 ++++++ packages/ui/src/lib/quota/model-families.ts | 110 ++++++ packages/ui/src/stores/useQuotaStore.ts | 105 +++++- packages/web/server/index.js | 108 ++++++ 8 files changed, 910 insertions(+), 69 deletions(-) create mode 100644 packages/ui/src/lib/quota/model-families.ts diff --git a/packages/ui/src/components/layout/Header.tsx b/packages/ui/src/components/layout/Header.tsx index 80259766..7a3dd280 100644 --- a/packages/ui/src/components/layout/Header.tsx +++ b/packages/ui/src/components/layout/Header.tsx @@ -32,6 +32,18 @@ import { ProviderLogo } from '@/components/ui/ProviderLogo'; import { formatPercent, formatWindowLabel, QUOTA_PROVIDERS } from '@/lib/quota'; import { UsageProgressBar } from '@/components/sections/usage/UsageProgressBar'; import { updateDesktopSettings } from '@/lib/persistence'; +import { + getAllModelFamilies, + groupModelsByFamily, + sortModelFamilies, +} from '@/lib/quota/model-families'; + +import { + Collapsible, + CollapsibleContent, + CollapsibleTrigger, +} from '@/components/ui/collapsible'; +import { RiArrowDownSLine, RiArrowRightSLine } from '@remixicon/react'; import type { UsageWindow } from '@/types'; import type { GitHubAuthStatus } from '@/lib/api/types'; import { DesktopHostSwitcherButton } from '@/components/desktop/DesktopHostSwitcher'; @@ -184,12 +196,23 @@ export const Header: React.FC = () => { const [isSwitchingGitHubAccount, setIsSwitchingGitHubAccount] = React.useState(false); const [isMobileRateLimitsOpen, setIsMobileRateLimitsOpen] = React.useState(false); useQuotaAutoRefresh(); + const selectedModels = useQuotaStore((state) => state.selectedModels); + const expandedFamilies = useQuotaStore((state) => state.expandedFamilies); + const toggleFamilyExpanded = useQuotaStore((state) => state.toggleFamilyExpanded); + + interface RateLimitGroup { + providerId: string; + providerName: string; + entries: Array<[string, UsageWindow]>; + modelFamilies?: Array<{ + familyId: string | null; + familyLabel: string; + models: Array<[string, UsageWindow]>; + }>; + } + const rateLimitGroups = React.useMemo(() => { - const groups: Array<{ - providerId: string; - providerName: string; - entries: Array<[string, UsageWindow]>; - }> = []; + const groups: RateLimitGroup[] = []; for (const provider of QUOTA_PROVIDERS) { if (!dropdownProviderIds.includes(provider.id)) { @@ -197,14 +220,91 @@ export const Header: React.FC = () => { } const result = quotaResults.find((entry) => entry.providerId === provider.id); const windows = (result?.usage?.windows ?? {}) as Record; + const models = result?.usage?.models; const entries = Object.entries(windows); - if (entries.length > 0) { - groups.push({ providerId: provider.id, providerName: provider.name, entries }); + + const group: RateLimitGroup = { + providerId: provider.id, + providerName: provider.name, + entries, + }; + + // Add model families if provider has per-model quotas + if (models && Object.keys(models).length > 0) { + const providerSelectedModels = selectedModels[provider.id] ?? []; + // hasExplicitSelection = true means user touched the selection (even if empty) + // hasExplicitSelection = false means no preference (key missing) → show all by default + const hasExplicitSelection = provider.id in selectedModels; + const modelGroups = groupModelsByFamily(models, provider.id); + const families = getAllModelFamilies(provider.id); + const sortedFamilies = sortModelFamilies(families); + + group.modelFamilies = []; + + // Add predefined families first + for (const family of sortedFamilies) { + const modelNames = modelGroups.get(family.id) ?? []; + if (modelNames.length === 0) continue; + + // Filter to selected models only, OR show all if nothing selected + const selectedModelNames = hasExplicitSelection + ? modelNames.filter((m: string) => providerSelectedModels.includes(m)) + : modelNames; + if (selectedModelNames.length === 0) continue; + + const familyModels: Array<[string, UsageWindow]> = []; + for (const modelName of selectedModelNames) { + const modelUsage = models[modelName] as { windows?: Record } | undefined; + if (modelUsage?.windows) { + const windowEntries = Object.entries(modelUsage.windows); + if (windowEntries.length > 0) { + familyModels.push([modelName, windowEntries[0][1]]); + } + } + } + + if (familyModels.length > 0) { + group.modelFamilies.push({ + familyId: family.id, + familyLabel: family.label, + models: familyModels, + }); + } + } + + // Add "Other" family for remaining models + const otherModelNames = modelGroups.get(null) ?? []; + const selectedOtherModels = hasExplicitSelection + ? otherModelNames.filter((m: string) => providerSelectedModels.includes(m)) + : otherModelNames; + if (selectedOtherModels.length > 0) { + const otherModels: Array<[string, UsageWindow]> = []; + for (const modelName of selectedOtherModels) { + const modelUsage = models[modelName] as { windows?: Record } | undefined; + if (modelUsage?.windows) { + const windowEntries = Object.entries(modelUsage.windows); + if (windowEntries.length > 0) { + otherModels.push([modelName, windowEntries[0][1]]); + } + } + } + if (otherModels.length > 0) { + group.modelFamilies.push({ + familyId: null, + familyLabel: 'Other', + models: otherModels, + }); + } + } + } + + if (entries.length > 0 || (group.modelFamilies && group.modelFamilies.length > 0)) { + groups.push(group); } } return groups; - }, [dropdownProviderIds, quotaResults]); + }, [dropdownProviderIds, quotaResults, selectedModels]); const hasRateLimits = rateLimitGroups.length > 0; React.useEffect(() => { void loadQuotaSettings(); @@ -747,54 +847,125 @@ export const Header: React.FC = () => { No rate limits available. )} - {rateLimitGroups.map((group, index) => ( - - - - {group.providerName} - - {group.entries.length === 0 ? ( - event.preventDefault()} - > - No rate limits reported. - - ) : ( - group.entries.map(([label, window]) => ( + {rateLimitGroups.map((group, index) => { + const providerExpandedFamilies = expandedFamilies[group.providerId] ?? []; + + return ( + + + + {group.providerName} + + + {/* Provider-level entries */} + {group.entries.length === 0 && (!group.modelFamilies || group.modelFamilies.length === 0) ? ( event.preventDefault()} > - - {(() => { - const displayPercent = quotaDisplayMode === 'remaining' - ? window.remainingPercent - : window.usedPercent; - return ( - <> - - {formatWindowLabel(label)} - - {formatPercent(displayPercent) === '-' ? '' : formatPercent(displayPercent)} - - - - - {window.resetAfterFormatted ?? window.resetAtFormatted ?? ''} - - - ); - })()} - + No rate limits reported. - )) - )} - {index < rateLimitGroups.length - 1 && } - - ))} + ) : ( + <> + {/* Provider-level windows */} + {group.entries.map(([label, window]) => ( + event.preventDefault()} + > + + {(() => { + const displayPercent = quotaDisplayMode === 'remaining' + ? window.remainingPercent + : window.usedPercent; + return ( + <> + + {formatWindowLabel(label)} + + {formatPercent(displayPercent) === '-' ? '' : formatPercent(displayPercent)} + + + + + {window.resetAfterFormatted ?? window.resetAtFormatted ?? ''} + + + ); + })()} + + + ))} + + {/* Model families with collapsible sections - default COLLAPSED */} + {group.modelFamilies && group.modelFamilies.length > 0 && ( +
+ {group.modelFamilies.map((family) => { + const isExpanded = providerExpandedFamilies.includes(family.familyId ?? 'other'); + + return ( + toggleFamilyExpanded(group.providerId, family.familyId ?? 'other')} + > + + + {family.familyLabel} + + {isExpanded ? ( + + ) : ( + + )} + + +
+ {family.models.map(([modelName, window]) => ( +
+
+ + {modelName} + {(() => { + const displayPercent = quotaDisplayMode === 'remaining' + ? window.remainingPercent + : window.usedPercent; + return ( + + {formatPercent(displayPercent) === '-' ? '' : formatPercent(displayPercent)} + + ); + })()} + + {(() => { + const displayPercent = quotaDisplayMode === 'remaining' + ? window.remainingPercent + : window.usedPercent; + return ( + + ); + })()} +
+
+ ))} +
+ + + ); + })} +
+ )} + + )} + {index < rateLimitGroups.length - 1 && } + + ); + })} @@ -1131,6 +1302,56 @@ export const Header: React.FC = () => { ); })} + {/* Model families with collapsible sections */} + {group.modelFamilies && group.modelFamilies.length > 0 && ( +
+ {group.modelFamilies.map((family) => { + const providerExpandedFamilies = expandedFamilies[group.providerId] ?? []; + const isExpanded = providerExpandedFamilies.includes(family.familyId ?? 'other'); + + return ( + toggleFamilyExpanded(group.providerId, family.familyId ?? 'other')} + > + + + {family.familyLabel} + + {isExpanded ? ( + + ) : ( + + )} + + +
+ {family.models.map(([modelName, window]) => { + const displayPercent = quotaDisplayMode === 'remaining' + ? window.remainingPercent + : window.usedPercent; + return ( +
+
+ + {modelName} + + + {formatPercent(displayPercent) === '-' ? '' : formatPercent(displayPercent)} + +
+ +
+ ); + })} +
+ + + ); + })} +
+ )} ))} diff --git a/packages/ui/src/components/sections/usage/UsageCard.tsx b/packages/ui/src/components/sections/usage/UsageCard.tsx index 882ad082..13f69458 100644 --- a/packages/ui/src/components/sections/usage/UsageCard.tsx +++ b/packages/ui/src/components/sections/usage/UsageCard.tsx @@ -3,14 +3,25 @@ import type { UsageWindow } from '@/types'; import { formatPercent, formatWindowLabel } from '@/lib/quota'; import { UsageProgressBar } from './UsageProgressBar'; import { useQuotaStore } from '@/stores/useQuotaStore'; +import { Switch } from '@/components/ui/switch'; interface UsageCardProps { title: string; window: UsageWindow; subtitle?: string | null; + showToggle?: boolean; + toggleEnabled?: boolean; + onToggle?: (enabled: boolean) => void; } -export const UsageCard: React.FC = ({ title, window, subtitle }) => { +export const UsageCard: React.FC = ({ + title, + window, + subtitle, + showToggle = false, + toggleEnabled = false, + onToggle, +}) => { const displayMode = useQuotaStore((state) => state.displayMode); const displayPercent = displayMode === 'remaining' ? window.remainingPercent : window.usedPercent; const barLabel = displayMode === 'remaining' ? 'remaining' : 'used'; @@ -21,13 +32,23 @@ export const UsageCard: React.FC = ({ title, window, subtitle }) return (
-
+
{windowLabel}
{subtitle && (
{subtitle}
)}
-
{percentLabel === '-' ? '' : percentLabel}
+ {showToggle ? ( + + ) : ( +
+ {percentLabel === '-' ? '' : percentLabel} +
+ )}
diff --git a/packages/ui/src/components/sections/usage/UsagePage.tsx b/packages/ui/src/components/sections/usage/UsagePage.tsx index 6e237b1f..675e9baf 100644 --- a/packages/ui/src/components/sections/usage/UsagePage.tsx +++ b/packages/ui/src/components/sections/usage/UsagePage.tsx @@ -6,6 +6,14 @@ import { useQuotaAutoRefresh, useQuotaStore } from '@/stores/useQuotaStore'; import { Switch } from '@/components/ui/switch'; import { updateDesktopSettings } from '@/lib/persistence'; import { ProviderLogo } from '@/components/ui/ProviderLogo'; +import { + Collapsible, + CollapsibleContent, + CollapsibleTrigger, +} from '@/components/ui/collapsible'; +import { RiArrowDownSLine, RiArrowRightSLine } from '@remixicon/react'; +import type { UsageWindows, QuotaProviderId } from '@/types'; +import { getAllModelFamilies, sortModelFamilies, groupModelsByFamilyWithGetter } from '@/lib/quota/model-families'; const formatTime = (timestamp: number | null) => { if (!timestamp) return '-'; @@ -19,6 +27,11 @@ const formatTime = (timestamp: number | null) => { } }; +interface ModelInfo { + name: string; + windows: UsageWindows; +} + export const UsagePage: React.FC = () => { const results = useQuotaStore((state) => state.results); const selectedProviderId = useQuotaStore((state) => state.selectedProviderId); @@ -30,6 +43,9 @@ export const UsagePage: React.FC = () => { const error = useQuotaStore((state) => state.error); const dropdownProviderIds = useQuotaStore((state) => state.dropdownProviderIds); const setDropdownProviderIds = useQuotaStore((state) => state.setDropdownProviderIds); + const selectedModels = useQuotaStore((state) => state.selectedModels); + const toggleModelSelected = useQuotaStore((state) => state.toggleModelSelected); + const applyDefaultSelections = useQuotaStore((state) => state.applyDefaultSelections); useQuotaAutoRefresh(); @@ -67,6 +83,69 @@ export const UsagePage: React.FC = () => { void updateDesktopSettings({ usageDropdownProviders: next }); }, [dropdownProviderIds, selectedProviderId, setDropdownProviderIds]); + // Get models for the selected provider + const providerModels = React.useMemo((): ModelInfo[] => { + if (!usage?.models) return []; + return Object.entries(usage.models) + .map(([name, modelUsage]) => ({ name, windows: modelUsage })) + .filter((model) => Object.keys(model.windows.windows).length > 0); + }, [usage?.models]); + + // Apply default selections on mount if no prior selections exist + React.useEffect(() => { + if (selectedProviderId && providerModels.length > 0) { + applyDefaultSelections(selectedProviderId, providerModels.map((m) => m.name)); + } + }, [selectedProviderId, providerModels, applyDefaultSelections]); + + // Group models by family + const modelsByFamily = React.useMemo(() => { + if (!selectedProviderId || providerModels.length === 0) { + return new Map(); + } + return groupModelsByFamilyWithGetter( + providerModels, + (model) => model.name, + selectedProviderId as QuotaProviderId + ); + }, [providerModels, selectedProviderId]); + + // Get sorted families + const sortedFamilies = React.useMemo(() => { + if (!selectedProviderId) return []; + const families = getAllModelFamilies(selectedProviderId as QuotaProviderId); + return sortModelFamilies(families); + }, [selectedProviderId]); + + // Collapsible state for family sections (persist per provider) + const [collapsedFamilies, setCollapsedFamilies] = React.useState>(() => { + // Default: all families start expanded (not collapsed) + return {}; + }); + + const toggleFamilyCollapsed = React.useCallback((familyId: string) => { + setCollapsedFamilies((prev) => ({ + ...prev, + [familyId]: !prev[familyId], + })); + }, []); + + const handleModelToggle = React.useCallback((modelName: string) => { + if (!selectedProviderId) return; + toggleModelSelected(selectedProviderId, modelName); + // Also update settings to persist + const currentSelected = selectedModels[selectedProviderId] ?? []; + const isSelected = currentSelected.includes(modelName); + const nextSelected = isSelected + ? currentSelected.filter((m) => m !== modelName) + : [...currentSelected, modelName]; + const nextSettings: Record = { ...selectedModels, [selectedProviderId]: nextSelected }; + void updateDesktopSettings({ usageSelectedModels: nextSettings }); + }, [selectedProviderId, selectedModels, toggleModelSelected]); + + // Get selected models for this provider + const providerSelectedModels = selectedProviderId ? (selectedModels[selectedProviderId] ?? []) : []; + if (!selectedProviderId) { return (
@@ -131,22 +210,119 @@ export const UsagePage: React.FC = () => {
)} - {usage?.models && Object.keys(usage.models).length > 0 && ( -
-
Model Quotas
- {Object.entries(usage.models).map(([modelName, modelUsage]) => { - const entries = Object.entries(modelUsage.windows); - if (entries.length === 0) { - return null; - } - const [label, window] = entries[0]; - return ; + {/* Models Section - Grouped by Family */} + {providerModels.length > 0 && ( +
+
+

Models

+

+ Toggle on to show in header dropdown +

+
+ + {/* Predefined families (Gemini, Claude) */} + {sortedFamilies.map((family) => { + const familyModels = modelsByFamily.get(family.id) ?? []; + if (familyModels.length === 0) return null; + + const isCollapsed = collapsedFamilies[family.id] ?? false; + + return ( + toggleFamilyCollapsed(family.id)} + > + +
+
{family.label}
+

+ {familyModels.length} model{familyModels.length !== 1 ? 's' : ''} +

+
+ {isCollapsed ? ( + + ) : ( + + )} + + + {familyModels.map((model) => { + const entries = Object.entries(model.windows.windows); + if (entries.length === 0) return null; + const [label, window] = entries[0]; + const isSelected = providerSelectedModels.includes(model.name); + + return ( + handleModelToggle(model.name)} + /> + ); + })} + + + ); })} + + {/* Other family */} + {(() => { + const otherModels = modelsByFamily.get(null) ?? []; + if (otherModels.length === 0) return null; + + const isCollapsed = collapsedFamilies['other'] ?? false; + + return ( + toggleFamilyCollapsed('other')} + > + +
+
Other
+

+ {otherModels.length} model{otherModels.length !== 1 ? 's' : ''} +

+
+ {isCollapsed ? ( + + ) : ( + + )} + + + {otherModels.map((model) => { + const entries = Object.entries(model.windows.windows); + if (entries.length === 0) return null; + const [label, window] = entries[0]; + const isSelected = providerSelectedModels.includes(model.name); + + return ( + handleModelToggle(model.name)} + /> + ); + })} + + + ); + })()}
)} {selectedResult?.configured && usage && Object.keys(usage.windows ?? {}).length === 0 && - Object.keys(usage.models ?? {}).length === 0 && ( + providerModels.length === 0 && (

No quota windows reported for this provider.

diff --git a/packages/ui/src/lib/desktop.ts b/packages/ui/src/lib/desktop.ts index 7c4720c8..d0901ab3 100644 --- a/packages/ui/src/lib/desktop.ts +++ b/packages/ui/src/lib/desktop.ts @@ -53,6 +53,14 @@ export type DesktopSettings = { usageRefreshIntervalMs?: number; usageDisplayMode?: 'usage' | 'remaining'; usageDropdownProviders?: string[]; + usageSelectedModels?: Record; // Map of providerId -> selected model names + usageCollapsedFamilies?: Record; // Map of providerId -> collapsed family IDs (UsagePage) + usageExpandedFamilies?: Record; // Map of providerId -> EXPANDED family IDs (header dropdown - inverted) + usageModelGroups?: Record; + modelAssignments?: Record; // modelName -> groupId + renamedGroups?: Record; // groupId -> custom label + }>; // Per-provider custom model groups configuration autoDeleteEnabled?: boolean; autoDeleteAfterDays?: number; defaultModel?: string; // format: "provider/model" diff --git a/packages/ui/src/lib/persistence.ts b/packages/ui/src/lib/persistence.ts index 33c68537..f7a813e1 100644 --- a/packages/ui/src/lib/persistence.ts +++ b/packages/ui/src/lib/persistence.ts @@ -421,6 +421,104 @@ const sanitizeWebSettings = (payload: unknown): DesktopSettings | null => { (entry): entry is string => typeof entry === 'string' && entry.length > 0 ); } + + // Parse usageSelectedModels (Record) + if (candidate.usageSelectedModels && typeof candidate.usageSelectedModels === 'object') { + const selectedModels: Record = {}; + for (const [providerId, models] of Object.entries(candidate.usageSelectedModels)) { + if (Array.isArray(models)) { + selectedModels[providerId] = models.filter((m): m is string => typeof m === 'string'); + } + } + if (Object.keys(selectedModels).length > 0) { + result.usageSelectedModels = selectedModels; + } + } + + // Parse usageCollapsedFamilies (Record) + if (candidate.usageCollapsedFamilies && typeof candidate.usageCollapsedFamilies === 'object') { + const collapsedFamilies: Record = {}; + for (const [providerId, families] of Object.entries(candidate.usageCollapsedFamilies)) { + if (Array.isArray(families)) { + collapsedFamilies[providerId] = families.filter((f): f is string => typeof f === 'string'); + } + } + if (Object.keys(collapsedFamilies).length > 0) { + result.usageCollapsedFamilies = collapsedFamilies; + } + } + + // Parse usageExpandedFamilies (Record) - inverted collapsed logic for header dropdown + if (candidate.usageExpandedFamilies && typeof candidate.usageExpandedFamilies === 'object') { + const expandedFamilies: Record = {}; + for (const [providerId, families] of Object.entries(candidate.usageExpandedFamilies)) { + if (Array.isArray(families)) { + expandedFamilies[providerId] = families.filter((f): f is string => typeof f === 'string'); + } + } + if (Object.keys(expandedFamilies).length > 0) { + result.usageExpandedFamilies = expandedFamilies; + } + } + + // Parse usageModelGroups - custom model groups configuration per provider + if (candidate.usageModelGroups && typeof candidate.usageModelGroups === 'object') { + const modelGroups: Record; + modelAssignments?: Record; + renamedGroups?: Record; + }> = {}; + for (const [providerId, config] of Object.entries(candidate.usageModelGroups)) { + if (config && typeof config === 'object') { + const typedConfig = config as Record; + const providerConfig: { + customGroups?: Array<{id: string; label: string; models: string[]; order: number}>; + modelAssignments?: Record; + renamedGroups?: Record; + } = {}; + + // Parse customGroups + if (Array.isArray(typedConfig.customGroups)) { + providerConfig.customGroups = typedConfig.customGroups + .filter((g): g is Record => g && typeof g === 'object') + .map((g) => ({ + id: String(g.id ?? ''), + label: String(g.label ?? ''), + models: Array.isArray(g.models) + ? g.models.filter((m): m is string => typeof m === 'string') + : [], + order: typeof g.order === 'number' ? g.order : 0, + })); + } + + // Parse modelAssignments + if (typedConfig.modelAssignments && typeof typedConfig.modelAssignments === 'object') { + providerConfig.modelAssignments = Object.fromEntries( + Object.entries(typedConfig.modelAssignments as Record) + .filter(([, v]) => typeof v === 'string') + .map(([k, v]) => [k, String(v)]) + ); + } + + // Parse renamedGroups + if (typedConfig.renamedGroups && typeof typedConfig.renamedGroups === 'object') { + providerConfig.renamedGroups = Object.fromEntries( + Object.entries(typedConfig.renamedGroups as Record) + .filter(([, v]) => typeof v === 'string') + .map(([k, v]) => [k, String(v)]) + ); + } + + if (Object.keys(providerConfig).length > 0) { + modelGroups[providerId] = providerConfig; + } + } + } + if (Object.keys(modelGroups).length > 0) { + result.usageModelGroups = modelGroups; + } + } + if ( typeof candidate.toolCallExpansion === 'string' && (candidate.toolCallExpansion === 'collapsed' diff --git a/packages/ui/src/lib/quota/model-families.ts b/packages/ui/src/lib/quota/model-families.ts new file mode 100644 index 00000000..3d9faf70 --- /dev/null +++ b/packages/ui/src/lib/quota/model-families.ts @@ -0,0 +1,110 @@ +import type { QuotaProviderId } from '@/types'; + +export interface ModelFamily { + id: string; + label: string; + matcher: (modelName: string) => boolean; + order: number; +} + +const GOOGLE_MODEL_FAMILIES: ModelFamily[] = [ + { + id: 'gemini', + label: 'Gemini', + matcher: (modelName) => modelName.toLowerCase().startsWith('gemini-'), + order: 1, + }, + { + id: 'claude', + label: 'Claude', + matcher: (modelName) => modelName.toLowerCase().startsWith('claude-'), + order: 2, + }, +]; + +export const PROVIDER_MODEL_FAMILIES: Record = { + google: GOOGLE_MODEL_FAMILIES, +}; + +export function getModelFamily(modelName: string, providerId: QuotaProviderId): ModelFamily | null { + const families = PROVIDER_MODEL_FAMILIES[providerId] ?? []; + for (const family of families) { + if (family.matcher(modelName)) { + return family; + } + } + return null; +} + +export function getAllModelFamilies(providerId: QuotaProviderId): ModelFamily[] { + return PROVIDER_MODEL_FAMILIES[providerId] ?? []; +} + +export function sortModelFamilies(families: ModelFamily[]): ModelFamily[] { + return [...families].sort((a, b) => a.order - b.order); +} + +/** + * Group model names by family (for backward compatibility with Header.tsx) + */ +export function groupModelsByFamily( + models: Record, + providerId: QuotaProviderId +): Map { + const groups = new Map(); + + for (const modelName of Object.keys(models)) { + const family = getModelFamily(modelName, providerId); + const familyId = family?.id ?? null; + + if (!groups.has(familyId)) { + groups.set(familyId, []); + } + groups.get(familyId)!.push(modelName); + } + + return groups; +} + +/** + * Group models by family with custom getter function (for UsagePage.tsx) + */ +export function groupModelsByFamilyWithGetter( + models: T[], + getModelName: (model: T) => string, + providerId: QuotaProviderId +): Map { + const groups = new Map(); + + for (const model of models) { + const modelName = getModelName(model); + const family = getModelFamily(modelName, providerId); + const familyId = family?.id ?? null; + + if (!groups.has(familyId)) { + groups.set(familyId, []); + } + groups.get(familyId)!.push(model); + } + + return groups; +} + +/** + * Get default models for a provider based on simple patterns. + * - Gemini 3.x models (starting with gemini-3-) + * - All Claude models + */ +export function getDefaultModels( + providerId: QuotaProviderId, + availableModels: string[] +): string[] { + return availableModels.filter((model) => { + const lower = model.toLowerCase(); + // Gemini 3.x + if (lower.startsWith('gemini-3-')) return true; + // All Claude models + if (lower.startsWith('claude-')) return true; + return false; + }); +} diff --git a/packages/ui/src/stores/useQuotaStore.ts b/packages/ui/src/stores/useQuotaStore.ts index 8d5b8cb4..2756210c 100644 --- a/packages/ui/src/stores/useQuotaStore.ts +++ b/packages/ui/src/stores/useQuotaStore.ts @@ -5,6 +5,8 @@ import type { ProviderResult, QuotaProviderId } from '@/types'; import { QUOTA_PROVIDERS } from '@/lib/quota'; import { isVSCodeRuntime } from '@/lib/desktop'; import { getRegisteredRuntimeAPIs } from '@/contexts/runtimeAPIRegistry'; +import { getDefaultModels } from '@/lib/quota/model-families'; +import { updateDesktopSettings } from '@/lib/persistence'; const DEFAULT_REFRESH_INTERVAL_MS = 60000; @@ -13,6 +15,8 @@ interface QuotaSettingsState { refreshIntervalMs: number; displayMode: 'usage' | 'remaining'; dropdownProviderIds: QuotaProviderId[]; + selectedModels: Record; // Map of providerId -> selected model names + expandedFamilies: Record; // Map of providerId -> EXPANDED family IDs (header dropdown - inverted) } interface QuotaStore extends QuotaSettingsState { @@ -31,6 +35,11 @@ interface QuotaStore extends QuotaSettingsState { setRefreshInterval: (intervalMs: number) => void; setDisplayMode: (mode: 'usage' | 'remaining') => void; setDropdownProviderIds: (providerIds: QuotaProviderId[]) => void; + setSelectedModels: (providerId: string, modelNames: string[]) => void; + toggleModelSelected: (providerId: string, modelName: string) => void; + setExpandedFamilies: (providerId: string, familyIds: string[]) => void; + toggleFamilyExpanded: (providerId: string, familyId: string) => void; + applyDefaultSelections: (providerId: string, availableModels: string[]) => void; } const parseSettings = (data: Record | null): QuotaSettingsState => { @@ -53,7 +62,36 @@ const parseSettings = (data: Record | null): QuotaSettingsState ) : allProviderIds; - return { autoRefresh, refreshIntervalMs, displayMode, dropdownProviderIds }; + // Parse selected models (providerId -> array of model names) + const selectedModels: Record = {}; + const rawSelectedModels = data?.usageSelectedModels; + if (rawSelectedModels && typeof rawSelectedModels === 'object') { + for (const [providerId, models] of Object.entries(rawSelectedModels)) { + if (Array.isArray(models)) { + selectedModels[providerId] = models.filter((m): m is string => typeof m === 'string'); + } + } + } + + // Parse expanded families (inverted collapsed logic for header dropdown) + const expandedFamilies: Record = {}; + const rawExpandedFamilies = data?.usageExpandedFamilies; + if (rawExpandedFamilies && typeof rawExpandedFamilies === 'object') { + for (const [providerId, families] of Object.entries(rawExpandedFamilies)) { + if (Array.isArray(families)) { + expandedFamilies[providerId] = families.filter((f): f is string => typeof f === 'string'); + } + } + } + + return { + autoRefresh, + refreshIntervalMs, + displayMode, + dropdownProviderIds, + selectedModels, + expandedFamilies, + }; }; const loadSettingsFromRuntime = async (): Promise => { @@ -83,7 +121,9 @@ const loadSettingsFromRuntime = async (): Promise => { autoRefresh: false, refreshIntervalMs: DEFAULT_REFRESH_INTERVAL_MS, displayMode: 'usage', - dropdownProviderIds: QUOTA_PROVIDERS.map((provider) => provider.id) + dropdownProviderIds: QUOTA_PROVIDERS.map((provider) => provider.id), + selectedModels: {}, + expandedFamilies: {}, }; }; @@ -100,6 +140,8 @@ export const useQuotaStore = create()( refreshIntervalMs: DEFAULT_REFRESH_INTERVAL_MS, displayMode: 'usage', dropdownProviderIds: QUOTA_PROVIDERS.map((provider) => provider.id), + selectedModels: {}, + expandedFamilies: {}, loadSettings: async () => { try { @@ -174,7 +216,64 @@ export const useQuotaStore = create()( set({ refreshIntervalMs: clamped }); }, setDisplayMode: (mode) => set({ displayMode: mode }), - setDropdownProviderIds: (providerIds) => set({ dropdownProviderIds: providerIds }) + setDropdownProviderIds: (providerIds) => set({ dropdownProviderIds: providerIds }), + + setSelectedModels: (providerId, modelNames) => { + set((state) => ({ + selectedModels: { ...state.selectedModels, [providerId]: modelNames } + })); + }, + + toggleModelSelected: (providerId, modelName) => { + set((state) => { + const currentSelected = state.selectedModels[providerId] ?? []; + const isSelected = currentSelected.includes(modelName); + const nextSelected = isSelected + ? currentSelected.filter((m) => m !== modelName) + : [...currentSelected, modelName]; + return { + selectedModels: { ...state.selectedModels, [providerId]: nextSelected } + }; + }); + }, + + setExpandedFamilies: (providerId, familyIds) => { + set((state) => ({ + expandedFamilies: { ...state.expandedFamilies, [providerId]: familyIds } + })); + // Persist + void updateDesktopSettings({ usageExpandedFamilies: get().expandedFamilies }); + }, + + toggleFamilyExpanded: (providerId, familyId) => { + set((state) => { + const currentExpanded = state.expandedFamilies[providerId] ?? []; + const isExpanded = currentExpanded.includes(familyId); + const nextExpanded = isExpanded + ? currentExpanded.filter((id) => id !== familyId) + : [...currentExpanded, familyId]; + return { + expandedFamilies: { ...state.expandedFamilies, [providerId]: nextExpanded } + }; + }); + // Persist + void updateDesktopSettings({ usageExpandedFamilies: get().expandedFamilies }); + }, + + applyDefaultSelections: (providerId, availableModels) => { + const state = get(); + // Only apply if no prior selections exist + if ((state.selectedModels[providerId]?.length ?? 0) > 0) return; + + const defaults = getDefaultModels(providerId as QuotaProviderId, availableModels); + if (defaults.length === 0) return; + + set((s) => ({ + selectedModels: { ...s.selectedModels, [providerId]: defaults }, + })); + // Persist + void updateDesktopSettings({ usageSelectedModels: get().selectedModels }); + }, }), { name: 'quota-store' } ) diff --git a/packages/web/server/index.js b/packages/web/server/index.js index a4334c79..e787155d 100644 --- a/packages/web/server/index.js +++ b/packages/web/server/index.js @@ -1128,6 +1128,114 @@ const sanitizeSettingsUpdate = (payload) => { result.skillCatalogs = skillCatalogs; } + // Usage model selections - which models appear in dropdown + if (candidate.usageSelectedModels && typeof candidate.usageSelectedModels === 'object') { + const sanitized = {}; + for (const [providerId, models] of Object.entries(candidate.usageSelectedModels)) { + if (typeof providerId === 'string' && Array.isArray(models)) { + const validModels = models.filter((m) => typeof m === 'string' && m.length > 0); + if (validModels.length > 0) { + sanitized[providerId] = validModels; + } + } + } + if (Object.keys(sanitized).length > 0) { + result.usageSelectedModels = sanitized; + } + } + + // Usage page collapsed families - for "Other Models" section + if (candidate.usageCollapsedFamilies && typeof candidate.usageCollapsedFamilies === 'object') { + const sanitized = {}; + for (const [providerId, families] of Object.entries(candidate.usageCollapsedFamilies)) { + if (typeof providerId === 'string' && Array.isArray(families)) { + const validFamilies = families.filter((f) => typeof f === 'string' && f.length > 0); + if (validFamilies.length > 0) { + sanitized[providerId] = validFamilies; + } + } + } + if (Object.keys(sanitized).length > 0) { + result.usageCollapsedFamilies = sanitized; + } + } + + // Header dropdown expanded families (inverted - stores EXPANDED, default all collapsed) + if (candidate.usageExpandedFamilies && typeof candidate.usageExpandedFamilies === 'object') { + const sanitized = {}; + for (const [providerId, families] of Object.entries(candidate.usageExpandedFamilies)) { + if (typeof providerId === 'string' && Array.isArray(families)) { + const validFamilies = families.filter((f) => typeof f === 'string' && f.length > 0); + if (validFamilies.length > 0) { + sanitized[providerId] = validFamilies; + } + } + } + if (Object.keys(sanitized).length > 0) { + result.usageExpandedFamilies = sanitized; + } + } + + // Custom model groups configuration + if (candidate.usageModelGroups && typeof candidate.usageModelGroups === 'object') { + const sanitized = {}; + for (const [providerId, config] of Object.entries(candidate.usageModelGroups)) { + if (typeof providerId !== 'string') continue; + + const providerConfig = {}; + + // customGroups: array of {id, label, models, order} + if (Array.isArray(config.customGroups)) { + const validGroups = config.customGroups + .filter((g) => g && typeof g.id === 'string' && typeof g.label === 'string') + .map((g) => ({ + id: g.id.slice(0, 64), + label: g.label.slice(0, 128), + models: Array.isArray(g.models) + ? g.models.filter((m) => typeof m === 'string').slice(0, 500) + : [], + order: typeof g.order === 'number' ? g.order : 0, + })); + if (validGroups.length > 0) { + providerConfig.customGroups = validGroups; + } + } + + // modelAssignments: Record + if (config.modelAssignments && typeof config.modelAssignments === 'object') { + const assignments = {}; + for (const [model, groupId] of Object.entries(config.modelAssignments)) { + if (typeof model === 'string' && typeof groupId === 'string') { + assignments[model] = groupId; + } + } + if (Object.keys(assignments).length > 0) { + providerConfig.modelAssignments = assignments; + } + } + + // renamedGroups: Record + if (config.renamedGroups && typeof config.renamedGroups === 'object') { + const renamed = {}; + for (const [groupId, label] of Object.entries(config.renamedGroups)) { + if (typeof groupId === 'string' && typeof label === 'string') { + renamed[groupId] = label.slice(0, 128); + } + } + if (Object.keys(renamed).length > 0) { + providerConfig.renamedGroups = renamed; + } + } + + if (Object.keys(providerConfig).length > 0) { + sanitized[providerId] = providerConfig; + } + } + if (Object.keys(sanitized).length > 0) { + result.usageModelGroups = sanitized; + } + } + return result; };