diff --git a/tests/test_downloads_tab_workflow_route.mjs b/tests/test_downloads_tab_workflow_route.mjs index 2709452..28464db 100644 --- a/tests/test_downloads_tab_workflow_route.mjs +++ b/tests/test_downloads_tab_workflow_route.mjs @@ -4913,6 +4913,80 @@ test('downloads tab shows active downloads from all workflow tabs', () => { ]); }); +test('browser refresh restores active backend downloads with their workflow context', async () => { + const restoreActiveDownloadsFromBackend = eval( + `(${extractMethod(resolveDownloadMethodsSource, 'restoreActiveDownloadsFromBackend')})` + ); + const progress = { + status: 'downloading', + progress: 42, + downloaded: 1024, + total_size: 2048, + filename: 'model.safetensors', + path: 'vae/model.safetensors', + directory: 'vae', + url: 'https://huggingface.co/example/repo/resolve/main/model.safetensors', + download_backend: 'huggingface_xet', + }; + const polled = []; + const removed = []; + const dialog = { + activeDownloads: {}, + _activeDownloadsRestorePromise: null, + fetchJson: async endpoint => { + assert.equal(endpoint, '/model_resolver/progress'); + return { + 'download-1': progress, + 'download-completed': { status: 'completed', filename: 'done.safetensors' }, + }; + }, + loadActiveDownloadRecovery() { + return { + 'download-1': { + missing: { + missing_key: 'workflow-model-1', + original_path: 'vae/original.safetensors', + node_id: 7, + widget_index: 0, + category: 'vae', + }, + workflowLabel: 'Video workflow', + workflowRouteKey: '#video-workflow', + sourceUrl: progress.url, + }, + stale: { filename: 'stale.safetensors' }, + }; + }, + isDownloadProgressStatus(status) { + return status === 'downloading'; + }, + rememberDownloadUiState(downloadId, info, nextProgress) { + info.lastProgress = nextProgress; + info.lastStatus = nextProgress.status; + return { downloadId, progress: nextProgress }; + }, + persistActiveDownloadRecovery() {}, + removeActiveDownloadRecovery(downloadId) { + removed.push(downloadId); + }, + pollDownloadProgress(downloadId) { + polled.push(downloadId); + }, + rebindActiveDownloadMissingModels() {}, + updateDownloadAllButtonState() {}, + updateQueuePanel() {}, + }; + + const restored = await restoreActiveDownloadsFromBackend.call(dialog); + + assert.equal(restored, 1); + assert.deepEqual(polled, ['download-1']); + assert.deepEqual(removed, ['stale']); + assert.equal(dialog.activeDownloads['download-1'].missing.missing_key, 'workflow-model-1'); + assert.equal(dialog.activeDownloads['download-1'].workflowLabel, 'Video workflow'); + assert.equal(dialog.activeDownloads['download-1'].lastProgress.progress, 42); +}); + test('downloads tab renders workflow label for active downloads', () => { const renderQueueDownloadsHtml = eval(`(${extractMethod(queueMethodsSource, 'renderQueueDownloadsHtml')})`); const getDownloadProgressPresentation = eval( diff --git a/tests/test_queue_storage_methods.mjs b/tests/test_queue_storage_methods.mjs index c0b1a5a..619c931 100644 --- a/tests/test_queue_storage_methods.mjs +++ b/tests/test_queue_storage_methods.mjs @@ -51,6 +51,45 @@ test('queue storage persists collapsed state and splitter width immediately', () } }); +test('queue storage persists active download recovery without transient DOM state', () => { + const storage = createStorage(); + const restoreStorage = installStorage(storage); + + try { + const context = { + ...queueStorageMethods, + activeDownloadsStorageKey: 'active-download-recovery', + }; + const progressDiv = { nodeType: 1 }; + const downloadBtn = { nodeType: 1 }; + context.persistActiveDownloadRecovery('download-1', { + missing: { + missing_key: 'model-1', + original_path: 'vae/model.safetensors', + matches: [{ path: 'old-match' }], + }, + progressDiv, + downloadBtn, + category: 'vae', + filename: 'model.safetensors', + sourceUrl: 'https://huggingface.co/example/model/resolve/main/model.safetensors', + workflowKey: 'workflow-1', + }); + + const persisted = JSON.parse(storage.getItem('active-download-recovery')); + assert.equal(persisted['download-1'].filename, 'model.safetensors'); + assert.equal(persisted['download-1'].missing.missing_key, 'model-1'); + assert.equal('matches' in persisted['download-1'].missing, false); + assert.equal('progressDiv' in persisted['download-1'], false); + assert.equal('downloadBtn' in persisted['download-1'], false); + + context.removeActiveDownloadRecovery('download-1'); + assert.deepEqual(JSON.parse(storage.getItem('active-download-recovery')), {}); + } finally { + restoreStorage(); + } +}); + test('download history replaces duplicate identities and keeps newest entries first', () => { const storage = createStorage(); const restoreStorage = installStorage(storage); diff --git a/web/resolver/actions/queue_storage_methods.js b/web/resolver/actions/queue_storage_methods.js index fad38e1..d875522 100644 --- a/web/resolver/actions/queue_storage_methods.js +++ b/web/resolver/actions/queue_storage_methods.js @@ -1,4 +1,86 @@ export const queueStorageMethods = { + getActiveDownloadRecoveryStorageKey() { + return this.activeDownloadsStorageKey || 'model_resolver_active_downloads'; + }, + + loadActiveDownloadRecovery() { + if (this._activeDownloadRecoveryLoaded) { + return this._activeDownloadRecovery || {}; + } + + this._activeDownloadRecoveryLoaded = true; + try { + const raw = localStorage.getItem(this.getActiveDownloadRecoveryStorageKey()); + const parsed = raw ? JSON.parse(raw) : {}; + this._activeDownloadRecovery = parsed && typeof parsed === 'object' && !Array.isArray(parsed) + ? parsed + : {}; + } catch (error) { + console.warn('Model Resolver: failed to load active download recovery', error); + this._activeDownloadRecovery = {}; + } + return this._activeDownloadRecovery; + }, + + saveActiveDownloadRecovery() { + const recovery = this.loadActiveDownloadRecovery(); + try { + localStorage.setItem( + this.getActiveDownloadRecoveryStorageKey(), + JSON.stringify(recovery) + ); + } catch (error) { + console.warn('Model Resolver: failed to save active download recovery', error); + } + }, + + persistActiveDownloadRecovery(downloadId, info = {}) { + if (!downloadId || !info) return; + + const recovery = this.loadActiveDownloadRecovery(); + const cloneForStorage = (value) => { + try { + return JSON.parse(JSON.stringify(value, (key, item) => ( + key === 'matches' || key === 'local_matches' ? undefined : item + ))); + } catch (_error) { + return null; + } + }; + const entry = { + missing: cloneForStorage(info.missing), + category: info.category || '', + subfolder: info.subfolder || '', + filename: info.filename || '', + downloadPath: info.downloadPath || '', + downloadDirectory: info.downloadDirectory || '', + baseDirectory: info.baseDirectory || '', + sourceUrl: info.sourceUrl || '', + workflowKey: info.workflowKey || '', + workflowRouteKey: info.workflowRouteKey || '', + workflowLabel: info.workflowLabel || '', + workflowSignature: info.workflowSignature || '', + workflowId: info.workflowId || '', + workflowTabId: info.workflowTabId || '', + workflowTabName: info.workflowTabName || '', + workflowTabAriaControls: info.workflowTabAriaControls || '', + workflowTabText: info.workflowTabText || '', + downloadBackend: info.downloadBackend || '' + }; + recovery[String(downloadId)] = entry; + this.saveActiveDownloadRecovery(); + }, + + removeActiveDownloadRecovery(downloadId) { + if (!downloadId) return; + + const recovery = this.loadActiveDownloadRecovery(); + const key = String(downloadId); + if (!(key in recovery)) return; + delete recovery[key]; + this.saveActiveDownloadRecovery(); + }, + persistQueueCollapsedState(collapsed, delayMs = 250) { if (this._queueCollapsedPersistTimer) { clearTimeout(this._queueCollapsedPersistTimer); diff --git a/web/resolver/actions/resolve_download_methods.js b/web/resolver/actions/resolve_download_methods.js index ecc1b9b..330d7eb 100644 --- a/web/resolver/actions/resolve_download_methods.js +++ b/web/resolver/actions/resolve_download_methods.js @@ -988,6 +988,7 @@ export const resolveDownloadMethods = { } delete this.activeDownloads?.[downloadId]; + this.removeActiveDownloadRecovery?.(downloadId); const elements = this.resolveDownloadUiElements?.(info) || {}; this.renderDownloadSnapshot?.(downloadId, cancelledSnapshot, elements); @@ -1758,6 +1759,181 @@ export const resolveDownloadMethods = { } }, + buildRecoveredDownloadMissing(downloadId, progress = {}, recovery = {}) { + const filename = progress.filename || recovery.filename || 'Download'; + const url = progress.url || recovery.sourceUrl || ''; + const category = progress.category || recovery.category || ''; + return { + missing_key: `download:${downloadId}`, + original_path: filename, + filename, + category, + node_type: 'Backend download', + download_source: { + source: 'backend', + url, + filename, + directory: category, + category + } + }; + }, + + rebindActiveDownloadMissingModels() { + const candidates = []; + const seen = new Set(); + const addCandidates = (models) => { + for (const model of Array.isArray(models) ? models : []) { + if (!model || seen.has(model)) continue; + seen.add(model); + candidates.push(model); + } + }; + addCandidates(this.missingModels); + addCandidates(this.cachedAnalysisData?.missing_models); + for (const cached of this.workflowAnalysisCaches?.values?.() || []) { + addCandidates(cached?.data?.missing_models); + } + if (!candidates.length) return; + + const normalizeUrl = (value = '') => String(value || '') + .trim() + .replace(/\/blob\//g, '/resolve/') + .replace(/[?#].*$/, '') + .toLowerCase(); + const getFilename = (value = '') => String(value || '') + .split(/[\\/]+/) + .pop() + .trim() + .toLowerCase(); + + for (const info of Object.values(this.activeDownloads || {})) { + if (!info?.missing) continue; + + const missingKey = this.getMissingModelKey?.(info.missing) || ''; + let replacement = candidates.find(model => ( + missingKey && this.getMissingModelKey?.(model) === missingKey + )); + + if (!replacement) { + const progress = info.lastProgress || {}; + const downloadUrl = normalizeUrl(progress.url || info.sourceUrl); + const filename = getFilename(progress.filename || info.filename); + const urlMatches = downloadUrl + ? candidates.filter(model => normalizeUrl( + model.download_source?.url + || model.download_source?.download_url + || model.url + ) === downloadUrl) + : []; + replacement = urlMatches.length === 1 ? urlMatches[0] : null; + if (!replacement && filename) { + const filenameMatches = candidates.filter(model => ( + getFilename(model.download_source?.filename) + || getFilename(model.original_path) + ) === filename); + replacement = filenameMatches.length === 1 ? filenameMatches[0] : null; + } + } + + if (replacement) info.missing = replacement; + } + }, + + async restoreActiveDownloadsFromBackend() { + if (this._activeDownloadsRestorePromise) { + return this._activeDownloadsRestorePromise; + } + + this._activeDownloadsRestorePromise = (async () => { + let allProgress; + try { + allProgress = await this.fetchJson( + '/model_resolver/progress', + { silent: true }, + 'Get active downloads' + ); + } catch (error) { + console.warn('Model Resolver: failed to restore active downloads:', error); + return 0; + } + + const terminalStatuses = new Set([ + 'completed', + 'completed_checking', + 'cancelled', + 'error', + 'refresh_error' + ]); + const recovery = this.loadActiveDownloadRecovery?.() || {}; + const activeIds = new Set(); + let restored = 0; + + for (const [downloadId, rawProgress] of Object.entries(allProgress || {})) { + const progress = rawProgress && typeof rawProgress === 'object' + ? rawProgress + : {}; + const status = String(progress.status || '').toLowerCase(); + if (terminalStatuses.has(status) || !this.isDownloadProgressStatus(status, true)) { + continue; + } + + activeIds.add(String(downloadId)); + const stored = recovery[downloadId] && typeof recovery[downloadId] === 'object' + ? recovery[downloadId] + : {}; + const existing = this.activeDownloads?.[downloadId]; + if (existing) { + existing.lastProgress = progress; + existing.lastStatus = status || existing.lastStatus || 'starting'; + this.rememberDownloadUiState?.(downloadId, existing, progress, { isActive: true }); + continue; + } + + const info = { + ...stored, + missing: stored.missing && typeof stored.missing === 'object' + ? stored.missing + : this.buildRecoveredDownloadMissing(downloadId, progress, stored), + progressDiv: null, + downloadBtn: null, + category: progress.category || stored.category || '', + subfolder: stored.subfolder || '', + filename: progress.filename || stored.filename || 'Download', + downloadPath: progress.path || stored.downloadPath || '', + downloadDirectory: progress.directory || stored.downloadDirectory || '', + baseDirectory: stored.baseDirectory || '', + sourceUrl: progress.url || stored.sourceUrl || '', + downloadBackend: progress.download_backend || stored.downloadBackend || '', + lastProgress: progress, + lastStatus: status || 'starting' + }; + this.activeDownloads[downloadId] = info; + this.rememberDownloadUiState?.(downloadId, info, progress, { isActive: true }); + this.persistActiveDownloadRecovery?.(downloadId, info); + this.pollDownloadProgress(downloadId); + restored += 1; + } + + for (const downloadId of Object.keys(recovery)) { + if (!activeIds.has(String(downloadId))) { + this.removeActiveDownloadRecovery?.(downloadId); + } + } + + this.rebindActiveDownloadMissingModels(); + this.updateDownloadAllButtonState?.(); + this.updateQueuePanel?.({ force: true }); + return restored; + })(); + + try { + return await this._activeDownloadsRestorePromise; + } finally { + this._activeDownloadsRestorePromise = null; + } + }, + async executeDownloadRequest(missing, { url, filename, @@ -1866,6 +2042,7 @@ export const resolveDownloadMethods = { workflowTabText, downloadBackend: data.download_backend || '' }; + this.persistActiveDownloadRecovery?.(downloadId, this.activeDownloads[downloadId]); const snapshot = this.rememberDownloadUiState( downloadId, @@ -2070,6 +2247,7 @@ export const resolveDownloadMethods = { this.rememberDownloadSubfolderAfterCompletion(info, progress); this.rememberCompletedDownloadHistory?.(downloadId, info, progress); delete this.activeDownloads[downloadId]; + this.removeActiveDownloadRecovery?.(downloadId); this.updateDownloadAllButtonState(); this.updateQueuePanel?.(); this.refreshLocalMatchesUiForMissing?.(missing); @@ -2102,6 +2280,7 @@ export const resolveDownloadMethods = { }); this.renderDownloadSnapshot(downloadId, errorSnapshot, { progressDiv, downloadBtn }); delete this.activeDownloads[downloadId]; + this.removeActiveDownloadRecovery?.(downloadId); this.updateDownloadAllButtonState(); this.updateQueuePanel?.(); this.refreshLocalMatchesUiForMissing?.(missing); @@ -2134,6 +2313,7 @@ export const resolveDownloadMethods = { this.renderDownloadSnapshot(downloadId, snapshot); } delete this.activeDownloads[downloadId]; + this.removeActiveDownloadRecovery?.(downloadId); this.updateDownloadAllButtonState(); this.updateQueuePanel?.(); } diff --git a/web/resolver/resolver_dialog.js b/web/resolver/resolver_dialog.js index 54a3205..53d9b50 100644 --- a/web/resolver/resolver_dialog.js +++ b/web/resolver/resolver_dialog.js @@ -48,6 +48,7 @@ export class ResolverManagerDialog extends ComfyDialog { this.workflowDownloadTargetSelectionCaches = new Map(); // workflow key -> suggested/user download folders this.activeDownloads = {}; // Track active downloads this.downloadHistoryStorageKey = 'model_resolver_download_history'; + this.activeDownloadsStorageKey = 'model_resolver_active_downloads'; this.downloadHistoryLimit = 200; this.downloadHistory = []; this.downloadProgressByMissingKey = new Map(); // missing model key -> last known download UI state diff --git a/web/resolver/shell/lifecycle_graph_methods.js b/web/resolver/shell/lifecycle_graph_methods.js index 5d04baa..7a5a58f 100644 --- a/web/resolver/shell/lifecycle_graph_methods.js +++ b/web/resolver/shell/lifecycle_graph_methods.js @@ -45,6 +45,7 @@ export const lifecycleGraphMethods = { this.updateTabButtonStates(); this.updateQueueVisibility(); this.syncWorkflowScopedQueue(workflow || this.getCurrentWorkflow()); + await this.restoreActiveDownloadsFromBackend?.(); if (this.activeTab === 'missing') { await this.loadWorkflowData(workflow); @@ -795,6 +796,7 @@ export const lifecycleGraphMethods = { reconnectActiveDownloads() { if (!this.contentElement) return; + this.rebindActiveDownloadMissingModels?.(); for (const [downloadId, info] of Object.entries(this.activeDownloads)) { if (!info?.missing) continue; const snapshot = this.rememberDownloadUiState?.( @@ -805,6 +807,8 @@ export const lifecycleGraphMethods = { ); this.renderDownloadSnapshot?.(downloadId, snapshot); } + this.updateDownloadAllButtonState?.(); + this.updateQueuePanel?.({ force: true }); }, };