diff --git a/src/components/sidebar/tabs/modelLibrary/DownloadItem.vue b/src/components/sidebar/tabs/modelLibrary/DownloadItem.vue index 6f855ede17..a121b5d8f5 100644 --- a/src/components/sidebar/tabs/modelLibrary/DownloadItem.vue +++ b/src/components/sidebar/tabs/modelLibrary/DownloadItem.vue @@ -3,7 +3,7 @@
{{ getDownloadLabel(download.savePath ?? '') }}
-
+
@@ -28,7 +28,7 @@ /> ' + } + }) +) + +vi.mock( + '@/platform/missingModel/composables/useMissingModelInteractions', + () => ({ + getModelStateKey: ( + modelName: string, + directory: string | null, + isAssetSupported: boolean + ) => + `${isAssetSupported ? 'supported' : 'unsupported'}::${directory ?? ''}::${modelName}`, + getNodeDisplayLabel: (nodeId: string | number) => `Node ${nodeId}`, + getComboValue: () => undefined, + useMissingModelInteractions: () => { + const store = useMissingModelStore() + + return { + toggleModelExpand: vi.fn(), + isModelExpanded: () => false, + getComboOptions: () => [], + handleComboSelect: (key: string, value: string | undefined) => { + if (value) { + store.selectedLibraryModel[key] = value + } + }, + isSelectionConfirmable: () => false, + cancelLibrarySelect: vi.fn(), + confirmLibrarySelect: vi.fn(), + getTypeMismatch: () => null, + getDownloadStatus: (key: string) => + store.downloadRefs[key]?.kind === 'electron-download' + ? { progress: 0, status: 'created' as const } + : null + } + } + }) +) + +vi.mock('@/platform/missingModel/missingModelDownload', () => ({ + downloadModel: (...args: unknown[]) => mockDownloadModel(...args), + fetchModelMetadata: (...args: unknown[]) => mockFetchModelMetadata(...args), + isModelDownloadable: () => true, + toBrowsableUrl: (url: string) => url +})) + +import MissingModelRow from './MissingModelRow.vue' + +const i18n = createI18n({ + legacy: false, + locale: 'en', + messages: { + en: { + g: { + download: 'Download' + }, + rightSidePanel: { + missingModels: { + copyModelName: 'Copy model name', + copyUrl: 'Copy URL', + confirmSelection: 'Confirm selection', + collapseNodes: 'Collapse nodes', + expandNodes: 'Expand nodes' + } + } + } + }, + missingWarn: false, + fallbackWarn: false +}) + +const model: MissingModelViewModel = { + name: 'z_image_turbo_bf16.safetensors', + representative: { + name: 'z_image_turbo_bf16.safetensors', + url: 'https://example.com/z_image_turbo_bf16.safetensors', + directory: 'checkpoints', + nodeId: '1', + nodeType: 'CheckpointLoaderSimple', + widgetName: 'ckpt_name', + isAssetSupported: false, + isMissing: true + }, + referencingNodes: [{ nodeId: '1', widgetName: 'ckpt_name' }] +} + +const modelKey = 'unsupported::checkpoints::z_image_turbo_bf16.safetensors' + +function renderComponent() { + return render(MissingModelRow, { + props: { + model, + directory: 'checkpoints', + showNodeIdBadge: false, + isAssetSupported: false + }, + global: { + plugins: [i18n] + } + }) +} + +describe('MissingModelRow', () => { + beforeEach(() => { + setActivePinia(createPinia()) + mockDownloadModel.mockReset() + mockDownloadModel.mockResolvedValue(true) + mockFetchModelMetadata.mockReset() + mockFetchModelMetadata.mockResolvedValue({ + fileSize: null, + gatedRepoUrl: null + }) + mockCopyToClipboard.mockReset() + + const store = useMissingModelStore() + store.folderPaths = { + checkpoints: ['/models/checkpoints'] + } + }) + + it('tracks and surfaces direct Electron downloads immediately after the button is clicked', async () => { + const user = userEvent.setup() + const store = useMissingModelStore() + renderComponent() + + await user.click(screen.getByTestId('missing-model-download')) + + expect(mockDownloadModel).toHaveBeenCalledWith( + { + name: model.representative.name, + url: model.representative.url, + directory: model.representative.directory + }, + store.folderPaths + ) + expect(store.downloadRefs[modelKey]).toEqual({ + kind: 'electron-download', + url: model.representative.url + }) + expect(store.selectedLibraryModel[modelKey]).toBe(model.representative.name) + expect(screen.getByTestId('missing-model-status-card')).toHaveTextContent( + 'created' + ) + }) + + it('does not create UI state when the Electron download does not start', async () => { + mockDownloadModel.mockResolvedValue(false) + const user = userEvent.setup() + const store = useMissingModelStore() + renderComponent() + + await user.click(screen.getByTestId('missing-model-download')) + + expect(store.downloadRefs[modelKey]).toBeUndefined() + expect(store.selectedLibraryModel[modelKey]).toBeUndefined() + expect( + screen.queryByTestId('missing-model-status-card') + ).not.toBeInTheDocument() + }) + + it('clears stale download refs when the user picks a library alternative', async () => { + const user = userEvent.setup() + const store = useMissingModelStore() + store.downloadRefs[modelKey] = { + kind: 'electron-download', + url: model.representative.url! + } + + renderComponent() + + await user.click(screen.getByTestId('missing-model-library-select')) + + expect(store.downloadRefs[modelKey]).toBeUndefined() + expect(store.selectedLibraryModel[modelKey]).toBe( + 'library-model.safetensors' + ) + }) +}) diff --git a/src/platform/missingModel/components/MissingModelRow.vue b/src/platform/missingModel/components/MissingModelRow.vue index a39e92652d..f941725776 100644 --- a/src/platform/missingModel/components/MissingModelRow.vue +++ b/src/platform/missingModel/components/MissingModelRow.vue @@ -175,7 +175,7 @@ :model-value="getComboValue(model.representative)" :options="comboOptions" :show-divider="isAssetSupported || downloadable" - @select="handleComboSelect(modelKey, $event)" + @select="handleLibraryModelSelect" />
@@ -284,18 +284,31 @@ const downloadLabel = computed(() => { return size ? `${base} (${formatSize(size)})` : base }) -function handleDownload() { +async function handleDownload() { const rep = model.representative if (rep.url && rep.directory) { - downloadModel( + const started = await downloadModel( { name: rep.name, url: rep.url, directory: rep.directory }, store.folderPaths ) + + if (started) { + store.downloadRefs[modelKey.value] = { + kind: 'electron-download', + url: rep.url + } + handleComboSelect(modelKey.value, rep.name) + } } else { console.warn('[MissingModelRow] Cannot download: missing url or directory') } } +function handleLibraryModelSelect(value: string | undefined) { + delete store.downloadRefs[modelKey.value] + handleComboSelect(modelKey.value, value) +} + const { toggleModelExpand, isModelExpanded, diff --git a/src/platform/missingModel/components/MissingModelStatusCard.vue b/src/platform/missingModel/components/MissingModelStatusCard.vue index 7c41445675..8b89001dd3 100644 --- a/src/platform/missingModel/components/MissingModelStatusCard.vue +++ b/src/platform/missingModel/components/MissingModelStatusCard.vue @@ -22,11 +22,21 @@ aria-hidden="true" class="icon-[lucide--circle-alert] size-5 text-destructive-background" /> +