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 c2a666d41..4b28d4302 100644 --- a/src/models/DownloadManager.ts +++ b/src/models/DownloadManager.ts @@ -9,31 +9,41 @@ 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. directoryPath: string; savePath: string; item: DownloadItem | null; + progress: number; + status: DownloadStatus; + message?: string; + receivedBytes: number; + totalBytes: number; } export interface DownloadState { + downloadId: string; 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; -} +export type StartDownloadResult = + | { ok: true; download: DownloadState } + | { ok: false; error: string; download: DownloadState }; /** * Singleton class that manages downloading model checkpoints for ComfyUI. @@ -41,177 +51,276 @@ interface DownloadReport { 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; - 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); 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; } - 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)); } }); }); } - 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: normalizedDirectoryPath, + savePath: localSavePath, filename, progress: 0, status: DownloadStatus.ERROR, message: 'Save path is not in models directory', - }); - return false; + state: DownloadStatus.ERROR, + receivedBytes: 0, + totalBytes: 0, + isPaused: 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: normalizedDirectoryPath, + savePath: localSavePath, filename, progress: 0, status: DownloadStatus.ERROR, - message: validationResult.error, - }); - return false; + message: errorMessage, + state: DownloadStatus.ERROR, + receivedBytes: 0, + totalBytes: 0, + isPaused: false, + }; + this.reportProgress(downloadState); + return { + ok: false, + error: errorMessage, + download: downloadState, + }; } if (fs.existsSync(localSavePath)) { log.info(`File ${filename} already exists, skipping download`); - return true; + const existingCompletedDownload = this.downloads.get(downloadId) ?? { + downloadId, + 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.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); + 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.item?.isPaused()) { - this.resumeDownload(url); + if (existingDownload.status === DownloadStatus.PAUSED) { + 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); + this.downloads.delete(downloadId); + } else { + return { ok: true, download: this.toDownloadState(existingDownload) }; } - return true; } log.info(`Starting download ${url} to ${localSavePath}`); - const tempPath = this.getTempPath(filename, normalizedDirectoryPath); - this.downloads.set(url, { + const download: Download = { + downloadId, 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(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); - if (!download?.item) return; + cancelDownload(downloadIdOrUrl: string): void { + const download = this.findDownload(downloadIdOrUrl); + 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 { - 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); - if (!download?.item) return; + resumeDownload(downloadIdOrUrl: string): void { + this.resumeDownloadWithState(downloadIdOrUrl); + } + + private resumeDownloadWithState(downloadIdOrUrl: string): StartDownloadResult | undefined { + const download = this.findDownload(downloadIdOrUrl); + 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.downloads.delete(url); - this.startDownload(download.url, download.directoryPath, download.filename); + this.deleteTempFile(download.tempPath); + this.downloads.delete(download.downloadId); + return this.startDownload(download.url, download.directoryPath, download.filename); } } @@ -239,40 +348,90 @@ export class DownloadManager { } catch (error) { log.error(`Failed to delete file ${tempPath}:`, error); } + this.downloads.delete(this.createDownloadId(localSavePath)); return true; } 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, - })); + 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 { + downloadId: download.downloadId, + 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 createDownloadId(savePath: string): string { + return 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 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; + private takePendingDownload(url: string): Download | undefined { + const pendingDownloadIds = this.pendingDownloadIdsByUrl.get(url); + 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); } + return [...this.downloads.values()].find( + (download) => download.url === url && download.status === DownloadStatus.PENDING && download.item === null + ); } - private getTempPath(filename: string, directoryPath: string): string { - return path.join(directoryPath, `Unconfirmed ${filename}.tmp`); + private getErrorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); + } + + 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)) { + fs.unlinkSync(tempPath); + } + } catch (error) { + log.error(`Failed to delete temp file ${tempPath}:`, error); + } } // Only allow .safetensors files to be downloaded. @@ -381,7 +540,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}` ); @@ -408,9 +567,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 f795f38e8..c3155f9f8 100644 --- a/src/preload.ts +++ b/src/preload.ts @@ -3,10 +3,10 @@ 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'; +import type { DownloadState, StartDownloadResult } from './main_types'; import type { DesktopInstallState, DesktopWindowStyle } 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 { @@ -198,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 884a34b75..386936d37 100644 --- a/tests/unit/models/DownloadManager.test.ts +++ b/tests/unit/models/DownloadManager.test.ts @@ -2,6 +2,9 @@ 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 type { Download, StartDownloadResult } from '@/models/DownloadManager'; + import { electronMock } from '../setup'; vi.mock('node:fs'); @@ -14,6 +17,60 @@ 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; +} + +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; @@ -62,16 +119,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', () => { @@ -80,26 +134,54 @@ 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 = getWillDownloadHandler(defaultSessionOn); + const firstItem = createMockDownloadItem(url); + const secondItem = createMockDownloadItem(url); + + 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(); }); @@ -108,9 +190,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(); }); @@ -121,13 +203,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'); }); @@ -138,19 +220,183 @@ 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); }); + 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); + + 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'), + progress: 0, + status: DownloadStatus.PENDING, + state: DownloadStatus.PENDING, + receivedBytes: 0, + totalBytes: 0, + isPaused: false, + }) + ); + }); + + 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'); + const manager = DownloadManager.getInstance(mainWindow as never, modelsDirectory); + mockExistingPaths(modelsDirectory, checkpointsDirectory); + const downloads = getDownloads(manager); + const resume = vi.fn(); + 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: () => true, + isPaused: () => false, + resume, + } as unknown as Download['item'], + }); + + expectStartOk(manager.startDownload(url, checkpointsDirectory, 'model.safetensors')); + + expect(resume).toHaveBeenCalledOnce(); + 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); 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(); }); @@ -170,9 +416,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(); }); @@ -192,7 +436,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(); }); @@ -202,39 +446,137 @@ 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; - 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, 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, - }, + } as unknown as Download['item'], }); - manager.resumeDownload(url); + manager.resumeDownload(downloadId); expect(resume).not.toHaveBeenCalled(); 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); + const url = 'https://example.com/model.safetensors'; + const savePath = path.join(modelsDirectory, 'checkpoints', 'model.safetensors'); + const downloads = getDownloads(manager); + const downloadId = savePath; + + downloads.set(downloadId, { + downloadId, + 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, + } as unknown as Download['item'], + }); + + expect(manager.getAllDownloads()).toEqual([ + { + downloadId, + url, + filename: 'model.safetensors', + savePath, + progress: 0.5, + status: DownloadStatus.IN_PROGRESS, + message: undefined, + state: DownloadStatus.IN_PROGRESS, + receivedBytes: 5, + totalBytes: 10, + isPaused: false, + }, + ]); + }); });