From 610c00ae81ff2b50ed4c59ce5d0c5038c57d192e Mon Sep 17 00:00:00 2001 From: Benjamin Lu Date: Fri, 17 Apr 2026 16:20:58 -0700 Subject: [PATCH 1/4] Align desktop download state contract --- src/models/DownloadManager.ts | 227 ++++++++++++++-------- src/preload.ts | 11 +- tests/unit/models/DownloadManager.test.ts | 157 +++++++++++++++ 3 files changed, 301 insertions(+), 94 deletions(-) diff --git a/src/models/DownloadManager.ts b/src/models/DownloadManager.ts index c2a666d41..def415329 100644 --- a/src/models/DownloadManager.ts +++ b/src/models/DownloadManager.ts @@ -15,26 +15,30 @@ export interface Download { directoryPath: string; savePath: string; item: DownloadItem | null; + progress: number; + status: DownloadStatus; + message?: string; + receivedBytes: number; + totalBytes: number; } export interface DownloadState { url: string; filename: string; + savePath: string; + progress: number; + status: DownloadStatus; + message?: string; + /** @deprecated Use `status` instead. */ state: DownloadStatus; + /** @deprecated Use `progress` instead. */ receivedBytes: number; + /** @deprecated Use `progress` instead. */ totalBytes: number; + /** @deprecated Use `status === DownloadStatus.PAUSED` instead. */ isPaused: boolean; } -interface DownloadReport { - url: string; - progress: number; - status: DownloadStatus; - filename: string; - savePath: string; - message?: string; -} - /** * Singleton class that manages downloading model checkpoints for ComfyUI. */ @@ -54,44 +58,43 @@ export class DownloadManager { const download = this.downloads.get(url); if (!download) return; - this.reportProgress({ - url, - filename: download.filename, - savePath: download.savePath, - progress: 0, - status: DownloadStatus.PENDING, - }); item.setSavePath(download.tempPath); download.item = item; + download.totalBytes = item.getTotalBytes(); log.info(`Setting save path to ${item.getSavePath()}`); item.on('updated', (event, state) => { if (state === 'interrupted') { log.info('Download is interrupted but can be resumed'); } else if (state === 'progressing') { - const progress = item.getReceivedBytes() / item.getTotalBytes(); + const receivedBytes = item.getReceivedBytes(); + const totalBytes = item.getTotalBytes(); + const progress = this.calculateProgress(receivedBytes, totalBytes); if (item.isPaused()) { log.info('Download is paused'); - this.reportProgress({ - url, - progress, - filename: download.filename, - savePath: download.savePath, - status: DownloadStatus.PAUSED, - }); + download.progress = progress; + download.status = DownloadStatus.PAUSED; + download.message = undefined; + download.receivedBytes = receivedBytes; + download.totalBytes = totalBytes; + this.reportProgress(this.toDownloadState(download)); } else { - this.reportProgress({ - url, - progress, - filename: download.filename, - savePath: download.savePath, - status: DownloadStatus.IN_PROGRESS, - }); + download.progress = progress; + download.status = DownloadStatus.IN_PROGRESS; + download.message = undefined; + download.receivedBytes = receivedBytes; + download.totalBytes = totalBytes; + this.reportProgress(this.toDownloadState(download)); } } }); item.once('done', (event, state) => { + const receivedBytes = item.getReceivedBytes(); + const totalBytes = item.getTotalBytes(); + download.receivedBytes = receivedBytes; + download.totalBytes = totalBytes; + if (state === 'completed') { try { fs.renameSync(download.tempPath, download.savePath); @@ -100,24 +103,25 @@ export class DownloadManager { log.error('Failed to rename downloaded file. Deleting temp file.', error); fs.unlinkSync(download.tempPath); } - this.reportProgress({ - url, - filename: download.filename, - savePath: download.savePath, - progress: 1, - status: DownloadStatus.COMPLETED, - }); - this.downloads.delete(url); + download.item = null; + download.progress = 1; + download.status = DownloadStatus.COMPLETED; + download.message = undefined; + this.reportProgress(this.toDownloadState(download)); + } else if (state === 'cancelled') { + log.info('Download cancelled'); + download.item = null; + download.progress = this.calculateProgress(receivedBytes, totalBytes); + download.status = DownloadStatus.CANCELLED; + download.message = undefined; + this.reportProgress(this.toDownloadState(download)); } else { log.info(`Download failed: ${state}`); - const progress = item.getReceivedBytes() / item.getTotalBytes(); - this.reportProgress({ - url, - filename: download.filename, - progress, - status: DownloadStatus.ERROR, - savePath: download.savePath, - }); + download.item = null; + download.progress = this.calculateProgress(receivedBytes, totalBytes); + download.status = DownloadStatus.ERROR; + download.message = 'Download interrupted'; + this.reportProgress(this.toDownloadState(download)); } }); }); @@ -130,11 +134,15 @@ export class DownloadManager { log.error(`Save path ${localSavePath} is not in models directory ${this.modelsDirectory}`); this.reportProgress({ url, - savePath: normalizedDirectoryPath, + savePath: localSavePath, filename, progress: 0, status: DownloadStatus.ERROR, message: 'Save path is not in models directory', + state: DownloadStatus.ERROR, + receivedBytes: 0, + totalBytes: 0, + isPaused: false, }); return false; } @@ -144,39 +152,76 @@ export class DownloadManager { log.error(validationResult.error); this.reportProgress({ url, - savePath: normalizedDirectoryPath, + savePath: localSavePath, filename, progress: 0, status: DownloadStatus.ERROR, message: validationResult.error, + state: DownloadStatus.ERROR, + receivedBytes: 0, + totalBytes: 0, + isPaused: false, }); return false; } if (fs.existsSync(localSavePath)) { log.info(`File ${filename} already exists, skipping download`); + const existingCompletedDownload = this.downloads.get(url) ?? { + url, + directoryPath: normalizedDirectoryPath, + savePath: localSavePath, + tempPath: this.getTempPath(filename, normalizedDirectoryPath), + filename, + item: null, + progress: 1, + status: DownloadStatus.COMPLETED, + message: undefined, + receivedBytes: 0, + totalBytes: 0, + }; + existingCompletedDownload.progress = 1; + existingCompletedDownload.status = DownloadStatus.COMPLETED; + existingCompletedDownload.message = undefined; + existingCompletedDownload.item = null; + this.downloads.set(url, existingCompletedDownload); + this.reportProgress(this.toDownloadState(existingCompletedDownload)); return true; } const existingDownload = this.downloads.get(url); if (existingDownload) { log.info('Download already exists'); - if (existingDownload.item?.isPaused()) { + if (existingDownload.status === DownloadStatus.PAUSED) { this.resumeDownload(url); + return true; + } else if ( + existingDownload.status === DownloadStatus.CANCELLED || + existingDownload.status === DownloadStatus.ERROR + ) { + this.deleteTempFile(existingDownload.tempPath); + this.downloads.delete(url); + } else { + return true; } - return true; } log.info(`Starting download ${url} to ${localSavePath}`); - const tempPath = this.getTempPath(filename, normalizedDirectoryPath); - this.downloads.set(url, { + const download: Download = { url, directoryPath: normalizedDirectoryPath, savePath: localSavePath, - tempPath, + tempPath: this.getTempPath(filename, normalizedDirectoryPath), filename, item: null, - }); + progress: 0, + status: DownloadStatus.PENDING, + message: undefined, + receivedBytes: 0, + totalBytes: 0, + }; + this.downloads.set(url, download); + this.reportProgress(this.toDownloadState(download)); // TODO(robinhuang): Add offset support for resuming downloads. // Can use https://www.electronjs.org/docs/latest/api/session#sescreateinterrupteddownloadoptions @@ -186,12 +231,17 @@ export class DownloadManager { cancelDownload(url: string): void { const download = this.downloads.get(url); - if (!download?.item) return; + if (!download) return; log.info('Cancelling download'); - download.item.cancel(); + if (download.item) { + download.item.cancel(); + return; + } - this.downloads.delete(url); + download.status = DownloadStatus.CANCELLED; + download.message = undefined; + this.reportProgress(this.toDownloadState(download)); } pauseDownload(url: string): void { @@ -243,38 +293,45 @@ export class DownloadManager { } getAllDownloads(): DownloadState[] { - return [...this.downloads.values()] - .filter((download) => download.item !== null) - .map((download) => ({ - url: download.url, - filename: download.filename, - tempPath: download.tempPath, - state: this.convertDownloadState(download.item?.getState()), - receivedBytes: download.item?.getReceivedBytes() || 0, - totalBytes: download.item?.getTotalBytes() || 0, - isPaused: download.item?.isPaused() || false, - })); - } - - private convertDownloadState(state?: 'progressing' | 'completed' | 'cancelled' | 'interrupted'): DownloadStatus { - switch (state) { - case 'progressing': - return DownloadStatus.IN_PROGRESS; - case 'completed': - return DownloadStatus.COMPLETED; - case 'cancelled': - return DownloadStatus.CANCELLED; - case 'interrupted': - return DownloadStatus.ERROR; - default: - return DownloadStatus.ERROR; - } + return [...this.downloads.values()].map((download) => this.toDownloadState(download)); } private getTempPath(filename: string, directoryPath: string): string { return path.join(directoryPath, `Unconfirmed ${filename}.tmp`); } + private calculateProgress(receivedBytes: number, totalBytes: number): number { + if (totalBytes <= 0) return 0; + return receivedBytes / totalBytes; + } + + private toDownloadState(download: Download): DownloadState { + const isPaused = download.status === DownloadStatus.PAUSED || download.item?.isPaused() || false; + + return { + url: download.url, + filename: download.filename, + savePath: download.savePath, + progress: download.progress, + status: download.status, + message: download.message, + state: download.status, + receivedBytes: download.receivedBytes, + totalBytes: download.totalBytes, + isPaused, + }; + } + + private deleteTempFile(tempPath: string): void { + try { + if (fs.existsSync(tempPath)) { + fs.unlinkSync(tempPath); + } + } catch (error) { + log.error(`Failed to delete temp file ${tempPath}:`, error); + } + } + // Only allow .safetensors files to be downloaded. private validateSafetensorsFile(url: string, filename: string): { isValid: boolean; error?: string } { try { @@ -381,7 +438,7 @@ export class DownloadManager { return process.platform === 'win32' ? targetPath.toLowerCase() : targetPath; } - private reportProgress(report: DownloadReport): void { + private reportProgress(report: DownloadState): void { log.info( `Download progress [${report.filename}]: ${report.progress}, status: ${report.status}, message: ${report.message}` ); diff --git a/src/preload.ts b/src/preload.ts index f795f38e8..38e473882 100644 --- a/src/preload.ts +++ b/src/preload.ts @@ -3,7 +3,7 @@ import path from 'node:path'; import { strictIpcRenderer as ipcRenderer } from '@/infrastructure/ipcChannels'; -import { DownloadStatus, ELECTRON_BRIDGE_API, IPC_CHANNELS, ProgressStatus } from './constants'; +import { ELECTRON_BRIDGE_API, IPC_CHANNELS, ProgressStatus } from './constants'; import type { RestrictedPathType } from './handlers/pathHandlers'; import type { InstallStageInfo } from './main-process/installStages'; import type { DownloadState } from './main_types'; @@ -45,14 +45,7 @@ export interface SystemPaths { defaultInstallPath: string; } -export interface DownloadProgressUpdate { - url: string; - filename: string; - savePath: string; - progress: number; - status: DownloadStatus; - message?: string; -} +export type DownloadProgressUpdate = DownloadState; /** @todo Type inference chain broken by comfyui-electron-types. This is duplication. */ export interface ElectronOverlayOptions { diff --git a/tests/unit/models/DownloadManager.test.ts b/tests/unit/models/DownloadManager.test.ts index 884a34b75..a47afac4d 100644 --- a/tests/unit/models/DownloadManager.test.ts +++ b/tests/unit/models/DownloadManager.test.ts @@ -2,6 +2,8 @@ import fs from 'node:fs'; import path from 'node:path'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { DownloadStatus, IPC_CHANNELS } from '@/constants'; + import { electronMock } from '../setup'; vi.mock('node:fs'); @@ -143,6 +145,85 @@ describe('DownloadManager', () => { expect(downloadURL).toHaveBeenCalledWith(url); }); + it('emits a canonical pending snapshot as soon as a download starts', () => { + const modelsDirectory = path.resolve('/mock/models'); + const checkpointsDirectory = path.join(modelsDirectory, 'checkpoints'); + const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); + const url = 'https://example.com/model.safetensors'; + mockExistingPaths(modelsDirectory, checkpointsDirectory); + + expect(manager.startDownload(url, checkpointsDirectory, 'model.safetensors')).toBe(true); + + expect(mainWindow.send).toHaveBeenCalledWith( + IPC_CHANNELS.DOWNLOAD_PROGRESS, + expect.objectContaining({ + url, + filename: 'model.safetensors', + savePath: path.join(checkpointsDirectory, 'model.safetensors'), + progress: 0, + status: DownloadStatus.PENDING, + state: DownloadStatus.PENDING, + receivedBytes: 0, + totalBytes: 0, + isPaused: false, + }) + ); + }); + + it('resumes paused downloads without starting a second download', () => { + const modelsDirectory = path.resolve('/mock/models'); + const checkpointsDirectory = path.join(modelsDirectory, 'checkpoints'); + const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); + mockExistingPaths(modelsDirectory, checkpointsDirectory); + const downloads = ( + manager as unknown as { + downloads: Map< + string, + { + url: string; + filename: string; + directoryPath: string; + savePath: string; + tempPath: string; + progress: number; + status: DownloadStatus; + message?: string; + receivedBytes: number; + totalBytes: number; + item: { + canResume: () => boolean; + resume: () => void; + } | null; + } + >; + } + ).downloads; + const resume = vi.fn(); + const url = 'https://example.com/model.safetensors'; + + downloads.set(url, { + url, + filename: 'model.safetensors', + directoryPath: checkpointsDirectory, + savePath: path.join(checkpointsDirectory, 'model.safetensors'), + tempPath: path.join(checkpointsDirectory, 'Unconfirmed model.safetensors.tmp'), + progress: 0.5, + status: DownloadStatus.PAUSED, + message: undefined, + receivedBytes: 5, + totalBytes: 10, + item: { + canResume: () => true, + resume, + }, + }); + + expect(manager.startDownload(url, checkpointsDirectory, 'model.safetensors')).toBe(true); + + expect(resume).toHaveBeenCalledOnce(); + expect(downloadURL).not.toHaveBeenCalled(); + }); + it('does not create directories outside models directory', () => { const modelsDirectory = path.resolve('/mock/models'); const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); @@ -212,6 +293,11 @@ describe('DownloadManager', () => { directoryPath: string; savePath: string; tempPath: string; + progress: number; + status: DownloadStatus; + message?: string; + receivedBytes: number; + totalBytes: number; item: { canResume: () => boolean; resume: () => void }; } >; @@ -226,6 +312,11 @@ describe('DownloadManager', () => { directoryPath: checkpointsDirectory, savePath: path.join(checkpointsDirectory, 'model.safetensors'), tempPath: path.join(checkpointsDirectory, 'Unconfirmed model.safetensors.tmp'), + progress: 0.5, + status: DownloadStatus.PAUSED, + message: undefined, + receivedBytes: 5, + totalBytes: 10, item: { canResume: () => false, resume, @@ -237,4 +328,70 @@ describe('DownloadManager', () => { expect(resume).not.toHaveBeenCalled(); expect(downloadURL).toHaveBeenCalledWith(url); }); + + it('returns canonical snapshots from getAllDownloads while preserving legacy fields', () => { + const modelsDirectory = path.resolve('/mock/models'); + const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); + const url = 'https://example.com/model.safetensors'; + const savePath = path.join(modelsDirectory, 'checkpoints', 'model.safetensors'); + const downloads = ( + manager as unknown as { + downloads: Map< + string, + { + url: string; + filename: string; + directoryPath: string; + savePath: string; + tempPath: string; + progress: number; + status: DownloadStatus; + message?: string; + receivedBytes: number; + totalBytes: number; + item: { + getState: () => 'progressing'; + getReceivedBytes: () => number; + getTotalBytes: () => number; + isPaused: () => boolean; + } | null; + } + >; + } + ).downloads; + + downloads.set(url, { + url, + filename: 'model.safetensors', + directoryPath: path.dirname(savePath), + savePath, + tempPath: path.join(path.dirname(savePath), 'Unconfirmed model.safetensors.tmp'), + progress: 0.5, + status: DownloadStatus.IN_PROGRESS, + message: undefined, + receivedBytes: 5, + totalBytes: 10, + item: { + getState: () => 'progressing', + getReceivedBytes: () => 5, + getTotalBytes: () => 10, + isPaused: () => false, + }, + }); + + expect(manager.getAllDownloads()).toEqual([ + { + url, + filename: 'model.safetensors', + savePath, + progress: 0.5, + status: DownloadStatus.IN_PROGRESS, + message: undefined, + state: DownloadStatus.IN_PROGRESS, + receivedBytes: 5, + totalBytes: 10, + isPaused: false, + }, + ]); + }); }); From dfae1a51b11fc723c2587397cf2df74033a47deb Mon Sep 17 00:00:00 2001 From: Benjamin Lu Date: Mon, 4 May 2026 08:19:42 -0700 Subject: [PATCH 2/4] fix: add stable download identity contract --- src/infrastructure/ipcChannels.ts | 10 +- src/main_types.ts | 2 +- src/models/DownloadManager.ts | 119 +++++++--- src/preload.ts | 16 +- .../post-install/downloadManager.spec.ts | 10 +- tests/unit/models/DownloadManager.test.ts | 221 +++++++++--------- 6 files changed, 225 insertions(+), 153 deletions(-) diff --git a/src/infrastructure/ipcChannels.ts b/src/infrastructure/ipcChannels.ts index 0735864dd..b3ce2a134 100644 --- a/src/infrastructure/ipcChannels.ts +++ b/src/infrastructure/ipcChannels.ts @@ -2,7 +2,7 @@ import { ipcMain, ipcRenderer } from 'electron'; import type { IPC_CHANNELS, ProgressStatus } from '@/constants'; import type { InstallStageInfo } from '@/main-process/installStages'; -import type { DownloadState } from '@/models/DownloadManager'; +import type { DownloadState, StartDownloadResult } from '@/models/DownloadManager'; import type { DownloadProgressUpdate, ElectronContextMenuOptions, @@ -264,21 +264,21 @@ export interface IpcChannels { [IPC_CHANNELS.START_DOWNLOAD]: { params: [details: { url: string; path: string; filename: string }]; - return: boolean; + return: StartDownloadResult; }; [IPC_CHANNELS.PAUSE_DOWNLOAD]: { - params: [url: string]; + params: [downloadId: string]; return: void; }; [IPC_CHANNELS.RESUME_DOWNLOAD]: { - params: [url: string]; + params: [downloadId: string]; return: void; }; [IPC_CHANNELS.CANCEL_DOWNLOAD]: { - params: [url: string]; + params: [downloadId: string]; return: void; }; diff --git a/src/main_types.ts b/src/main_types.ts index a4026002b..81868ab20 100644 --- a/src/main_types.ts +++ b/src/main_types.ts @@ -1,5 +1,5 @@ export * from './constants'; -export type { DownloadState } from './models/DownloadManager'; +export type { DownloadState, StartDownloadResult } from './models/DownloadManager'; export type { InstallStageInfo, InstallStageName } from './main-process/installStages'; export type { ElectronAPI, diff --git a/src/models/DownloadManager.ts b/src/models/DownloadManager.ts index def415329..2e995fe0c 100644 --- a/src/models/DownloadManager.ts +++ b/src/models/DownloadManager.ts @@ -9,6 +9,7 @@ import { DownloadStatus, IPC_CHANNELS } from '../constants'; import type { AppWindow } from '../main-process/appWindow'; export interface Download { + downloadId: string; url: string; filename: string; tempPath: string; // Temporary filename until the download is complete. @@ -23,6 +24,7 @@ export interface Download { } export interface DownloadState { + downloadId: string; url: string; filename: string; savePath: string; @@ -39,23 +41,29 @@ export interface DownloadState { isPaused: boolean; } +export type StartDownloadResult = + | { ok: true; download: DownloadState } + | { ok: false; error: string; download: DownloadState }; + /** * Singleton class that manages downloading model checkpoints for ComfyUI. */ export class DownloadManager { private static instance: DownloadManager; private readonly downloads: Map; + private readonly pendingDownloadIdsByUrl: Map; private constructor( private readonly mainWindow: AppWindow, private readonly modelsDirectory: string ) { this.downloads = new Map(); + this.pendingDownloadIdsByUrl = new Map(); session.defaultSession.on('will-download', (event, item) => { const url = item.getURLChain()[0]; // Get the original URL in case of redirects. log.info('Will-download event', url); - const download = this.downloads.get(url); + const download = this.takePendingDownload(url); if (!download) return; item.setSavePath(download.tempPath); @@ -127,12 +135,14 @@ export class DownloadManager { }); } - startDownload(url: string, directoryPath: string, filename: string): boolean { + startDownload(url: string, directoryPath: string, filename: string): StartDownloadResult { const normalizedDirectoryPath = this.normalizeDirectoryPath(directoryPath); const localSavePath = this.getLocalSavePath(filename, normalizedDirectoryPath); + const downloadId = this.createDownloadId(localSavePath); if (!this.ensureDownloadTargetDirectory(localSavePath)) { log.error(`Save path ${localSavePath} is not in models directory ${this.modelsDirectory}`); - this.reportProgress({ + const downloadState: DownloadState = { + downloadId, url, savePath: localSavePath, filename, @@ -143,31 +153,44 @@ export class DownloadManager { receivedBytes: 0, totalBytes: 0, isPaused: false, - }); - return false; + }; + this.reportProgress(downloadState); + return { + ok: false, + error: 'Save path is not in models directory', + download: downloadState, + }; } const validationResult = this.validateSafetensorsFile(url, filename); if (!validationResult.isValid) { log.error(validationResult.error); - this.reportProgress({ + const errorMessage = validationResult.error ?? 'Invalid download'; + const downloadState: DownloadState = { + downloadId, url, savePath: localSavePath, filename, progress: 0, status: DownloadStatus.ERROR, - message: validationResult.error, + message: errorMessage, state: DownloadStatus.ERROR, receivedBytes: 0, totalBytes: 0, isPaused: false, - }); - return false; + }; + this.reportProgress(downloadState); + return { + ok: false, + error: errorMessage, + download: downloadState, + }; } if (fs.existsSync(localSavePath)) { log.info(`File ${filename} already exists, skipping download`); - const existingCompletedDownload = this.downloads.get(url) ?? { + const existingCompletedDownload = this.downloads.get(downloadId) ?? { + downloadId, url, directoryPath: normalizedDirectoryPath, savePath: localSavePath, @@ -184,30 +207,32 @@ export class DownloadManager { existingCompletedDownload.status = DownloadStatus.COMPLETED; existingCompletedDownload.message = undefined; existingCompletedDownload.item = null; - this.downloads.set(url, existingCompletedDownload); - this.reportProgress(this.toDownloadState(existingCompletedDownload)); - return true; + this.downloads.set(downloadId, existingCompletedDownload); + const downloadState = this.toDownloadState(existingCompletedDownload); + this.reportProgress(downloadState); + return { ok: true, download: downloadState }; } - const existingDownload = this.downloads.get(url); + const existingDownload = this.downloads.get(downloadId); if (existingDownload) { log.info('Download already exists'); if (existingDownload.status === DownloadStatus.PAUSED) { - this.resumeDownload(url); - return true; + this.resumeDownload(downloadId); + return { ok: true, download: this.toDownloadState(existingDownload) }; } else if ( existingDownload.status === DownloadStatus.CANCELLED || existingDownload.status === DownloadStatus.ERROR ) { this.deleteTempFile(existingDownload.tempPath); - this.downloads.delete(url); + this.downloads.delete(downloadId); } else { - return true; + return { ok: true, download: this.toDownloadState(existingDownload) }; } } log.info(`Starting download ${url} to ${localSavePath}`); const download: Download = { + downloadId, url, directoryPath: normalizedDirectoryPath, savePath: localSavePath, @@ -220,17 +245,19 @@ export class DownloadManager { receivedBytes: 0, totalBytes: 0, }; - this.downloads.set(url, download); - this.reportProgress(this.toDownloadState(download)); + this.downloads.set(downloadId, download); + const downloadState = this.toDownloadState(download); + this.reportProgress(downloadState); // TODO(robinhuang): Add offset support for resuming downloads. // Can use https://www.electronjs.org/docs/latest/api/session#sescreateinterrupteddownloadoptions + this.enqueuePendingDownload(url, downloadId); session.defaultSession.downloadURL(url); - return true; + return { ok: true, download: downloadState }; } - cancelDownload(url: string): void { - const download = this.downloads.get(url); + cancelDownload(downloadIdOrUrl: string): void { + const download = this.findDownload(downloadIdOrUrl); if (!download) return; log.info('Cancelling download'); @@ -244,23 +271,23 @@ export class DownloadManager { this.reportProgress(this.toDownloadState(download)); } - pauseDownload(url: string): void { - const download = this.downloads.get(url); + pauseDownload(downloadIdOrUrl: string): void { + const download = this.findDownload(downloadIdOrUrl); if (!download?.item) return; log.info('Pausing download'); download.item.pause(); } - resumeDownload(url: string): void { - const download = this.downloads.get(url); + resumeDownload(downloadIdOrUrl: string): void { + const download = this.findDownload(downloadIdOrUrl); if (!download?.item) return; if (download.item.canResume()) { log.info('Resuming download'); download.item.resume(); } else { - this.downloads.delete(url); + this.downloads.delete(download.downloadId); this.startDownload(download.url, download.directoryPath, download.filename); } } @@ -309,6 +336,7 @@ export class DownloadManager { const isPaused = download.status === DownloadStatus.PAUSED || download.item?.isPaused() || false; return { + downloadId: download.downloadId, url: download.url, filename: download.filename, savePath: download.savePath, @@ -322,6 +350,35 @@ export class DownloadManager { }; } + private createDownloadId(savePath: string): string { + return this.getPathForComparison(path.resolve(savePath)); + } + + private enqueuePendingDownload(url: string, downloadId: string): void { + const pendingDownloadIds = this.pendingDownloadIdsByUrl.get(url) ?? []; + pendingDownloadIds.push(downloadId); + this.pendingDownloadIdsByUrl.set(url, pendingDownloadIds); + } + + private takePendingDownload(url: string): Download | undefined { + const pendingDownloadIds = this.pendingDownloadIdsByUrl.get(url); + const downloadId = pendingDownloadIds?.shift(); + if (pendingDownloadIds?.length === 0) { + this.pendingDownloadIdsByUrl.delete(url); + } + if (downloadId) { + return this.downloads.get(downloadId); + } + return [...this.downloads.values()].find((download) => download.url === url && download.item === null); + } + + private findDownload(downloadIdOrUrl: string): Download | undefined { + return ( + this.downloads.get(downloadIdOrUrl) ?? + [...this.downloads.values()].find((download) => download.url === downloadIdOrUrl) + ); + } + private deleteTempFile(tempPath: string): void { try { if (fs.existsSync(tempPath)) { @@ -465,9 +522,9 @@ export class DownloadManager { ipcMain.handle(IPC_CHANNELS.START_DOWNLOAD, (event, { url, path, filename }: DownloadDetails) => this.startDownload(url, path, filename) ); - ipcMain.handle(IPC_CHANNELS.PAUSE_DOWNLOAD, (event, url: string) => this.pauseDownload(url)); - ipcMain.handle(IPC_CHANNELS.RESUME_DOWNLOAD, (event, url: string) => this.resumeDownload(url)); - ipcMain.handle(IPC_CHANNELS.CANCEL_DOWNLOAD, (event, url: string) => this.cancelDownload(url)); + ipcMain.handle(IPC_CHANNELS.PAUSE_DOWNLOAD, (event, downloadId: string) => this.pauseDownload(downloadId)); + ipcMain.handle(IPC_CHANNELS.RESUME_DOWNLOAD, (event, downloadId: string) => this.resumeDownload(downloadId)); + ipcMain.handle(IPC_CHANNELS.CANCEL_DOWNLOAD, (event, downloadId: string) => this.cancelDownload(downloadId)); ipcMain.handle(IPC_CHANNELS.GET_ALL_DOWNLOADS, () => this.getAllDownloads()); ipcMain.handle(IPC_CHANNELS.DELETE_MODEL, (event, { filename, path }: FileAndPath) => diff --git a/src/preload.ts b/src/preload.ts index 38e473882..c3155f9f8 100644 --- a/src/preload.ts +++ b/src/preload.ts @@ -6,7 +6,7 @@ import { strictIpcRenderer as ipcRenderer } from '@/infrastructure/ipcChannels'; import { ELECTRON_BRIDGE_API, IPC_CHANNELS, ProgressStatus } from './constants'; import type { RestrictedPathType } from './handlers/pathHandlers'; import type { InstallStageInfo } from './main-process/installStages'; -import type { DownloadState } from './main_types'; +import type { DownloadState, StartDownloadResult } from './main_types'; import type { DesktopInstallState, DesktopWindowStyle } from './main_types'; /** @@ -191,18 +191,18 @@ const electronAPI = { onDownloadProgress: (callback: (progress: DownloadProgressUpdate) => void) => { ipcRenderer.on(IPC_CHANNELS.DOWNLOAD_PROGRESS, (_event, progress) => callback(progress)); }, - startDownload: (url: string, path: string, filename: string): Promise => { + startDownload: (url: string, path: string, filename: string): Promise => { console.log(`Sending start download message to main process`, { url, path, filename }); return ipcRenderer.invoke(IPC_CHANNELS.START_DOWNLOAD, { url, path, filename }); }, - cancelDownload: (url: string): Promise => { - return ipcRenderer.invoke(IPC_CHANNELS.CANCEL_DOWNLOAD, url); + cancelDownload: (downloadId: string): Promise => { + return ipcRenderer.invoke(IPC_CHANNELS.CANCEL_DOWNLOAD, downloadId); }, - pauseDownload: (url: string): Promise => { - return ipcRenderer.invoke(IPC_CHANNELS.PAUSE_DOWNLOAD, url); + pauseDownload: (downloadId: string): Promise => { + return ipcRenderer.invoke(IPC_CHANNELS.PAUSE_DOWNLOAD, downloadId); }, - resumeDownload: (url: string): Promise => { - return ipcRenderer.invoke(IPC_CHANNELS.RESUME_DOWNLOAD, url); + resumeDownload: (downloadId: string): Promise => { + return ipcRenderer.invoke(IPC_CHANNELS.RESUME_DOWNLOAD, downloadId); }, deleteModel: (filename: string, path: string): Promise => { return ipcRenderer.invoke(IPC_CHANNELS.DELETE_MODEL, { filename, path }); diff --git a/tests/integration/post-install/downloadManager.spec.ts b/tests/integration/post-install/downloadManager.spec.ts index 93a04b699..54bd97351 100644 --- a/tests/integration/post-install/downloadManager.spec.ts +++ b/tests/integration/post-install/downloadManager.spec.ts @@ -3,13 +3,15 @@ import { createServer } from 'node:http'; import path from 'node:path'; import { addRandomSuffix, pathExists } from 'tests/shared/utils'; +import type { StartDownloadResult } from '@/models/DownloadManager'; + import { expect, test } from '../testExtensions'; interface RendererElectronApi { electronAPI: { getBasePath: () => Promise; DownloadManager: { - startDownload: (url: string, directoryPath: string, filename: string) => Promise; + startDownload: (url: string, directoryPath: string, filename: string) => Promise; }; }; } @@ -62,7 +64,7 @@ test.describe('DownloadManager', () => { await rm(expectedFilePath, { force: true }); - const started = await window.evaluate( + const startResult = await window.evaluate( async ({ directoryPath, filename, url }) => { const api = (globalThis as typeof globalThis & RendererElectronApi).electronAPI; return await api.DownloadManager.startDownload(url, directoryPath, filename); @@ -74,7 +76,9 @@ test.describe('DownloadManager', () => { } ); - expect(started).toBe(true); + expect(startResult.ok).toBe(true); + expect(startResult.download.downloadId).toBe(expectedFilePath); + expect(startResult.download.savePath).toBe(expectedFilePath); await expect .poll( diff --git a/tests/unit/models/DownloadManager.test.ts b/tests/unit/models/DownloadManager.test.ts index a47afac4d..3eebdbbad 100644 --- a/tests/unit/models/DownloadManager.test.ts +++ b/tests/unit/models/DownloadManager.test.ts @@ -3,6 +3,7 @@ import path from 'node:path'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { DownloadStatus, IPC_CHANNELS } from '@/constants'; +import type { Download, StartDownloadResult } from '@/models/DownloadManager'; import { electronMock } from '../setup'; @@ -16,6 +17,26 @@ const mockExistingPaths = (...paths: string[]) => { vi.mocked(fs.existsSync).mockImplementation((targetPath) => existingPaths.has(path.resolve(String(targetPath)))); }; +function expectStartOk(result: StartDownloadResult) { + expect(result.ok).toBe(true); + if (!result.ok) { + throw new Error(result.error); + } + return result.download; +} + +function expectStartFailed(result: StartDownloadResult) { + expect(result.ok).toBe(false); + if (result.ok) { + throw new Error('Expected download start to fail'); + } + return result; +} + +function getDownloads(manager: unknown): Map { + return (manager as { downloads: Map }).downloads; +} + describe('DownloadManager', () => { let DownloadManager: typeof import('@/models/DownloadManager').DownloadManager; let defaultSessionOn: ReturnType; @@ -64,16 +85,13 @@ describe('DownloadManager', () => { const savePath = path.join(modelsDirectory, 'ipadapter'); mockExistingPaths(modelsDirectory, savePath); - expect(manager.startDownload(url, savePath, 'model.safetensors')).toBe(true); + const download = expectStartOk(manager.startDownload(url, savePath, 'model.safetensors')); expect(downloadURL).toHaveBeenCalledWith(url); - const downloads = ( - manager as unknown as { - downloads: Map; - } - ).downloads; - expect(downloads.get(url)?.savePath).toBe(path.join(savePath, 'model.safetensors')); - expect(downloads.get(url)?.tempPath).toBe(path.join(savePath, 'Unconfirmed model.safetensors.tmp')); + const downloads = getDownloads(manager); + expect(download.downloadId).toBe(path.join(savePath, 'model.safetensors')); + expect(downloads.get(download.downloadId)?.savePath).toBe(path.join(savePath, 'model.safetensors')); + expect(downloads.get(download.downloadId)?.tempPath).toBe(path.join(savePath, 'Unconfirmed model.safetensors.tmp')); }); it('normalizes relative save paths from legacy callers under the models directory', () => { @@ -82,26 +100,78 @@ describe('DownloadManager', () => { const url = 'https://example.com/model.safetensors'; mockExistingPaths(modelsDirectory, path.join(modelsDirectory, 'checkpoints')); - expect(manager.startDownload(url, 'checkpoints', 'model.safetensors')).toBe(true); + const download = expectStartOk(manager.startDownload(url, 'checkpoints', 'model.safetensors')); expect(downloadURL).toHaveBeenCalledWith(url); - const downloads = ( - manager as unknown as { - downloads: Map; - } - ).downloads; - expect(downloads.get(url)?.savePath).toBe(path.join(modelsDirectory, 'checkpoints', 'model.safetensors')); - expect(downloads.get(url)?.tempPath).toBe( + const downloads = getDownloads(manager); + expect(download.downloadId).toBe(path.join(modelsDirectory, 'checkpoints', 'model.safetensors')); + expect(downloads.get(download.downloadId)?.savePath).toBe( + path.join(modelsDirectory, 'checkpoints', 'model.safetensors') + ); + expect(downloads.get(download.downloadId)?.tempPath).toBe( path.join(modelsDirectory, 'checkpoints', 'Unconfirmed model.safetensors.tmp') ); }); + it('tracks same-source downloads separately by target save path', () => { + const modelsDirectory = path.resolve('/mock/models'); + const checkpointsDirectory = path.join(modelsDirectory, 'checkpoints'); + const lorasDirectory = path.join(modelsDirectory, 'loras'); + const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); + const url = 'https://example.com/model.safetensors'; + mockExistingPaths(modelsDirectory, checkpointsDirectory, lorasDirectory); + + const checkpointDownload = expectStartOk(manager.startDownload(url, checkpointsDirectory, 'model.safetensors')); + const loraDownload = expectStartOk(manager.startDownload(url, lorasDirectory, 'model.safetensors')); + + expect(checkpointDownload.downloadId).toBe(path.join(checkpointsDirectory, 'model.safetensors')); + expect(loraDownload.downloadId).toBe(path.join(lorasDirectory, 'model.safetensors')); + expect(downloadURL).toHaveBeenCalledTimes(2); + expect([...getDownloads(manager).keys()]).toEqual([checkpointDownload.downloadId, loraDownload.downloadId]); + + const willDownload = defaultSessionOn.mock.calls[0][1] as ( + event: unknown, + item: { + getURLChain: () => string[]; + getTotalBytes: () => number; + getSavePath: () => string; + setSavePath: ReturnType; + on: ReturnType; + once: ReturnType; + } + ) => void; + const firstItem = { + getURLChain: () => [url], + getTotalBytes: () => 10, + getSavePath: () => '', + setSavePath: vi.fn(), + on: vi.fn(), + once: vi.fn(), + }; + const secondItem = { + getURLChain: () => [url], + getTotalBytes: () => 10, + getSavePath: () => '', + setSavePath: vi.fn(), + on: vi.fn(), + once: vi.fn(), + }; + + willDownload({}, firstItem); + willDownload({}, secondItem); + + expect(firstItem.setSavePath).toHaveBeenCalledWith( + path.join(checkpointsDirectory, 'Unconfirmed model.safetensors.tmp') + ); + expect(secondItem.setSavePath).toHaveBeenCalledWith(path.join(lorasDirectory, 'Unconfirmed model.safetensors.tmp')); + }); + it('rejects relative save paths that escape the models directory', () => { const modelsDirectory = path.resolve('/mock/models'); const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); mockExistingPaths(modelsDirectory); - expect(manager.startDownload('https://example.com/model.safetensors', '../tmp', 'model.safetensors')).toBe(false); + expectStartFailed(manager.startDownload('https://example.com/model.safetensors', '../tmp', 'model.safetensors')); expect(downloadURL).not.toHaveBeenCalled(); }); @@ -110,9 +180,9 @@ describe('DownloadManager', () => { const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); mockExistingPaths(modelsDirectory, path.resolve('/tmp')); - expect( + expectStartFailed( manager.startDownload('https://example.com/model.safetensors', path.resolve('/tmp'), 'model.safetensors') - ).toBe(false); + ); expect(downloadURL).not.toHaveBeenCalled(); }); @@ -123,13 +193,13 @@ describe('DownloadManager', () => { const manager = DownloadManager.getInstance(mainWindow as never, path.resolve('/Mock/Models')); mockExistingPaths(path.resolve('/Mock/Models'), path.resolve('/mock/models/ipadapter')); - expect( + expectStartOk( manager.startDownload( 'https://example.com/model.safetensors', path.resolve('/mock/models/ipadapter'), 'model.safetensors' ) - ).toBe(true); + ); expect(downloadURL).toHaveBeenCalledWith('https://example.com/model.safetensors'); }); @@ -140,7 +210,7 @@ describe('DownloadManager', () => { const newSubdir = path.join(modelsDirectory, 'latent_upscale_models'); mockExistingPaths(modelsDirectory); - expect(manager.startDownload(url, newSubdir, 'model.safetensors')).toBe(true); + expectStartOk(manager.startDownload(url, newSubdir, 'model.safetensors')); expect(fs.mkdirSync).toHaveBeenCalledWith(newSubdir, { recursive: true }); expect(downloadURL).toHaveBeenCalledWith(url); }); @@ -152,11 +222,12 @@ describe('DownloadManager', () => { const url = 'https://example.com/model.safetensors'; mockExistingPaths(modelsDirectory, checkpointsDirectory); - expect(manager.startDownload(url, checkpointsDirectory, 'model.safetensors')).toBe(true); + const download = expectStartOk(manager.startDownload(url, checkpointsDirectory, 'model.safetensors')); expect(mainWindow.send).toHaveBeenCalledWith( IPC_CHANNELS.DOWNLOAD_PROGRESS, expect.objectContaining({ + downloadId: download.downloadId, url, filename: 'model.safetensors', savePath: path.join(checkpointsDirectory, 'model.safetensors'), @@ -175,33 +246,13 @@ describe('DownloadManager', () => { const checkpointsDirectory = path.join(modelsDirectory, 'checkpoints'); const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); mockExistingPaths(modelsDirectory, checkpointsDirectory); - const downloads = ( - manager as unknown as { - downloads: Map< - string, - { - url: string; - filename: string; - directoryPath: string; - savePath: string; - tempPath: string; - progress: number; - status: DownloadStatus; - message?: string; - receivedBytes: number; - totalBytes: number; - item: { - canResume: () => boolean; - resume: () => void; - } | null; - } - >; - } - ).downloads; + const downloads = getDownloads(manager); const resume = vi.fn(); const url = 'https://example.com/model.safetensors'; + const downloadId = path.join(checkpointsDirectory, 'model.safetensors'); - downloads.set(url, { + downloads.set(downloadId, { + downloadId, url, filename: 'model.safetensors', directoryPath: checkpointsDirectory, @@ -215,10 +266,10 @@ describe('DownloadManager', () => { item: { canResume: () => true, resume, - }, + } as unknown as Download['item'], }); - expect(manager.startDownload(url, checkpointsDirectory, 'model.safetensors')).toBe(true); + expectStartOk(manager.startDownload(url, checkpointsDirectory, 'model.safetensors')); expect(resume).toHaveBeenCalledOnce(); expect(downloadURL).not.toHaveBeenCalled(); @@ -229,9 +280,9 @@ describe('DownloadManager', () => { const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); mockExistingPaths(modelsDirectory, path.resolve('/tmp')); - expect( + expectStartFailed( manager.startDownload('https://example.com/model.safetensors', path.resolve('/tmp/evil'), 'model.safetensors') - ).toBe(false); + ); expect(fs.mkdirSync).not.toHaveBeenCalled(); expect(downloadURL).not.toHaveBeenCalled(); }); @@ -251,9 +302,7 @@ describe('DownloadManager', () => { }); const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); - expect(manager.startDownload('https://example.com/model.safetensors', symlinkPath, 'model.safetensors')).toBe( - false - ); + expectStartFailed(manager.startDownload('https://example.com/model.safetensors', symlinkPath, 'model.safetensors')); expect(downloadURL).not.toHaveBeenCalled(); }); @@ -273,7 +322,7 @@ describe('DownloadManager', () => { }); const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); - expect(manager.startDownload('https://example.com/model.safetensors', nestedPath, 'model.safetensors')).toBe(false); + expectStartFailed(manager.startDownload('https://example.com/model.safetensors', nestedPath, 'model.safetensors')); expect(fs.mkdirSync).not.toHaveBeenCalled(); expect(downloadURL).not.toHaveBeenCalled(); }); @@ -283,30 +332,13 @@ describe('DownloadManager', () => { const checkpointsDirectory = path.join(modelsDirectory, 'checkpoints'); const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); mockExistingPaths(modelsDirectory, checkpointsDirectory); - const downloads = ( - manager as unknown as { - downloads: Map< - string, - { - url: string; - filename: string; - directoryPath: string; - savePath: string; - tempPath: string; - progress: number; - status: DownloadStatus; - message?: string; - receivedBytes: number; - totalBytes: number; - item: { canResume: () => boolean; resume: () => void }; - } - >; - } - ).downloads; + const downloads = getDownloads(manager); const resume = vi.fn(); const url = 'https://example.com/model.safetensors'; + const downloadId = path.join(checkpointsDirectory, 'model.safetensors'); - downloads.set(url, { + downloads.set(downloadId, { + downloadId, url, filename: 'model.safetensors', directoryPath: checkpointsDirectory, @@ -320,10 +352,10 @@ describe('DownloadManager', () => { item: { canResume: () => false, resume, - }, + } as unknown as Download['item'], }); - manager.resumeDownload(url); + manager.resumeDownload(downloadId); expect(resume).not.toHaveBeenCalled(); expect(downloadURL).toHaveBeenCalledWith(url); @@ -334,33 +366,11 @@ describe('DownloadManager', () => { const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); const url = 'https://example.com/model.safetensors'; const savePath = path.join(modelsDirectory, 'checkpoints', 'model.safetensors'); - const downloads = ( - manager as unknown as { - downloads: Map< - string, - { - url: string; - filename: string; - directoryPath: string; - savePath: string; - tempPath: string; - progress: number; - status: DownloadStatus; - message?: string; - receivedBytes: number; - totalBytes: number; - item: { - getState: () => 'progressing'; - getReceivedBytes: () => number; - getTotalBytes: () => number; - isPaused: () => boolean; - } | null; - } - >; - } - ).downloads; + const downloads = getDownloads(manager); + const downloadId = savePath; - downloads.set(url, { + downloads.set(downloadId, { + downloadId, url, filename: 'model.safetensors', directoryPath: path.dirname(savePath), @@ -376,11 +386,12 @@ describe('DownloadManager', () => { getReceivedBytes: () => 5, getTotalBytes: () => 10, isPaused: () => false, - }, + } as unknown as Download['item'], }); expect(manager.getAllDownloads()).toEqual([ { + downloadId, url, filename: 'model.safetensors', savePath, From 95abf67662119d9763ffcf33d143a0e2e587b25a Mon Sep 17 00:00:00 2001 From: Benjamin Lu Date: Mon, 4 May 2026 08:26:19 -0700 Subject: [PATCH 3/4] fix: preserve download id path casing --- src/models/DownloadManager.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/DownloadManager.ts b/src/models/DownloadManager.ts index 2e995fe0c..d41705bac 100644 --- a/src/models/DownloadManager.ts +++ b/src/models/DownloadManager.ts @@ -351,7 +351,7 @@ export class DownloadManager { } private createDownloadId(savePath: string): string { - return this.getPathForComparison(path.resolve(savePath)); + return path.resolve(savePath); } private enqueuePendingDownload(url: string, downloadId: string): void { From 961b390e3aed0c2b3b909772dd090fcfe288e57b Mon Sep 17 00:00:00 2001 From: Benjamin Lu Date: Mon, 4 May 2026 10:35:27 -0700 Subject: [PATCH 4/4] fix: harden desktop download state transitions --- src/models/DownloadManager.ts | 65 +++++- tests/unit/models/DownloadManager.test.ts | 228 +++++++++++++++++++--- 2 files changed, 256 insertions(+), 37 deletions(-) diff --git a/src/models/DownloadManager.ts b/src/models/DownloadManager.ts index d41705bac..4b28d4302 100644 --- a/src/models/DownloadManager.ts +++ b/src/models/DownloadManager.ts @@ -109,7 +109,13 @@ export class DownloadManager { log.info(`Successfully renamed ${download.tempPath} to ${download.savePath}`); } catch (error) { log.error('Failed to rename downloaded file. Deleting temp file.', error); - fs.unlinkSync(download.tempPath); + this.deleteTempFile(download.tempPath); + download.item = null; + download.progress = this.calculateProgress(receivedBytes, totalBytes); + download.status = DownloadStatus.ERROR; + download.message = `Failed to finalize downloaded file: ${this.getErrorMessage(error)}`; + this.reportProgress(this.toDownloadState(download)); + return; } download.item = null; download.progress = 1; @@ -203,10 +209,17 @@ export class DownloadManager { receivedBytes: 0, totalBytes: 0, }; + existingCompletedDownload.url = url; + existingCompletedDownload.directoryPath = normalizedDirectoryPath; + existingCompletedDownload.savePath = localSavePath; + existingCompletedDownload.tempPath = this.getTempPath(filename, normalizedDirectoryPath); + existingCompletedDownload.filename = filename; existingCompletedDownload.progress = 1; existingCompletedDownload.status = DownloadStatus.COMPLETED; existingCompletedDownload.message = undefined; existingCompletedDownload.item = null; + existingCompletedDownload.receivedBytes = 0; + existingCompletedDownload.totalBytes = 0; this.downloads.set(downloadId, existingCompletedDownload); const downloadState = this.toDownloadState(existingCompletedDownload); this.reportProgress(downloadState); @@ -217,10 +230,13 @@ export class DownloadManager { if (existingDownload) { log.info('Download already exists'); if (existingDownload.status === DownloadStatus.PAUSED) { - this.resumeDownload(downloadId); - return { ok: true, download: this.toDownloadState(existingDownload) }; + const resumedDownload = this.resumeDownloadWithState(downloadId); + if (resumedDownload) return resumedDownload; + this.deleteTempFile(existingDownload.tempPath); + this.downloads.delete(downloadId); } else if ( existingDownload.status === DownloadStatus.CANCELLED || + existingDownload.status === DownloadStatus.COMPLETED || existingDownload.status === DownloadStatus.ERROR ) { this.deleteTempFile(existingDownload.tempPath); @@ -280,15 +296,31 @@ export class DownloadManager { } resumeDownload(downloadIdOrUrl: string): void { + this.resumeDownloadWithState(downloadIdOrUrl); + } + + private resumeDownloadWithState(downloadIdOrUrl: string): StartDownloadResult | undefined { const download = this.findDownload(downloadIdOrUrl); - if (!download?.item) return; + if (!download) return undefined; + + if (!download.item) { + this.deleteTempFile(download.tempPath); + this.downloads.delete(download.downloadId); + return this.startDownload(download.url, download.directoryPath, download.filename); + } if (download.item.canResume()) { log.info('Resuming download'); download.item.resume(); + download.status = DownloadStatus.IN_PROGRESS; + download.message = undefined; + const downloadState = this.toDownloadState(download); + this.reportProgress(downloadState); + return { ok: true, download: downloadState }; } else { + this.deleteTempFile(download.tempPath); this.downloads.delete(download.downloadId); - this.startDownload(download.url, download.directoryPath, download.filename); + return this.startDownload(download.url, download.directoryPath, download.filename); } } @@ -316,6 +348,7 @@ export class DownloadManager { } catch (error) { log.error(`Failed to delete file ${tempPath}:`, error); } + this.downloads.delete(this.createDownloadId(localSavePath)); return true; } @@ -362,14 +395,26 @@ export class DownloadManager { private takePendingDownload(url: string): Download | undefined { const pendingDownloadIds = this.pendingDownloadIdsByUrl.get(url); - const downloadId = pendingDownloadIds?.shift(); + while (pendingDownloadIds?.length) { + const downloadId = pendingDownloadIds.shift(); + const pendingDownload = downloadId ? this.downloads.get(downloadId) : undefined; + if (pendingDownload?.status === DownloadStatus.PENDING && pendingDownload.item === null) { + if (pendingDownloadIds.length === 0) { + this.pendingDownloadIdsByUrl.delete(url); + } + return pendingDownload; + } + } if (pendingDownloadIds?.length === 0) { this.pendingDownloadIdsByUrl.delete(url); } - if (downloadId) { - return this.downloads.get(downloadId); - } - return [...this.downloads.values()].find((download) => download.url === url && download.item === null); + return [...this.downloads.values()].find( + (download) => download.url === url && download.status === DownloadStatus.PENDING && download.item === null + ); + } + + private getErrorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); } private findDownload(downloadIdOrUrl: string): Download | undefined { diff --git a/tests/unit/models/DownloadManager.test.ts b/tests/unit/models/DownloadManager.test.ts index 3eebdbbad..386936d37 100644 --- a/tests/unit/models/DownloadManager.test.ts +++ b/tests/unit/models/DownloadManager.test.ts @@ -37,6 +37,40 @@ function getDownloads(manager: unknown): Map { return (manager as { downloads: Map }).downloads; } +interface MockDownloadItem { + getURLChain: () => string[]; + getTotalBytes: () => number; + getReceivedBytes: () => number; + getSavePath: () => string; + setSavePath: ReturnType; + on: ReturnType; + once: ReturnType; +} + +function getWillDownloadHandler(defaultSessionOn: ReturnType) { + return defaultSessionOn.mock.calls[0][1] as (event: unknown, item: MockDownloadItem) => void; +} + +function createMockDownloadItem(url: string, receivedBytes = 0, totalBytes = 10): MockDownloadItem { + return { + getURLChain: () => [url], + getTotalBytes: () => totalBytes, + getReceivedBytes: () => receivedBytes, + getSavePath: () => '', + setSavePath: vi.fn(), + on: vi.fn(), + once: vi.fn(), + }; +} + +function getDoneHandler(item: MockDownloadItem) { + const doneCall = item.once.mock.calls.find(([eventName]) => eventName === 'done'); + if (!doneCall) { + throw new Error('Expected done handler to be registered'); + } + return doneCall[1] as (event: unknown, state: string) => void; +} + describe('DownloadManager', () => { let DownloadManager: typeof import('@/models/DownloadManager').DownloadManager; let defaultSessionOn: ReturnType; @@ -129,33 +163,9 @@ describe('DownloadManager', () => { expect(downloadURL).toHaveBeenCalledTimes(2); expect([...getDownloads(manager).keys()]).toEqual([checkpointDownload.downloadId, loraDownload.downloadId]); - const willDownload = defaultSessionOn.mock.calls[0][1] as ( - event: unknown, - item: { - getURLChain: () => string[]; - getTotalBytes: () => number; - getSavePath: () => string; - setSavePath: ReturnType; - on: ReturnType; - once: ReturnType; - } - ) => void; - const firstItem = { - getURLChain: () => [url], - getTotalBytes: () => 10, - getSavePath: () => '', - setSavePath: vi.fn(), - on: vi.fn(), - once: vi.fn(), - }; - const secondItem = { - getURLChain: () => [url], - getTotalBytes: () => 10, - getSavePath: () => '', - setSavePath: vi.fn(), - on: vi.fn(), - once: vi.fn(), - }; + const willDownload = getWillDownloadHandler(defaultSessionOn); + const firstItem = createMockDownloadItem(url); + const secondItem = createMockDownloadItem(url); willDownload({}, firstItem); willDownload({}, secondItem); @@ -241,6 +251,78 @@ describe('DownloadManager', () => { ); }); + it('reports an error when a completed download cannot be finalized', () => { + const modelsDirectory = path.resolve('/mock/models'); + const checkpointsDirectory = path.join(modelsDirectory, 'checkpoints'); + const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); + const url = 'https://example.com/model.safetensors'; + const tempPath = path.join(checkpointsDirectory, 'Unconfirmed model.safetensors.tmp'); + mockExistingPaths(modelsDirectory, checkpointsDirectory, tempPath); + vi.mocked(fs.renameSync).mockImplementation(() => { + throw new Error('rename failed'); + }); + + const download = expectStartOk(manager.startDownload(url, checkpointsDirectory, 'model.safetensors')); + const item = createMockDownloadItem(url, 5, 10); + getWillDownloadHandler(defaultSessionOn)({}, item); + + getDoneHandler(item)({}, 'completed'); + + expect(fs.unlinkSync).toHaveBeenCalledWith(tempPath); + expect(mainWindow.send).toHaveBeenLastCalledWith( + IPC_CHANNELS.DOWNLOAD_PROGRESS, + expect.objectContaining({ + downloadId: download.downloadId, + progress: 0.5, + status: DownloadStatus.ERROR, + state: DownloadStatus.ERROR, + message: 'Failed to finalize downloaded file: rename failed', + }) + ); + }); + + it('refreshes completed snapshots when an existing file is requested again', () => { + const modelsDirectory = path.resolve('/mock/models'); + const checkpointsDirectory = path.join(modelsDirectory, 'checkpoints'); + const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); + const url = 'https://example.com/model.safetensors'; + const filename = 'model.safetensors'; + const savePath = path.join(checkpointsDirectory, filename); + const downloadId = savePath; + mockExistingPaths(modelsDirectory, checkpointsDirectory, savePath); + + getDownloads(manager).set(downloadId, { + downloadId, + url: 'https://example.com/stale.safetensors', + filename: 'stale.safetensors', + directoryPath: path.join(modelsDirectory, 'loras'), + savePath: path.join(modelsDirectory, 'loras', 'stale.safetensors'), + tempPath: path.join(modelsDirectory, 'loras', 'Unconfirmed stale.safetensors.tmp'), + progress: 0.4, + status: DownloadStatus.COMPLETED, + message: 'stale', + receivedBytes: 4, + totalBytes: 10, + item: null, + }); + + const download = expectStartOk(manager.startDownload(url, checkpointsDirectory, filename)); + + expect(download).toEqual( + expect.objectContaining({ + downloadId, + url, + filename, + savePath, + progress: 1, + status: DownloadStatus.COMPLETED, + receivedBytes: 0, + totalBytes: 0, + }) + ); + expect(downloadURL).not.toHaveBeenCalled(); + }); + it('resumes paused downloads without starting a second download', () => { const modelsDirectory = path.resolve('/mock/models'); const checkpointsDirectory = path.join(modelsDirectory, 'checkpoints'); @@ -265,6 +347,7 @@ describe('DownloadManager', () => { totalBytes: 10, item: { canResume: () => true, + isPaused: () => false, resume, } as unknown as Download['item'], }); @@ -275,6 +358,37 @@ describe('DownloadManager', () => { expect(downloadURL).not.toHaveBeenCalled(); }); + it('restarts completed downloads when the model file was deleted in the same session', () => { + const modelsDirectory = path.resolve('/mock/models'); + const checkpointsDirectory = path.join(modelsDirectory, 'checkpoints'); + const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); + const url = 'https://example.com/model.safetensors'; + const filename = 'model.safetensors'; + const savePath = path.join(checkpointsDirectory, filename); + const downloadId = savePath; + mockExistingPaths(modelsDirectory, checkpointsDirectory); + + getDownloads(manager).set(downloadId, { + downloadId, + url, + filename, + directoryPath: checkpointsDirectory, + savePath, + tempPath: path.join(checkpointsDirectory, 'Unconfirmed model.safetensors.tmp'), + progress: 1, + status: DownloadStatus.COMPLETED, + message: undefined, + receivedBytes: 10, + totalBytes: 10, + item: null, + }); + + const download = expectStartOk(manager.startDownload(url, checkpointsDirectory, filename)); + + expect(download.status).toBe(DownloadStatus.PENDING); + expect(downloadURL).toHaveBeenCalledWith(url); + }); + it('does not create directories outside models directory', () => { const modelsDirectory = path.resolve('/mock/models'); const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); @@ -361,6 +475,66 @@ describe('DownloadManager', () => { expect(downloadURL).toHaveBeenCalledWith(url); }); + it('returns the restarted snapshot when startDownload sees a paused download that cannot resume', () => { + const modelsDirectory = path.resolve('/mock/models'); + const checkpointsDirectory = path.join(modelsDirectory, 'checkpoints'); + const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); + mockExistingPaths(modelsDirectory, checkpointsDirectory); + const downloads = getDownloads(manager); + const url = 'https://example.com/model.safetensors'; + const downloadId = path.join(checkpointsDirectory, 'model.safetensors'); + + downloads.set(downloadId, { + downloadId, + url, + filename: 'model.safetensors', + directoryPath: checkpointsDirectory, + savePath: path.join(checkpointsDirectory, 'model.safetensors'), + tempPath: path.join(checkpointsDirectory, 'Unconfirmed model.safetensors.tmp'), + progress: 0.5, + status: DownloadStatus.PAUSED, + message: undefined, + receivedBytes: 5, + totalBytes: 10, + item: { + canResume: () => false, + resume: vi.fn(), + } as unknown as Download['item'], + }); + + const download = expectStartOk(manager.startDownload(url, checkpointsDirectory, 'model.safetensors')); + + expect(download.status).toBe(DownloadStatus.PENDING); + expect(downloadURL).toHaveBeenCalledWith(url); + }); + + it('does not bind will-download events to terminal rows with matching URLs', () => { + const modelsDirectory = path.resolve('/mock/models'); + const checkpointsDirectory = path.join(modelsDirectory, 'checkpoints'); + const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); + const url = 'https://example.com/model.safetensors'; + const item = createMockDownloadItem(url); + + getDownloads(manager).set(path.join(checkpointsDirectory, 'model.safetensors'), { + downloadId: path.join(checkpointsDirectory, 'model.safetensors'), + url, + filename: 'model.safetensors', + directoryPath: checkpointsDirectory, + savePath: path.join(checkpointsDirectory, 'model.safetensors'), + tempPath: path.join(checkpointsDirectory, 'Unconfirmed model.safetensors.tmp'), + progress: 1, + status: DownloadStatus.COMPLETED, + message: undefined, + receivedBytes: 10, + totalBytes: 10, + item: null, + }); + + getWillDownloadHandler(defaultSessionOn)({}, item); + + expect(item.setSavePath).not.toHaveBeenCalled(); + }); + it('returns canonical snapshots from getAllDownloads while preserving legacy fields', () => { const modelsDirectory = path.resolve('/mock/models'); const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory);