import { useCallback, useEffect, useMemo, useState } from "react"; import { mergeConnectorModelEntries, type ConnectorModelEntry, } from "@shared/connector-models"; import { isConnectorExecutionReady } from "@shared/connector-policy"; import type { AppLocale } from "@shared/locale"; import type { ConnectorId, ProfileSnapshot } from "@shared/types"; import { unwrap } from "../../lib/desktop"; import type { TranslationFn } from "../../i18n"; interface UseConnectorModelSelectionResult { executionReadyConnectors: ProfileSnapshot["connectors"]; hasAtLeastOneReady: boolean; selectedConnectorId: ConnectorId | ""; setSelectedConnectorId: (value: ConnectorId | "") => void; selectedConnector: ProfileSnapshot["connectors"][number] | null; selectedModel: string; setSelectedModel: (value: string) => void; connectorModels: ConnectorModelEntry[]; connectorModelPools: Array<{ connector: ProfileSnapshot["connectors"][number]; models: ConnectorModelEntry[]; }>; modelOptions: Array<{ id: string; label: string }>; loadingModels: boolean; totalEnabledModels: number; selectedConnectorEnabledModels: number; toggleConnectorModel: (connectorId: ConnectorId, modelId: string) => Promise; } export function useConnectorModelSelection( connectors: ProfileSnapshot["connectors"], currentLocale: AppLocale, t: TranslationFn, setError: (error: string | null) => void, ): UseConnectorModelSelectionResult { const executionReadyConnectors = useMemo( () => connectors.filter((connector) => isConnectorExecutionReady(connector)), [connectors], ); const hasAtLeastOneReady = executionReadyConnectors.length > 0; const [selectedConnectorId, setSelectedConnectorId] = useState( executionReadyConnectors[0]?.id ?? "", ); const [selectedModel, setSelectedModel] = useState(""); const [connectorModelsById, setConnectorModelsById] = useState>({}); const [loadingModels, setLoadingModels] = useState(false); useEffect(() => { if (!selectedConnectorId && executionReadyConnectors[0]) { setSelectedConnectorId(executionReadyConnectors[0].id); return; } if (selectedConnectorId && !executionReadyConnectors.some((connector) => connector.id === selectedConnectorId)) { setSelectedConnectorId(executionReadyConnectors[0]?.id ?? ""); } }, [executionReadyConnectors, selectedConnectorId]); const selectedConnector = executionReadyConnectors.find((connector) => connector.id === selectedConnectorId) ?? executionReadyConnectors[0] ?? null; const loadConnectorModels = useCallback(async (connector: ProfileSnapshot["connectors"][number]) => { if (connector.id.startsWith("claude_compat")) { const configuredModels = (connector.model ?? "") .split(",") .map((model) => model.trim()) .filter(Boolean) .map((model) => ({ id: model, label: model })); const disabledResult = await window.agentCompany.getConnectorDisabledModels(connector.id); const disabledModels = disabledResult.ok ? disabledResult.data : []; return mergeConnectorModelEntries(configuredModels, disabledModels); } const [enabledResult, disabledResult] = await Promise.all([ window.agentCompany.listConnectorModels(connector.id), window.agentCompany.getConnectorDisabledModels(connector.id), ]); const enabledModels = enabledResult.ok ? enabledResult.data : []; const disabledModels = disabledResult.ok ? disabledResult.data : []; return mergeConnectorModelEntries(enabledModels, disabledModels); }, []); useEffect(() => { if (!hasAtLeastOneReady) { setConnectorModelsById({}); setSelectedModel(""); return; } let cancelled = false; setLoadingModels(true); const loadModels = async () => { try { const entries = await Promise.all( executionReadyConnectors.map(async (connector) => [connector.id, await loadConnectorModels(connector)] as const), ); if (cancelled) return; setConnectorModelsById(Object.fromEntries(entries)); } catch { if (cancelled) return; setConnectorModelsById({}); setSelectedModel(""); } finally { if (!cancelled) { setLoadingModels(false); } } }; void loadModels(); return () => { cancelled = true; }; }, [executionReadyConnectors, hasAtLeastOneReady, loadConnectorModels]); const connectorModels = useMemo( () => connectorModelsById[selectedConnectorId] ?? [], [connectorModelsById, selectedConnectorId], ); useEffect(() => { if (!selectedConnectorId) { setSelectedModel(""); return; } setSelectedModel((currentModel) => currentModel && connectorModels.some((entry) => entry.enabled && entry.id === currentModel) ? currentModel : "", ); }, [connectorModels, selectedConnectorId]); const modelOptions = useMemo( () => connectorModels .filter((model) => model.enabled) .map((model) => ({ id: model.id, label: model.label })), [connectorModels], ); const connectorModelPools = useMemo( () => executionReadyConnectors.map((connector) => ({ connector, models: connectorModelsById[connector.id] ?? [], })), [connectorModelsById, executionReadyConnectors], ); const totalEnabledModels = useMemo( () => Object.values(connectorModelsById).reduce( (count, models) => count + models.filter((model) => model.enabled).length, 0, ), [connectorModelsById], ); const selectedConnectorEnabledModels = useMemo( () => connectorModels.filter((model) => model.enabled).length, [connectorModels], ); const toggleConnectorModel = useCallback(async (connectorId: ConnectorId, modelId: string) => { const connectorModels = connectorModelsById[connectorId]; if (!connectorModels) { return; } const currentEntry = connectorModels.find((model) => model.id === modelId); if (!currentEntry) { return; } const updatedModels = connectorModels.map((model) => model.id === modelId ? { ...model, enabled: !model.enabled } : model, ); const updatedModelMap = { ...connectorModelsById, [connectorId]: updatedModels, }; const nextEnabledCount = Object.values(updatedModelMap).reduce( (count, models) => count + models.filter((model) => model.enabled).length, 0, ); if (currentEntry.enabled && nextEnabledCount === 0) { setError( currentLocale === "zh" ? "至少保留一个全局启用模型,CEO 才能继续为团队分配工作。" : "Keep at least one enabled model across all providers so the CEO can keep staffing the team.", ); return; } const previousModels = connectorModelsById; setConnectorModelsById(updatedModelMap); setError(null); try { unwrap( await window.agentCompany.setConnectorDisabledModels({ connectorId, disabledModels: updatedModels .filter((model) => !model.enabled) .map((model) => ({ id: model.id, label: model.label })), }), ); } catch (error) { setConnectorModelsById(previousModels); setError(error instanceof Error ? error.message : t("onboarding.unableToBootstrap")); } }, [connectorModelsById, currentLocale, setError, t]); return { executionReadyConnectors, hasAtLeastOneReady, selectedConnectorId, setSelectedConnectorId, selectedConnector, selectedModel, setSelectedModel, connectorModels, connectorModelPools, modelOptions, loadingModels, totalEnabledModels, selectedConnectorEnabledModels, toggleConnectorModel, }; }