From f10d5798bb7fd5b8b57ff735991ec9983b69a33b Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 9 Jun 2026 06:07:43 -0700 Subject: [PATCH 1/3] Clean up unique browser partitions on delete + drive-aware data defaults Deleting an install never removed its browser partition. Unique-partition installs each own a persist: bucket under userData/Partitions/ that nothing else reuses, so every delete leaked one forever. handleDelete now removes the install record first, then best-effort clears the session (timeout-bounded) and deletes the partition dir (force + retries) so it can never hang or lock up the delete. Shared partitions are left untouched. Also make the large data dirs follow the drive the user picks in the Windows installer: when the app is installed to a non-system drive, installDir, the shared models/input/output root, and the download cache default to that drive instead of always landing on the system drive. Non-Windows and same-drive installs keep the existing home/userData defaults. Amp-Thread-ID: https://ampcode.com/threads/T-019eac5b-d0b9-7090-a4dd-491a74ab1fbe Co-authored-by: Amp --- .../sessionActions/delete.integration.test.ts | 188 ++++++++++++++++++ src/main/lib/ipc/sessionActions/delete.ts | 3 + src/main/lib/ipc/shared.ts | 30 ++- src/main/lib/paths.test.ts | 78 ++++++++ src/main/lib/paths.ts | 40 +++- src/main/settings.ts | 8 +- 6 files changed, 341 insertions(+), 6 deletions(-) create mode 100644 src/main/lib/ipc/sessionActions/delete.integration.test.ts create mode 100644 src/main/lib/paths.test.ts diff --git a/src/main/lib/ipc/sessionActions/delete.integration.test.ts b/src/main/lib/ipc/sessionActions/delete.integration.test.ts new file mode 100644 index 00000000..351f7fb1 --- /dev/null +++ b/src/main/lib/ipc/sessionActions/delete.integration.test.ts @@ -0,0 +1,188 @@ +// @vitest-environment node +// Integration test for the delete path's browser-partition cleanup. Runs the +// real handleDelete + deleteBrowserPartition against a seeded install tree and +// a tmp userData/Partitions dir, asserting unique partitions are removed (and +// their session cleared) while shared partitions are left untouched. +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' +import fs from 'fs' +import os from 'os' +import path from 'path' +import { EventEmitter } from 'events' +import type { InstallationRecord } from '../../../installations' + +const installationsStore = new Map() + +const h = vi.hoisted(() => { + const clearStorageData = vi.fn(async () => {}) + const fromPartition = vi.fn((_partition: string) => ({ clearStorageData })) + return { userDataDir: { value: '' }, clearStorageData, fromPartition } +}) +const { clearStorageData, fromPartition } = h + +vi.mock('electron', () => ({ + app: { + isPackaged: false, + getPath: (name: string) => (name === 'userData' ? h.userDataDir.value : os.tmpdir()), + getVersion: () => '0.0.0-test', + getLocale: () => 'en', + }, + ipcMain: { handle: vi.fn(), on: vi.fn(), off: vi.fn() }, + dialog: {}, + shell: {}, + session: { fromPartition: h.fromPartition }, + BrowserWindow: { getAllWindows: () => [] }, + nativeTheme: { on: vi.fn(), shouldUseDarkColors: false }, +})) + +vi.mock('../../i18n', () => ({ + t: (key: string, params?: Record) => + params ? `${key}:${JSON.stringify(params)}` : key, + init: vi.fn(async () => {}), + getMessages: () => ({}), + getLocale: () => 'en', + getAvailableLocales: () => [], +})) + +vi.mock('../../../settings', () => ({ + get: vi.fn((_key: string): unknown => undefined), + set: vi.fn(async () => {}), + getAll: vi.fn(() => ({})), + getMirrorConfig: vi.fn(() => ({ pypiMirror: undefined, useChineseMirrors: false })), + defaults: { modelsDirs: ['/unused'], inputDir: '/unused', outputDir: '/unused' }, +})) + +vi.mock('../../../installations', () => ({ + installationEvents: new EventEmitter(), + list: vi.fn(async () => Array.from(installationsStore.values())), + get: vi.fn(async (id: string) => installationsStore.get(id) ?? null), + update: vi.fn(async (id: string, data: Record) => { + const cur = installationsStore.get(id) + if (!cur) return null + const next = { ...cur, ...data } as InstallationRecord + installationsStore.set(id, next) + return next + }), + remove: vi.fn(async (id: string) => { installationsStore.delete(id) }), +})) + +vi.mock('../../snapshots', () => ({ + saveSnapshot: vi.fn(async () => 'noop.json'), + getSnapshotCount: vi.fn(async () => 0), + deduplicatePreUpdateSnapshot: vi.fn(async () => false), +})) +vi.mock('../../../lib/pip', () => ({ + installFilteredRequirements: vi.fn(async () => 0), +})) + +import { handleDelete } from './delete' +import { MARKER_FILE } from '../shared' + +function makeSender(): Electron.WebContents { + return { isDestroyed: () => false, send: vi.fn() } as unknown as Electron.WebContents +} + +function seedInstall(installPath: string, id: string): void { + fs.mkdirSync(installPath, { recursive: true }) + fs.writeFileSync(path.join(installPath, MARKER_FILE), id) + // A little content so deleteDir has something to walk. + fs.mkdirSync(path.join(installPath, 'ComfyUI'), { recursive: true }) + fs.writeFileSync(path.join(installPath, 'ComfyUI', 'main.py'), '# stub\n') +} + +function partitionDir(id: string): string { + return path.join(h.userDataDir.value, 'Partitions', id) +} + +describe('handleDelete browser-partition cleanup', () => { + let tmpRoot: string + + beforeEach(() => { + tmpRoot = fs.mkdtempSync(path.join(os.tmpdir(), 'delete-partition-')) + h.userDataDir.value = path.join(tmpRoot, 'userData') + fs.mkdirSync(h.userDataDir.value, { recursive: true }) + installationsStore.clear() + clearStorageData.mockClear() + fromPartition.mockClear() + }) + + afterEach(() => { + fs.rmSync(tmpRoot, { recursive: true, force: true }) + }) + + it('removes the unique install partition dir and clears its session', async () => { + const id = 'inst-unique' + const installPath = path.join(tmpRoot, 'unique-install') + seedInstall(installPath, id) + fs.mkdirSync(partitionDir(id), { recursive: true }) + fs.writeFileSync(path.join(partitionDir(id), 'cookies.db'), 'x') + + const inst = { id, name: 'unique', sourceId: 'standalone', installPath, status: 'installed', browserPartition: 'unique', createdAt: new Date(0).toISOString() } as InstallationRecord + installationsStore.set(id, inst) + + const result = await handleDelete({ event: { sender: makeSender() } as unknown as Electron.IpcMainInvokeEvent, installationId: id, inst }) + + expect(result.ok).toBe(true) + expect(fs.existsSync(installPath)).toBe(false) + expect(fs.existsSync(partitionDir(id))).toBe(false) + expect(fromPartition).toHaveBeenCalledWith(`persist:${id}`) + expect(clearStorageData).toHaveBeenCalledTimes(1) + expect(installationsStore.has(id)).toBe(false) + }) + + it('leaves the shared partition untouched and never clears a session', async () => { + const id = 'inst-shared' + const installPath = path.join(tmpRoot, 'shared-install') + seedInstall(installPath, id) + const sharedDir = partitionDir('shared') + fs.mkdirSync(sharedDir, { recursive: true }) + fs.writeFileSync(path.join(sharedDir, 'cookies.db'), 'x') + + const inst = { id, name: 'shared', sourceId: 'standalone', installPath, status: 'installed', browserPartition: 'shared', createdAt: new Date(0).toISOString() } as InstallationRecord + installationsStore.set(id, inst) + + const result = await handleDelete({ event: { sender: makeSender() } as unknown as Electron.IpcMainInvokeEvent, installationId: id, inst }) + + expect(result.ok).toBe(true) + expect(fs.existsSync(installPath)).toBe(false) + expect(fs.existsSync(sharedDir)).toBe(true) + expect(fromPartition).not.toHaveBeenCalled() + expect(clearStorageData).not.toHaveBeenCalled() + expect(installationsStore.has(id)).toBe(false) + }) + + it('still completes the delete when clearing the session storage fails', async () => { + clearStorageData.mockRejectedValueOnce(new Error('session busy')) + const id = 'inst-clear-fail' + const installPath = path.join(tmpRoot, 'clear-fail-install') + seedInstall(installPath, id) + fs.mkdirSync(partitionDir(id), { recursive: true }) + + const inst = { id, name: 'clear-fail', sourceId: 'standalone', installPath, status: 'installed', browserPartition: 'unique', createdAt: new Date(0).toISOString() } as InstallationRecord + installationsStore.set(id, inst) + + const result = await handleDelete({ event: { sender: makeSender() } as unknown as Electron.IpcMainInvokeEvent, installationId: id, inst }) + + expect(result.ok).toBe(true) + expect(fs.existsSync(installPath)).toBe(false) + // Record removal must not depend on partition cleanup succeeding. + expect(installationsStore.has(id)).toBe(false) + // rm still runs even though clearStorageData rejected. + expect(fs.existsSync(partitionDir(id))).toBe(false) + }) + + it('cleans up the unique partition even when the install dir is already gone', async () => { + const id = 'inst-gone' + const installPath = path.join(tmpRoot, 'missing-install') + fs.mkdirSync(partitionDir(id), { recursive: true }) + + const inst = { id, name: 'gone', sourceId: 'standalone', installPath, status: 'installed', browserPartition: 'unique', createdAt: new Date(0).toISOString() } as InstallationRecord + installationsStore.set(id, inst) + + const result = await handleDelete({ event: { sender: makeSender() } as unknown as Electron.IpcMainInvokeEvent, installationId: id, inst }) + + expect(result.ok).toBe(true) + expect(fs.existsSync(partitionDir(id))).toBe(false) + expect(fromPartition).toHaveBeenCalledWith(`persist:${id}`) + expect(installationsStore.has(id)).toBe(false) + }) +}) diff --git a/src/main/lib/ipc/sessionActions/delete.ts b/src/main/lib/ipc/sessionActions/delete.ts index 7e87b56b..d988fc25 100644 --- a/src/main/lib/ipc/sessionActions/delete.ts +++ b/src/main/lib/ipc/sessionActions/delete.ts @@ -5,6 +5,7 @@ import { findLockingProcesses, MARKER_FILE, makeSendProgress, + deleteBrowserPartition, } from '../shared' import type { ActionContext, ActionResult } from './types' import { withAbortableSessionAction } from './withAbortable' @@ -46,6 +47,7 @@ export async function handleDelete(ctx: ActionContext): Promise { await cleanupAdoptedLegacyDir(adoptedBaseDir, noopProgress) } await installations.remove(installationId) + await deleteBrowserPartition(inst.id, inst.browserPartition as string | undefined) return { ok: true, navigate: 'list' } } const markerPath = path.join(inst.installPath, MARKER_FILE) @@ -94,6 +96,7 @@ export async function handleDelete(ctx: ActionContext): Promise { throw err } await installations.remove(installationId) + await deleteBrowserPartition(inst.id, inst.browserPartition as string | undefined) return { ok: true, navigate: 'list' } }) } diff --git a/src/main/lib/ipc/shared.ts b/src/main/lib/ipc/shared.ts index 469fffec..efde099f 100644 --- a/src/main/lib/ipc/shared.ts +++ b/src/main/lib/ipc/shared.ts @@ -2,7 +2,7 @@ import path from 'path' import fs from 'fs' import os from 'os' import { EventEmitter } from 'events' -import { app, ipcMain, dialog, shell, BrowserWindow, nativeTheme } from 'electron' +import { app, ipcMain, dialog, shell, BrowserWindow, nativeTheme, session } from 'electron' import { execFile, spawn, execFileSync } from 'child_process' import type { ChildProcess } from 'child_process' import sources from '../../sources/index' @@ -350,6 +350,34 @@ export async function copyBrowserPartition(sourceId: string, destId: string, sou } } +/** Delete the on-disk browser partition for a deleted install. Unique-partition + * installs each own a `persist:${id}` bucket under userData/Partitions/ + * (created lazily by Electron, deep-copied on install-copy); nothing else ever + * reuses it, so it must be removed when the install is deleted or it leaks + * forever. No-op for shared-partition installs (they all share `persist:shared`, + * which must never be deleted). Best-effort: clears the session first to release + * file handles (Windows locks LevelDB/IndexedDB while open). */ +export async function deleteBrowserPartition(id: string, browserPartition?: string): Promise { + if (browserPartition !== 'unique') return + // Best-effort, fully bounded: this runs after the install record is already + // removed, so it must never hang the delete operation or hold its lock. + // clearStorageData has no hard completion guarantee, so race it against a + // timeout; rm fails fast (force) with a few transient-lock retries. + try { + const cleared = session.fromPartition(`persist:${id}`).clearStorageData() + const timeout = new Promise((resolve) => setTimeout(resolve, 5000)) + await Promise.race([cleared, timeout]) + } catch (err) { + console.warn('Failed to clear browser partition storage:', (err as Error).message) + } + const partitionDir = path.join(app.getPath('userData'), 'Partitions', id) + try { + await fs.promises.rm(partitionDir, { recursive: true, force: true, maxRetries: 3, retryDelay: 100 }) + } catch (err) { + console.warn('Failed to delete browser partition:', (err as Error).message) + } +} + export async function performCopy( inst: InstallationRecord, name: string, diff --git a/src/main/lib/paths.test.ts b/src/main/lib/paths.test.ts new file mode 100644 index 00000000..3d305028 --- /dev/null +++ b/src/main/lib/paths.test.ts @@ -0,0 +1,78 @@ +import path from 'path' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +// Drive-aware default helpers in paths.ts. The drive-redirect branch only +// activates on Windows (NSIS lets the user pick an install drive); on other +// platforms the helpers fall back to the home dir. + +const HOME = path.join('/mock', 'home') + +let exePath = '' + +beforeEach(() => { + vi.resetModules() + vi.doMock('electron', () => ({ + app: { + getPath: (name: string) => { + if (name === 'home') return HOME + if (name === 'exe') return exePath + return HOME // userData fallback (unused by these helpers on win) + }, + }, + })) +}) + +afterEach(() => { + vi.unstubAllGlobals() + vi.doUnmock('electron') +}) + +async function loadPaths() { + return await import('./paths') +} + +function stubPlatform(platform: NodeJS.Platform): void { + vi.stubGlobal('process', { ...process, platform }) +} + +describe('drive-aware defaults', () => { + it('redirects data dirs to the app drive when the app is on a non-home drive (win32)', async () => { + stubPlatform('win32') + exePath = 'D:\\Programs\\Comfy Desktop\\Comfy Desktop.exe' + const p = await loadPaths() + + expect(p.defaultDataRoot()).toBe('D:\\') + expect(p.builtinDefaultInstallDir()).toBe(path.join('D:\\', 'ComfyUI-Installs')) + expect(p.defaultDownloadCacheDir()).toBe(path.join('D:\\', 'ComfyUI-Cache', 'download-cache')) + }) + + it('falls back to home when the app is on the same drive as home (win32)', async () => { + stubPlatform('win32') + // Home parsed by win32 has no drive root here, so emulate a real same-drive + // case by putting both on C:. + exePath = 'C:\\Program Files\\Comfy Desktop\\Comfy Desktop.exe' + vi.resetModules() + vi.doMock('electron', () => ({ + app: { + getPath: (name: string) => { + if (name === 'home') return 'C:\\Users\\test' + if (name === 'exe') return exePath + return 'C:\\Users\\test' + }, + }, + })) + const p = await import('./paths') + + expect(p.defaultDataRoot()).toBe('C:\\Users\\test') + expect(p.builtinDefaultInstallDir()).toBe(path.join('C:\\Users\\test', 'ComfyUI-Installs')) + }) + + it('never redirects on non-Windows platforms', async () => { + stubPlatform('linux') + exePath = '/opt/Comfy Desktop/comfy-desktop' + const p = await loadPaths() + + expect(p.defaultDataRoot()).toBe(HOME) + expect(p.builtinDefaultInstallDir()).toBe(path.join(HOME, 'ComfyUI-Installs')) + }) +}) diff --git a/src/main/lib/paths.ts b/src/main/lib/paths.ts index a6d03453..9c29e752 100644 --- a/src/main/lib/paths.ts +++ b/src/main/lib/paths.ts @@ -49,9 +49,47 @@ export function stateDir(): string { return app.getPath("userData"); } +/** Windows only: the root of the drive the app was installed onto, when it + * differs from the user's home drive. The NSIS installer lets the user pick an + * install location (e.g. `D:\...`); this lets ComfyUI's large data dirs follow + * that drive instead of always landing on the system drive. Returns null on + * non-Windows, or when the app lives on the home drive (home is already the + * right default there). */ +function selectedInstallDrive(): string | null { + if (process.platform !== "win32") return null; + try { + // Explicit win32 parsing so the drive is extracted correctly regardless of + // the host the code is exercised on (matches the platform `path` at runtime). + const exeDrive = path.win32.parse(app.getPath("exe")).root; // e.g. "D:\\" + const homeDrive = path.win32.parse(app.getPath("home")).root; // e.g. "C:\\" + if (!exeDrive) return null; + if (exeDrive.toLowerCase() === homeDrive.toLowerCase()) return null; + return exeDrive; + } catch { + return null; + } +} + +/** Base directory under which ComfyUI's large data dirs (installs, shared + * models/input/output) live by default. On Windows, when the app was installed + * to a non-system drive, this is that drive's root so data follows the user's + * chosen drive; otherwise the user's home dir. */ +export function defaultDataRoot(): string { + return selectedInstallDrive() ?? app.getPath("home"); +} + +/** Default location for the multi-GB download cache. Follows the user's chosen + * install drive on Windows; otherwise the platform cache dir (userData on + * Windows/macOS, XDG cache on Linux). */ +export function defaultDownloadCacheDir(): string { + const drive = selectedInstallDrive(); + if (drive) return path.join(drive, "ComfyUI-Cache", "download-cache"); + return path.join(cacheDir(), "download-cache"); +} + /** Built-in fallback when no `installDir` setting is configured. */ export function builtinDefaultInstallDir(): string { - return path.join(app.getPath("home"), "ComfyUI-Installs"); + return path.join(defaultDataRoot(), "ComfyUI-Installs"); } /** Resolver for the user's configured install location, injected by settings.ts. diff --git a/src/main/settings.ts b/src/main/settings.ts index f1d845f3..5df1040b 100644 --- a/src/main/settings.ts +++ b/src/main/settings.ts @@ -1,7 +1,7 @@ import path from 'path' import fs from 'fs' import { app } from 'electron' -import { configDir, cacheDir, homeDir, setInstallDirResolver } from './lib/paths' +import { configDir, homeDir, defaultDataRoot, defaultDownloadCacheDir, setInstallDirResolver } from './lib/paths' import { MODEL_FOLDER_TYPES } from './lib/models' import { readFileSafe, writeFileSafe } from './lib/safe-file' @@ -64,7 +64,7 @@ type SettingsDefaults = Pick const dataPath = path.join(configDir(), "settings.json") -const SHARED_ROOT = path.join(homeDir(), "ComfyUI-Shared") +const SHARED_ROOT = path.join(defaultDataRoot(), "ComfyUI-Shared") const SETTINGS_SCHEMA = { cacheDir: { nullable: false }, @@ -107,14 +107,14 @@ function isNullableKnownSettingKey(key: KnownSettingKey): key is NullableKnownSe } export const defaults: SettingsDefaults = { - cacheDir: path.join(cacheDir(), "download-cache"), + cacheDir: defaultDownloadCacheDir(), maxCachedDownloads: 1, // Docking-to-tray is disabled (createTray() is currently a no-op). onAppClose: "quit", modelsDirs: [path.join(SHARED_ROOT, "models")], inputDir: path.join(SHARED_ROOT, "input"), outputDir: path.join(SHARED_ROOT, "output"), - installDir: path.join(homeDir(), "ComfyUI-Installs"), + installDir: path.join(defaultDataRoot(), "ComfyUI-Installs"), } const systemDefault = defaults.modelsDirs[0]! From 02b5ac77d91b448a9d8dc0d3ba0af7045e3cf6f0 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 9 Jun 2026 06:15:59 -0700 Subject: [PATCH 2/3] Group drive-redirected data dirs under a single Comfy-Desktop folder When the app is installed to a non-system drive, installs, shared models/input/output, and the download cache now live under one \\Comfy-Desktop parent instead of three separate folders at the drive root. Home-dir and non-Windows defaults are unchanged. Amp-Thread-ID: https://ampcode.com/threads/T-019eac5b-d0b9-7090-a4dd-491a74ab1fbe Co-authored-by: Amp --- src/main/lib/paths.test.ts | 7 ++++--- src/main/lib/paths.ts | 21 ++++++++++++--------- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/src/main/lib/paths.test.ts b/src/main/lib/paths.test.ts index 3d305028..2d8e35ce 100644 --- a/src/main/lib/paths.test.ts +++ b/src/main/lib/paths.test.ts @@ -41,9 +41,10 @@ describe('drive-aware defaults', () => { exePath = 'D:\\Programs\\Comfy Desktop\\Comfy Desktop.exe' const p = await loadPaths() - expect(p.defaultDataRoot()).toBe('D:\\') - expect(p.builtinDefaultInstallDir()).toBe(path.join('D:\\', 'ComfyUI-Installs')) - expect(p.defaultDownloadCacheDir()).toBe(path.join('D:\\', 'ComfyUI-Cache', 'download-cache')) + const dataRoot = path.join('D:\\', 'Comfy-Desktop') + expect(p.defaultDataRoot()).toBe(dataRoot) + expect(p.builtinDefaultInstallDir()).toBe(path.join(dataRoot, 'ComfyUI-Installs')) + expect(p.defaultDownloadCacheDir()).toBe(path.join(dataRoot, 'ComfyUI-Cache', 'download-cache')) }) it('falls back to home when the app is on the same drive as home (win32)', async () => { diff --git a/src/main/lib/paths.ts b/src/main/lib/paths.ts index 9c29e752..9b321d98 100644 --- a/src/main/lib/paths.ts +++ b/src/main/lib/paths.ts @@ -71,19 +71,22 @@ function selectedInstallDrive(): string | null { } /** Base directory under which ComfyUI's large data dirs (installs, shared - * models/input/output) live by default. On Windows, when the app was installed - * to a non-system drive, this is that drive's root so data follows the user's - * chosen drive; otherwise the user's home dir. */ + * models/input/output, download cache) live by default. On Windows, when the + * app was installed to a non-system drive, this is a single `Comfy-Desktop` + * folder on that drive so everything stays grouped under one parent instead of + * scattering folders at the drive root; otherwise the user's home dir. */ export function defaultDataRoot(): string { - return selectedInstallDrive() ?? app.getPath("home"); + const drive = selectedInstallDrive(); + return drive ? path.join(drive, "Comfy-Desktop") : app.getPath("home"); } -/** Default location for the multi-GB download cache. Follows the user's chosen - * install drive on Windows; otherwise the platform cache dir (userData on - * Windows/macOS, XDG cache on Linux). */ +/** Default location for the multi-GB download cache. On a redirected drive it + * lives under the same `Comfy-Desktop` parent as the other data dirs; otherwise + * the platform cache dir (userData on Windows/macOS, XDG cache on Linux). */ export function defaultDownloadCacheDir(): string { - const drive = selectedInstallDrive(); - if (drive) return path.join(drive, "ComfyUI-Cache", "download-cache"); + if (selectedInstallDrive()) { + return path.join(defaultDataRoot(), "ComfyUI-Cache", "download-cache"); + } return path.join(cacheDir(), "download-cache"); } From c86c254c879bf85e3c90e21f3699aba75e1489e3 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 9 Jun 2026 17:31:16 -0700 Subject: [PATCH 3/3] feat(storage): default new Windows installs to %LOCALAPPDATA% on the system drive On a system-drive Windows install, the large data dirs (installs, shared models/input/output, download cache) previously defaulted to the user home root and the cache to roaming AppData. New installs now group under %LOCALAPPDATA%\Comfy-Desktop instead, while existing installs keep their home-root layout. The choice is classified from a home-root footprint and pinned with a one-time marker so it never flips between launches. Amp-Thread-ID: https://ampcode.com/threads/T-019eac3f-cda6-74ea-ba76-49da53cadea2 Co-authored-by: Amp --- src/main/index.ts | 3 +- src/main/lib/oem.test.ts | 4 ++ src/main/lib/paths.test.ts | 112 ++++++++++++++++++++++++++++++------- src/main/lib/paths.ts | 77 +++++++++++++++++++++++-- src/main/settings.test.ts | 4 ++ 5 files changed, 175 insertions(+), 25 deletions(-) diff --git a/src/main/index.ts b/src/main/index.ts index b59aaddd..5f05b327 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -16,7 +16,7 @@ import * as updater from './lib/updater' import * as settings from './settings' import { installAppMenu } from './menu' import * as i18n from './lib/i18n' -import { migrateXdgPaths } from './lib/paths' +import { migrateXdgPaths, persistWinDataRootChoice } from './lib/paths' import { saveWindowBounds } from './lib/windowState' import { clearLastActiveSurface, @@ -1264,6 +1264,7 @@ if (app.isPackaged && !app.requestSingleInstanceLock()) { registerPanelViewIpc() migrateXdgPaths() + persistWinDataRootChoice() registerProcessErrorHandlers() // Strip Electron's default menu before any BrowserWindow opens so diff --git a/src/main/lib/oem.test.ts b/src/main/lib/oem.test.ts index d284c6fd..467ab945 100644 --- a/src/main/lib/oem.test.ts +++ b/src/main/lib/oem.test.ts @@ -33,6 +33,10 @@ beforeEach(() => { oemRoot = path.join(programDataPath, 'Comfy Desktop', 'OEM') fs.mkdirSync(homePath, { recursive: true }) + // Home-root footprint → existing install, so Windows large-data defaults use + // the home layout these tests assert (a clean machine would default to + // %LOCALAPPDATA%\Comfy-Desktop). + fs.mkdirSync(path.join(homePath, 'ComfyUI-Installs'), { recursive: true }) fs.mkdirSync(userDataPath, { recursive: true }) fs.mkdirSync(programDataPath, { recursive: true }) diff --git a/src/main/lib/paths.test.ts b/src/main/lib/paths.test.ts index 2d8e35ce..073f62cb 100644 --- a/src/main/lib/paths.test.ts +++ b/src/main/lib/paths.test.ts @@ -1,9 +1,11 @@ +import fs from 'fs' +import os from 'os' import path from 'path' -import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -// Drive-aware default helpers in paths.ts. The drive-redirect branch only -// activates on Windows (NSIS lets the user pick an install drive); on other -// platforms the helpers fall back to the home dir. +// Drive-aware default helpers in paths.ts. The drive-redirect and +// system-drive branches only activate on Windows; on other platforms the +// helpers fall back to the home dir. const HOME = path.join('/mock', 'home') @@ -47,33 +49,105 @@ describe('drive-aware defaults', () => { expect(p.defaultDownloadCacheDir()).toBe(path.join(dataRoot, 'ComfyUI-Cache', 'download-cache')) }) - it('falls back to home when the app is on the same drive as home (win32)', async () => { - stubPlatform('win32') - // Home parsed by win32 has no drive root here, so emulate a real same-drive - // case by putting both on C:. - exePath = 'C:\\Program Files\\Comfy Desktop\\Comfy Desktop.exe' + it('never redirects on non-Windows platforms', async () => { + stubPlatform('linux') + exePath = '/opt/Comfy Desktop/comfy-desktop' + const p = await loadPaths() + + expect(p.defaultDataRoot()).toBe(HOME) + expect(p.builtinDefaultInstallDir()).toBe(path.join(HOME, 'ComfyUI-Installs')) + }) +}) + +describe('windows system-drive defaults', () => { + // Real temp dirs (the repo convention) for the home footprint and the marker + // location, so selectedInstallDrive() sees a same-drive install and the + // classifier reads genuine on-disk state. + const tmpRoot = fs.mkdtempSync(path.join(os.tmpdir(), 'comfyui-desktop-2-paths-')) + const driveRoot = path.win32.parse(tmpRoot).root // same drive as home → system-drive case + const LOCAL = path.join(tmpRoot, 'AppData', 'Local') + let homeDir = '' + let userDataDir = '' + let prevLocalAppData: string | undefined + + beforeEach(() => { + homeDir = fs.mkdtempSync(path.join(tmpRoot, 'home-')) + userDataDir = fs.mkdtempSync(path.join(tmpRoot, 'userdata-')) + exePath = path.join(driveRoot, 'Program Files', 'Comfy Desktop', 'Comfy Desktop.exe') + prevLocalAppData = process.env.LOCALAPPDATA + process.env.LOCALAPPDATA = LOCAL + vi.stubGlobal('process', { ...process, platform: 'win32' }) + vi.resetModules() vi.doMock('electron', () => ({ app: { getPath: (name: string) => { - if (name === 'home') return 'C:\\Users\\test' + if (name === 'home') return homeDir if (name === 'exe') return exePath - return 'C:\\Users\\test' + return userDataDir // userData → marker location }, }, })) + }) + + afterEach(() => { + if (prevLocalAppData === undefined) delete process.env.LOCALAPPDATA + else process.env.LOCALAPPDATA = prevLocalAppData + }) + + afterAll(() => { + fs.rmSync(tmpRoot, { recursive: true, force: true }) + }) + + function writeMarker(mode: string): void { + fs.writeFileSync(path.join(userDataDir, 'data-location.json'), JSON.stringify({ mode }), 'utf-8') + } + + it('new install (no footprint) groups under %LOCALAPPDATA%\\Comfy-Desktop', async () => { const p = await import('./paths') - expect(p.defaultDataRoot()).toBe('C:\\Users\\test') - expect(p.builtinDefaultInstallDir()).toBe(path.join('C:\\Users\\test', 'ComfyUI-Installs')) + const root = path.join(LOCAL, 'Comfy-Desktop') + expect(p.defaultDataRoot()).toBe(root) + expect(p.builtinDefaultInstallDir()).toBe(path.join(root, 'ComfyUI-Installs')) + expect(p.defaultDownloadCacheDir()).toBe(path.join(root, 'ComfyUI-Cache', 'download-cache')) }) - it('never redirects on non-Windows platforms', async () => { - stubPlatform('linux') - exePath = '/opt/Comfy Desktop/comfy-desktop' - const p = await loadPaths() + it('existing install (home footprint) keeps home root and roaming cache', async () => { + fs.mkdirSync(path.join(homeDir, 'ComfyUI-Installs'), { recursive: true }) + const p = await import('./paths') - expect(p.defaultDataRoot()).toBe(HOME) - expect(p.builtinDefaultInstallDir()).toBe(path.join(HOME, 'ComfyUI-Installs')) + expect(p.defaultDataRoot()).toBe(homeDir) + expect(p.builtinDefaultInstallDir()).toBe(path.join(homeDir, 'ComfyUI-Installs')) + // Roaming userData/download-cache, not the grouped Comfy-Desktop cache. + expect(p.defaultDownloadCacheDir()).toBe(path.join(userDataDir, 'download-cache')) + }) + + it('a local-appdata marker wins even when legacy folders later appear', async () => { + writeMarker('local-appdata') + fs.mkdirSync(path.join(homeDir, 'ComfyUI-Shared'), { recursive: true }) + const p = await import('./paths') + + expect(p.defaultDataRoot()).toBe(path.join(LOCAL, 'Comfy-Desktop')) + }) + + it('a legacy-home marker wins even when no legacy footprint exists', async () => { + writeMarker('legacy-home') + const p = await import('./paths') + + expect(p.defaultDataRoot()).toBe(homeDir) + }) + + it('persistWinDataRootChoice writes the classified mode once', async () => { + const p = await import('./paths') + + const markerPath = path.join(userDataDir, 'data-location.json') + expect(fs.existsSync(markerPath)).toBe(false) + p.persistWinDataRootChoice() + expect(JSON.parse(fs.readFileSync(markerPath, 'utf-8'))).toEqual({ mode: 'local-appdata' }) + + // A second call must not overwrite the first decision. + fs.mkdirSync(path.join(homeDir, 'ComfyUI-Installs'), { recursive: true }) + p.persistWinDataRootChoice() + expect(JSON.parse(fs.readFileSync(markerPath, 'utf-8'))).toEqual({ mode: 'local-appdata' }) }) }) diff --git a/src/main/lib/paths.ts b/src/main/lib/paths.ts index 9b321d98..34d8a9f6 100644 --- a/src/main/lib/paths.ts +++ b/src/main/lib/paths.ts @@ -70,23 +70,90 @@ function selectedInstallDrive(): string | null { } } +/** Windows system-drive layout for the large data dirs. New installs are + * grouped under `%LOCALAPPDATA%\Comfy-Desktop` (the standard per-user, + * non-roaming spot for large app data — the home root and roaming AppData are + * both poor fits); pre-existing installs keep their original home-root layout + * so an upgrade never strands a user's data. */ +type WinSystemDriveMode = "legacy-home" | "local-appdata"; + +const DATA_ROOT_MARKER = "data-location.json"; + +let cachedWinMode: WinSystemDriveMode | undefined; + +function winDataRootMarkerPath(): string { + return path.join(dataDir(), DATA_ROOT_MARKER); +} + +/** Classify a Windows system-drive install. A persisted marker wins so the + * choice can never flip between launches; otherwise an existing home-root + * footprint means an upgrade (keep home), and a clean machine is a new install + * (use LocalAppData). */ +function classifyWinSystemDriveMode(): WinSystemDriveMode { + try { + const raw = fs.readFileSync(winDataRootMarkerPath(), "utf-8"); + const mode = (JSON.parse(raw) as { mode?: unknown }).mode; + if (mode === "legacy-home" || mode === "local-appdata") return mode; + } catch {} + const home = app.getPath("home"); + const hasLegacyFootprint = + fs.existsSync(path.join(home, "ComfyUI-Installs")) || + fs.existsSync(path.join(home, "ComfyUI-Shared")); + return hasLegacyFootprint ? "legacy-home" : "local-appdata"; +} + +function winSystemDriveRoot(): string { + cachedWinMode ??= classifyWinSystemDriveMode(); + if (cachedWinMode === "legacy-home") return app.getPath("home"); + const localAppData = + process.env.LOCALAPPDATA || path.join(app.getPath("home"), "AppData", "Local"); + return path.join(localAppData, "Comfy-Desktop"); +} + +/** Persist the resolved Windows system-drive data-root choice once, so a new + * user's location can never flip on a later launch (e.g. if the legacy folders + * appear externally). No-op off win32, on a redirected install drive, or when a + * marker already exists. Call once after startup — never at module import. */ +export function persistWinDataRootChoice(): void { + if (process.platform !== "win32") return; + if (selectedInstallDrive()) return; // redirected drive is already deterministic + const markerPath = winDataRootMarkerPath(); + if (fs.existsSync(markerPath)) return; + try { + fs.mkdirSync(path.dirname(markerPath), { recursive: true }); + const mode = cachedWinMode ?? classifyWinSystemDriveMode(); + fs.writeFileSync(markerPath, JSON.stringify({ mode }), "utf-8"); + } catch {} +} + /** Base directory under which ComfyUI's large data dirs (installs, shared * models/input/output, download cache) live by default. On Windows, when the * app was installed to a non-system drive, this is a single `Comfy-Desktop` * folder on that drive so everything stays grouped under one parent instead of - * scattering folders at the drive root; otherwise the user's home dir. */ + * scattering folders at the drive root. On a system-drive Windows install, + * new users get `%LOCALAPPDATA%\Comfy-Desktop` while existing users keep home. + * Non-Windows uses the user's home dir. */ export function defaultDataRoot(): string { const drive = selectedInstallDrive(); - return drive ? path.join(drive, "Comfy-Desktop") : app.getPath("home"); + if (drive) return path.join(drive, "Comfy-Desktop"); + if (process.platform === "win32") return winSystemDriveRoot(); + return app.getPath("home"); } -/** Default location for the multi-GB download cache. On a redirected drive it - * lives under the same `Comfy-Desktop` parent as the other data dirs; otherwise - * the platform cache dir (userData on Windows/macOS, XDG cache on Linux). */ +/** Default location for the multi-GB download cache. When the large data dirs + * are grouped under a `Comfy-Desktop` parent (a redirected drive, or a new + * Windows system-drive install), the cache lives there too; otherwise the + * platform cache dir (roaming userData on Windows/macOS, XDG cache on Linux). */ export function defaultDownloadCacheDir(): string { if (selectedInstallDrive()) { return path.join(defaultDataRoot(), "ComfyUI-Cache", "download-cache"); } + if (process.platform === "win32") { + const root = winSystemDriveRoot(); + if (root !== app.getPath("home")) { + return path.join(root, "ComfyUI-Cache", "download-cache"); + } + } return path.join(cacheDir(), "download-cache"); } diff --git a/src/main/settings.test.ts b/src/main/settings.test.ts index 3d9060cd..46c4edb8 100644 --- a/src/main/settings.test.ts +++ b/src/main/settings.test.ts @@ -17,6 +17,10 @@ const originalXdgCacheHome = process.env.XDG_CACHE_HOME process.env.XDG_CONFIG_HOME = xdgConfigHome process.env.XDG_CACHE_HOME = xdgCacheHome fs.mkdirSync(homePath, { recursive: true }) +// A home-root footprint marks this as an existing install, so on Windows the +// large-data defaults resolve to the home layout these tests assert (a clean +// machine would instead default to %LOCALAPPDATA%\Comfy-Desktop). +fs.mkdirSync(path.join(homePath, 'ComfyUI-Installs'), { recursive: true }) fs.mkdirSync(userDataPath, { recursive: true }) fs.mkdirSync(adminHomePath, { recursive: true }) fs.mkdirSync(adminUserDataPath, { recursive: true })