diff --git a/frontend/libs/components/model-selector/src/lib/classy-model-selector-store.spec.ts b/frontend/libs/components/model-selector/src/lib/classy-model-selector-store.spec.ts new file mode 100644 index 0000000000..a349c835d8 --- /dev/null +++ b/frontend/libs/components/model-selector/src/lib/classy-model-selector-store.spec.ts @@ -0,0 +1,160 @@ +import type { Model } from "@storyteller/model-list"; +import { beforeEach, describe, expect, it } from "vitest"; +import { useClassyModelSelectorStore } from "./classy-model-selector-store"; +import { ModelPage } from "./model-pages"; + +const PAGE = ModelPage.ImageTo3DObject; + +describe("model selection catalog reconciliation", () => { + beforeEach(() => { + useClassyModelSelectorStore.setState({ + selectedModels: {}, + modelSelectionSources: {}, + selectedProviders: {}, + }); + }); + + it("promotes an automatic fallback to the hydrated catalog default", () => { + const bootstrapFallback = fakeModel("bootstrap-fallback", "static"); + const hydratedFallback = fakeModel("bootstrap-fallback", "hydrated"); + const serverDefault = fakeModel("server-default", "hydrated"); + const store = useClassyModelSelectorStore.getState(); + + store.reconcileSelectedModelFromCatalog( + PAGE, + [bootstrapFallback], + bootstrapFallback, + ); + store.reconcileSelectedModelFromCatalog( + PAGE, + [serverDefault, hydratedFallback], + serverDefault, + ); + + const state = useClassyModelSelectorStore.getState(); + expect(state.selectedModels[PAGE]).toBe(serverDefault); + expect(state.modelSelectionSources[PAGE]).toBe("automatic"); + }); + + it("preserves an explicit choice that is absent from the hydrated catalog", () => { + const bootstrapFallback = fakeModel("bootstrap-fallback", "static"); + const explicitChoice = fakeModel("artist-choice", "static"); + const serverDefault = fakeModel("server-default", "hydrated"); + const store = useClassyModelSelectorStore.getState(); + + store.reconcileSelectedModelFromCatalog( + PAGE, + [bootstrapFallback], + bootstrapFallback, + ); + store.setSelectedModel(PAGE, explicitChoice); + store.reconcileSelectedModelFromCatalog( + PAGE, + [serverDefault], + serverDefault, + ); + + const state = useClassyModelSelectorStore.getState(); + expect(state.selectedModels[PAGE]).toBe(explicitChoice); + expect(state.modelSelectionSources[PAGE]).toBe("explicit"); + }); + + it("refreshes an explicit same-id instance without losing its provenance", () => { + const explicitChoice = fakeModel("artist-choice", "static-capabilities"); + const hydratedChoice = fakeModel("artist-choice", "hydrated-capabilities"); + const serverDefault = fakeModel("server-default", "hydrated"); + const store = useClassyModelSelectorStore.getState(); + + store.setSelectedModel(PAGE, explicitChoice); + store.reconcileSelectedModelFromCatalog( + PAGE, + [serverDefault, hydratedChoice], + serverDefault, + ); + + const state = useClassyModelSelectorStore.getState(); + expect(state.selectedModels[PAGE]).toBe(hydratedChoice); + expect(state.selectedModels[PAGE]?.selectorDescription).toBe( + "hydrated-capabilities", + ); + expect(state.modelSelectionSources[PAGE]).toBe("explicit"); + }); + + it("refreshes an automatic same-id instance after hydration", () => { + const bootstrapDefault = fakeModel( + "catalog-default", + "static-capabilities", + ); + const hydratedDefault = fakeModel( + "catalog-default", + "hydrated-capabilities", + ); + const store = useClassyModelSelectorStore.getState(); + + store.reconcileSelectedModelFromCatalog( + PAGE, + [bootstrapDefault], + bootstrapDefault, + ); + store.reconcileSelectedModelFromCatalog( + PAGE, + [hydratedDefault], + hydratedDefault, + ); + + const state = useClassyModelSelectorStore.getState(); + expect(state.selectedModels[PAGE]).toBe(hydratedDefault); + expect(state.modelSelectionSources[PAGE]).toBe("automatic"); + }); + + it("keeps the current selection when the catalog is empty", () => { + const explicitChoice = fakeModel("artist-choice", "static"); + const store = useClassyModelSelectorStore.getState(); + + store.setSelectedModel(PAGE, explicitChoice); + store.reconcileSelectedModelFromCatalog(PAGE, [], undefined); + + const state = useClassyModelSelectorStore.getState(); + expect(state.selectedModels[PAGE]).toBe(explicitChoice); + expect(state.modelSelectionSources[PAGE]).toBe("explicit"); + }); + + it("keeps an automatic fallback when hydration yields no models", () => { + const bootstrapFallback = fakeModel("bootstrap-fallback", "static"); + const store = useClassyModelSelectorStore.getState(); + + store.reconcileSelectedModelFromCatalog( + PAGE, + [bootstrapFallback], + bootstrapFallback, + ); + store.reconcileSelectedModelFromCatalog(PAGE, [], undefined); + + const state = useClassyModelSelectorStore.getState(); + expect(state.selectedModels[PAGE]).toBe(bootstrapFallback); + expect(state.modelSelectionSources[PAGE]).toBe("automatic"); + }); + + it("uses an available fallback when the requested default is unavailable", () => { + const firstAvailable = fakeModel("first-available", "hydrated"); + const unavailableDefault = fakeModel("unavailable-default", "static"); + const store = useClassyModelSelectorStore.getState(); + + store.reconcileSelectedModelFromCatalog( + PAGE, + [firstAvailable], + unavailableDefault, + ); + + expect(useClassyModelSelectorStore.getState().selectedModels[PAGE]).toBe( + firstAvailable, + ); + }); +}); + +const fakeModel = (tauriId: string, selectorDescription: string): Model => + ({ + id: `ui-${tauriId}`, + tauriId, + selectorDescription, + }) as unknown as Model; diff --git a/frontend/libs/components/model-selector/src/lib/classy-model-selector-store.ts b/frontend/libs/components/model-selector/src/lib/classy-model-selector-store.ts index 266eb7d87e..b511bfc097 100644 --- a/frontend/libs/components/model-selector/src/lib/classy-model-selector-store.ts +++ b/frontend/libs/components/model-selector/src/lib/classy-model-selector-store.ts @@ -3,10 +3,20 @@ import { ModelPage } from "./model-pages"; import { ImageModel, Model, VideoModel } from "@storyteller/model-list"; import { GenerationProvider } from "@storyteller/api-enums"; +type ModelSelectionSource = "automatic" | "explicit"; + interface ClassyModelSelectorState { selectedModels: { [page in ModelPage]?: Model }; + modelSelectionSources: { + [page in ModelPage]?: ModelSelectionSource; + }; selectedProviders: { [page in ModelPage]?: { [modelId: string]: GenerationProvider } }; setSelectedModel: (page: ModelPage, model: Model) => void; + reconcileSelectedModelFromCatalog: ( + page: ModelPage, + models: Model[], + defaultModel: Model | undefined + ) => void; setSelectedProvider: ( page: ModelPage, modelId: string, @@ -17,6 +27,7 @@ interface ClassyModelSelectorState { export const useClassyModelSelectorStore = create( (set) => ({ selectedModels: {}, + modelSelectionSources: {}, selectedProviders: {}, setSelectedModel: (page, model) => set((state) => ({ @@ -24,7 +35,45 @@ export const useClassyModelSelectorStore = create( ...state.selectedModels, [page]: model, }, + modelSelectionSources: { + ...state.modelSelectionSources, + [page]: "explicit", + }, })), + reconcileSelectedModelFromCatalog: (page, models, defaultModel) => + set((state) => { + if (models.length === 0) return state; + + const selected = state.selectedModels[page]; + const refreshed = selected + ? models.find((model) => model.tauriId === selected.tauriId) + : undefined; + const availableDefault = defaultModel + ? (models.find((model) => model.tauriId === defaultModel.tauriId) ?? + models[0]) + : models[0]; + const automatic = state.modelSelectionSources[page] === "automatic"; + const nextSelected = selected + ? automatic + ? (availableDefault ?? refreshed ?? selected) + : (refreshed ?? selected) + : availableDefault; + + if (!nextSelected || nextSelected === selected) return state; + + return { + selectedModels: { + ...state.selectedModels, + [page]: nextSelected, + }, + modelSelectionSources: selected + ? state.modelSelectionSources + : { + ...state.modelSelectionSources, + [page]: "automatic", + }, + }; + }), setSelectedProvider: (page, modelId, provider) => set((state) => ({ selectedProviders: { diff --git a/frontend/libs/components/model-selector/src/lib/classy-model-selector.tsx b/frontend/libs/components/model-selector/src/lib/classy-model-selector.tsx index 288d933fac..fcf5eb9c50 100644 --- a/frontend/libs/components/model-selector/src/lib/classy-model-selector.tsx +++ b/frontend/libs/components/model-selector/src/lib/classy-model-selector.tsx @@ -20,7 +20,10 @@ import { import { ChevronUpIcon, CircleCheckIcon } from "lucide-react"; import { GenerationProvider } from "@storyteller/api-enums"; import { Tooltip } from "@storyteller/ui-tooltip"; -import { defaultModelForPage } from "./defaultModelForPage"; +import { + availableDefaultModelForPage, + defaultModelForPage, +} from "./defaultModelForPage"; interface ClassyModelSelectorProps { items: Omit[]; @@ -139,39 +142,43 @@ export function ClassyModelSelector({ showProviderSelection = true, ...popoverProps }: ClassyModelSelectorProps) { - const { selectedModels, setSelectedModel, setSelectedProvider } = - useClassyModelSelectorStore(); - const itemModels: Model[] = items - .map((item) => item.model) - .filter((model) => model !== undefined); + const { + selectedModels, + setSelectedModel, + reconcileSelectedModelFromCatalog, + setSelectedProvider, + } = useClassyModelSelectorStore(); + const itemModels = useMemo( + () => + items + .map((item) => item.model) + .filter((model): model is Model => model !== undefined), + [items], + ); + const storedSelectedModel = selectedModels[page]; const selectedModel = - selectedModels[page] || defaultModelForPage(itemModels, page); + storedSelectedModel || defaultModelForPage(itemModels, page); const selectedProvider = useSelectedProviderForModel(page, selectedModel?.id); const selectedProvidersByModel = useClassyModelSelectorStore( (s) => s.selectedProviders[page] ?? {}, ); - // For the first mount, make sure the selected model is set for other components to listen + // Seed an automatic default for downstream consumers. When the asynchronous + // backend catalog arrives, automatic selections may follow its available + // default while explicit choices keep their id. Matching selections still + // receive the rebuilt instance so capability-driven UI uses hydrated data. useEffect(() => { - // Initialize selected model if not set - if (!selectedModels[page] && items[0]) { - setSelectedModel(page, defaultModelForPage(itemModels, page)); - } - }, []); - - // The backend listing hydrates asynchronously and rebuilds the model - // instances with API capabilities. Swap a stale selected instance for the - // fresh one so capability-driven UI (keyframes, references, pickers) - // reflects the API data without needing a manual re-select. - useEffect(() => { - const selected = selectedModels[page]; - if (!selected) return; - const fresh = itemModels.find((m) => m.tauriId === selected.tauriId); - if (fresh !== undefined && fresh !== selected) { - setSelectedModel(page, fresh); - } - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [items, selectedModels, page]); + reconcileSelectedModelFromCatalog( + page, + itemModels, + availableDefaultModelForPage(itemModels, page), + ); + }, [ + itemModels, + page, + reconcileSelectedModelFromCatalog, + storedSelectedModel, + ]); // Initialize a default provider for each model so we can render icons even when not selected useEffect(() => { diff --git a/frontend/libs/components/model-selector/src/lib/defaultModelForPage.ts b/frontend/libs/components/model-selector/src/lib/defaultModelForPage.ts index 60aca7a6bd..73666ffb2a 100644 --- a/frontend/libs/components/model-selector/src/lib/defaultModelForPage.ts +++ b/frontend/libs/components/model-selector/src/lib/defaultModelForPage.ts @@ -16,6 +16,17 @@ const DEFAULT_MODEL_ID_FOR_PAGE: Partial> = { [ModelPage.Angles]: "flux_2_lora_angles", }; +export const availableDefaultModelForPage = ( + models: Model[], + page: ModelPage, +): Model | undefined => { + const defaultId = DEFAULT_MODEL_ID_FOR_PAGE[page]; + return ( + models.find((m) => m.id === defaultId || m.tauriId === defaultId) ?? + models[0] + ); +}; + export const defaultModelForPage = ( models: Model[], page: ModelPage, diff --git a/frontend/libs/components/pagescene/src/lib/Stage3DBody.tsx b/frontend/libs/components/pagescene/src/lib/Stage3DBody.tsx index 1ad4e645aa..8bda5653e7 100644 --- a/frontend/libs/components/pagescene/src/lib/Stage3DBody.tsx +++ b/frontend/libs/components/pagescene/src/lib/Stage3DBody.tsx @@ -22,7 +22,7 @@ import { import { useStage3dPageModelList, ModelPage, - defaultModelForPage, + availableDefaultModelForPage, useClassyModelSelectorStore, useSelectedImageModel, useSelectedProviderForModel, @@ -171,24 +171,27 @@ export const Stage3DBody = ({ const setSelectedModel = useClassyModelSelectorStore( (s) => s.setSelectedModel, ); + const reconcileSelectedModelFromCatalog = useClassyModelSelectorStore( + (s) => s.reconcileSelectedModelFromCatalog, + ); - // Seed the default model on mount when we're the only model picker - // on the page. ClassyModelSelector does this itself on mount, but in - // the prompt-box placement we don't render it — so without this - // effect the store stays empty and the trigger has no icon until the - // user opens the popover and picks a model manually. + // Reconcile the default when we're the only model picker on the page. + // ClassyModelSelector normally owns this catalog-hydration step, but the + // prompt-box placement renders its picker directly. useEffect(() => { if (modelSelectorPlacement !== "prompt-box") return; - if (selectedImageModel) return; const models = stage3dModelList.map((i) => i.model).filter( (m): m is NonNullable => m !== undefined, ); - const def = defaultModelForPage(models, PAGE_ID); - if (def) setSelectedModel(PAGE_ID, def); + reconcileSelectedModelFromCatalog( + PAGE_ID, + models, + availableDefaultModelForPage(models, PAGE_ID), + ); }, [ modelSelectorPlacement, + reconcileSelectedModelFromCatalog, selectedImageModel, - setSelectedModel, stage3dModelList, ]);