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; };