diff --git a/tests/test_downloads_tab_workflow_route.mjs b/tests/test_downloads_tab_workflow_route.mjs index 32a9769..c648a2b 100644 --- a/tests/test_downloads_tab_workflow_route.mjs +++ b/tests/test_downloads_tab_workflow_route.mjs @@ -237,6 +237,62 @@ test('node context integration preserves existing menu hooks and refreshes after assert.deepEqual(calls, ['original-menu', 'original-widget:ckpt_name', 'refresh']); }); +test('queued workflow model selection survives missing data and completes after analysis', () => { + const queueWorkflowModelReferenceSelection = eval(`(${extractMethod(missingBrowserMethodsSource, 'queueWorkflowModelReferenceSelection')})`); + const applyPendingWorkflowModelSelection = eval(`(${extractMethod(missingBrowserMethodsSource, 'applyPendingWorkflowModelSelection')})`); + const selectedModel = { + node_id: 163, + widget_index: 1, + original_path: 'FLUX/ae.safetensors', + }; + let analysisReady = false; + const dialog = { + pendingWorkflowModelSelection: null, + queueWorkflowModelReferenceSelection, + selectWorkflowModelReference() { + return analysisReady ? selectedModel : null; + }, + }; + + const request = queueWorkflowModelReferenceSelection.call(dialog, selectedModel); + + assert.equal(request.status, 'pending'); + assert.equal(applyPendingWorkflowModelSelection.call(dialog, {}), null); + assert.equal(dialog.pendingWorkflowModelSelection, request); + + analysisReady = true; + assert.equal(applyPendingWorkflowModelSelection.call(dialog, {}), selectedModel); + assert.equal(request.status, 'selected'); + assert.equal(request.selected, selectedModel); + assert.equal(dialog.pendingWorkflowModelSelection, null); +}); + +test('a newer workflow model selection supersedes the previous pending request', () => { + const queueWorkflowModelReferenceSelection = eval(`(${extractMethod(missingBrowserMethodsSource, 'queueWorkflowModelReferenceSelection')})`); + const dialog = { pendingWorkflowModelSelection: null }; + + const first = queueWorkflowModelReferenceSelection.call(dialog, { node_id: 1 }); + const second = queueWorkflowModelReferenceSelection.call(dialog, { node_id: 2 }); + + assert.equal(first.status, 'superseded'); + assert.equal(second.status, 'pending'); + assert.equal(dialog.pendingWorkflowModelSelection, second); +}); + +test('workflow model selection wait resolves only after the queued request completes', async () => { + const waitForWorkflowModelSelection = eval(`(${extractMethod(modelResolverSource, 'waitForWorkflowModelSelection')})`); + const request = { status: 'pending' }; + const resolver = { dialog: { pendingWorkflowModelSelection: request } }; + + setTimeout(() => { + request.status = 'selected'; + }, 20); + + const completed = await waitForWorkflowModelSelection.call(resolver, request, 500); + assert.equal(completed, request); + assert.equal(completed.status, 'selected'); +}); + test('Local Database source ignores installed local model matches before search', () => { const hasMissingSourceSearchAttempt = eval(`(${extractMethod(missingBrowserMethodsSource, 'hasMissingSourceSearchAttempt')})`); const getMissingSourceResultStatus = eval(`(${extractMethod(missingBrowserMethodsSource, 'getMissingSourceResultStatus')})`); diff --git a/web/resolver/model_resolver.js b/web/resolver/model_resolver.js index b56f58a..2d92a83 100644 --- a/web/resolver/model_resolver.js +++ b/web/resolver/model_resolver.js @@ -483,9 +483,33 @@ export class ModelResolver { }); } + waitForWorkflowModelSelection(request, timeoutMs = 120000) { + const startedAt = Date.now(); + return new Promise((resolve) => { + const check = () => { + if (!request || request.status !== 'pending') { + resolve(request || null); + return; + } + if (Date.now() - startedAt >= timeoutMs) { + request.status = 'timeout'; + if (this.dialog?.pendingWorkflowModelSelection === request) { + this.dialog.pendingWorkflowModelSelection = null; + } + resolve(request); + return; + } + setTimeout(check, 50); + }; + check(); + }); + } + async showResolvedNodeModelInResolver(reference) { if (!this.dialog) return; + const wasVisible = Boolean(this.dialog.isVisible?.()); + const wasOnMissingTab = this.dialog.activeTab === 'missing'; const { signature } = this.getNodeContextWorkflowState(); const analysisData = this.getCurrentNodeContextAnalysis(signature); if (signature && analysisData) { @@ -493,11 +517,11 @@ export class ModelResolver { this.dialog.cachedAnalysisData = analysisData; } - this.dialog.activeTab = 'missing'; this.dialog.persistActiveTab?.('missing'); this.dialog.showResolvedModels = true; this.dialog.missingModelsTypeFilter = 'all'; this.dialog.missingModelsTypeFilterMenuOpen = false; + const selectionRequest = this.dialog.queueWorkflowModelReferenceSelection?.(reference); try { localStorage.setItem(this.dialog.showResolvedModelsStorageKey, '1'); @@ -505,30 +529,30 @@ export class ModelResolver { log.debug('Model Resolver: failed to persist resolved-model visibility.', error); } - if (!this.dialog.isVisible?.()) { + if (!wasVisible) { + this.dialog.activeTab = 'missing'; this.activateResolverButton(); } if (!await this.waitForResolverDialogReady()) { + if (selectionRequest?.status === 'pending') { + selectionRequest.status = 'open-failed'; + } + if (this.dialog.pendingWorkflowModelSelection === selectionRequest) { + this.dialog.pendingWorkflowModelSelection = null; + } showNotification('Could not open Model Resolver.', 'error'); return; } - this.dialog.switchTab?.('missing', { force: true }); - await this.dialog.loadWorkflowData?.(); - - let selected = this.dialog.selectWorkflowModelReference?.( - reference, - this.dialog.cachedAnalysisData - ); - if (!selected) { - await this.dialog.loadWorkflowData?.(null, { force: true, forceRescan: false }); - selected = this.dialog.selectWorkflowModelReference?.( - reference, - this.dialog.cachedAnalysisData - ); + if (wasVisible && !wasOnMissingTab) { + await this.dialog.switchTab?.('missing', { force: true }); + } else if (wasVisible) { + await this.dialog.loadWorkflowData?.(); } - if (!selected) { - showNotification('The selected model is no longer present in the current workflow.', 'warning'); + + const completedRequest = await this.waitForWorkflowModelSelection(selectionRequest); + if (completedRequest?.status === 'timeout') { + showNotification('Could not select the model after workflow analysis completed.', 'warning'); } } diff --git a/web/resolver/resolver_dialog.js b/web/resolver/resolver_dialog.js index 0043ade..df6bbda 100644 --- a/web/resolver/resolver_dialog.js +++ b/web/resolver/resolver_dialog.js @@ -54,6 +54,7 @@ export class ResolverManagerDialog extends ComfyDialog { this.cachedLoadedModelsData = null; this.cachedLoadedModelsSignature = null; this.selectedMissingModelKey = null; + this.pendingWorkflowModelSelection = null; this.batchSelectedMissingKeys = new Set(); this.lastBatchSelectedMissingKey = null; this.activeFooterMenu = null; diff --git a/web/resolver/shell/lifecycle_graph_methods.js b/web/resolver/shell/lifecycle_graph_methods.js index 2db1ee6..fc0d618 100644 --- a/web/resolver/shell/lifecycle_graph_methods.js +++ b/web/resolver/shell/lifecycle_graph_methods.js @@ -157,9 +157,10 @@ export const lifecycleGraphMethods = { if (shouldRenderMissingModels()) { await this.ensureDownloadDirectoriesLoaded(); this.displayMissingModels(this.contentElement, this.cachedAnalysisData); + this.applyPendingWorkflowModelSelection?.(this.cachedAnalysisData); this.reconnectActiveDownloads(); } - return; + return this.cachedAnalysisData; } if (!this.workflowHasNodes(workflow)) { @@ -179,8 +180,9 @@ export const lifecycleGraphMethods = { if (shouldRenderMissingModels()) { await this.ensureDownloadDirectoriesLoaded(); this.displayMissingModels(this.contentElement, data); + this.applyPendingWorkflowModelSelection?.(data); } - return; + return data; } const analysisId = `an-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; @@ -217,10 +219,12 @@ export const lifecycleGraphMethods = { if (shouldRenderMissingModels()) { await this.ensureDownloadDirectoriesLoaded(); this.displayMissingModels(this.contentElement, data); + this.applyPendingWorkflowModelSelection?.(data); // Reconnect any active downloads to their new progress divs this.reconnectActiveDownloads(); } + return data; } catch (error) { if (!loadToken || this._workflowDataLoadToken === loadToken) { @@ -232,6 +236,7 @@ export const lifecycleGraphMethods = { } else if (!loadToken && this.activeTab === 'missing' && this.contentElement) { this.contentElement.innerHTML = `

Error: ${error.message}

`; } + return null; } }, diff --git a/web/resolver/views/missing_browser_methods.js b/web/resolver/views/missing_browser_methods.js index c8792c3..ec30d68 100644 --- a/web/resolver/views/missing_browser_methods.js +++ b/web/resolver/views/missing_browser_methods.js @@ -508,6 +508,36 @@ export const missingBrowserMethods = { }; }, + queueWorkflowModelReferenceSelection(reference = {}) { + if (this.pendingWorkflowModelSelection?.status === 'pending') { + this.pendingWorkflowModelSelection.status = 'superseded'; + } + + const request = { + id: `workflow-model-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`, + reference: { ...reference }, + status: 'pending', + selected: null + }; + this.pendingWorkflowModelSelection = request; + return request; + }, + + applyPendingWorkflowModelSelection(data = this.cachedAnalysisData || {}) { + const request = this.pendingWorkflowModelSelection; + if (!request || request.status !== 'pending') return null; + + const selected = this.selectWorkflowModelReference(request.reference, data); + if (!selected) return null; + + request.status = 'selected'; + request.selected = selected; + if (this.pendingWorkflowModelSelection === request) { + this.pendingWorkflowModelSelection = null; + } + return selected; + }, + selectWorkflowModelReference(reference = {}, data = this.cachedAnalysisData || {}) { const selected = this.getResolvedWorkflowModels(data) .find(model => matchesWorkflowModelReference(model, reference)); diff --git a/web/resolver/views/tabs_loaded_methods.js b/web/resolver/views/tabs_loaded_methods.js index 9aeea39..48ca0e7 100644 --- a/web/resolver/views/tabs_loaded_methods.js +++ b/web/resolver/views/tabs_loaded_methods.js @@ -59,13 +59,13 @@ export const tabsLoadedMethods = { this.contentElement.style.overflowY = 'auto'; } this.setMissingFooterControlsVisible(true); - this.loadWorkflowData(); + return this.loadWorkflowData(); } else if (this.activeTab === 'loaded') { if (this.contentElement) { this.contentElement.style.overflowY = 'auto'; } this.setMissingFooterControlsVisible(false); - this.loadLoadedModels(); + return this.loadLoadedModels(); } else { if (this.contentElement) { this.contentElement.style.overflowY = 'hidden'; @@ -73,6 +73,7 @@ export const tabsLoadedMethods = { this.setMissingFooterControlsVisible(false); this.displayOptions(); } + return null; }, async loadLoadedModels(workflow = null, { force = false } = {}) {