Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add installation validation #616

Merged
merged 2 commits into from
Jan 20, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
152 changes: 107 additions & 45 deletions src/install/installationManager.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { app, ipcMain } from 'electron';
import { app, dialog, ipcMain } from 'electron';
import log from 'electron-log/main';

import { IPC_CHANNELS } from '../constants';
Expand All @@ -21,32 +21,99 @@ export class InstallationManager {
* Ensures that ComfyUI is installed and ready to run.
*
* First checks for an existing installation and validates it. If missing or invalid, a fresh install is started.
* Will not resolve until the installation is valid.
* @returns A valid {@link ComfyInstallation} object.
*/
async ensureInstalled(): Promise<ComfyInstallation> {
const installation = ComfyInstallation.fromConfig();
log.verbose(`Install state: ${installation?.state ?? 'not installed'}`);
log.info(`Install state: ${installation?.state ?? 'not installed'}`);

// Fresh install
if (!installation) return await this.freshInstall();

// Validate installation
const state = await installation.validate();
log.verbose(`Validated install state: ${state}`);
if (state !== 'installed') await this.resumeInstallation(installation);
try {
// Send updates to renderer
this.#setupIpc(installation);

// Resolve issues and re-run validation
if (installation.issues.size > 0) {
await this.resolveIssues(installation);
await installation.validate();
// Validate installation
const state = await installation.validate();
if (state !== 'installed') await this.resumeInstallation(installation);

// Resolve issues and re-run validation
if (installation.hasIssues) {
while (!(await this.resolveIssues(installation))) {
// Re-run validation
log.verbose('Re-validating installation.');
}
}

// Return validated installation
return installation;
} finally {
delete installation.onUpdate;
this.#removeIpcHandlers();
}
}

/** Removes all handlers created by {@link #setupIpc} */
#removeIpcHandlers() {
ipcMain.removeHandler(IPC_CHANNELS.GET_VALIDATION_STATE);
ipcMain.removeHandler(IPC_CHANNELS.VALIDATE_INSTALLATION);
ipcMain.removeHandler(IPC_CHANNELS.UV_INSTALL_REQUIREMENTS);
ipcMain.removeHandler(IPC_CHANNELS.UV_CLEAR_CACHE);
ipcMain.removeHandler(IPC_CHANNELS.UV_RESET_VENV);
}

// TODO: Confirm this is no longer possible after resolveIssues and remove.
if (!installation.basePath) throw new Error('Base path was invalid after installation validation.');
if (installation.issues.size > 0) throw new Error('Installation issues remain after validation.');
/** Set to `true` the first time an error is found during validation. @todo Move to app state singleton once impl. */
#onMaintenancePage = false;

/** Creates IPC handlers for the installation instance. */
#setupIpc(installation: ComfyInstallation) {
this.#onMaintenancePage = false;
installation.onUpdate = (data) => {
this.appWindow.send(IPC_CHANNELS.VALIDATION_UPDATE, data);

// Load maintenance page the first time any error is found.
if (!this.#onMaintenancePage && Object.values(data).includes('error')) {
this.#onMaintenancePage = true;

log.info('Validation error - loading maintenance page.');
this.appWindow.loadRenderer('maintenance').catch((error) => {
log.error('Error loading maintenance page.', error);
const message = `An error was detected with your installation, and the maintenance page could not be loaded to resolve it. The app will close now. Please reinstall if this issue persists.\n\nError message:\n\n${error}`;
dialog.showErrorBox('Critical Error', message);
app.quit();
});
}
};
const sendLogIpc = (data: string) => this.appWindow.send(IPC_CHANNELS.LOG_MESSAGE, data);

// Return validated installation
return installation;
ipcMain.handle(IPC_CHANNELS.GET_VALIDATION_STATE, () => {
installation.onUpdate?.(installation.validation);
return installation.validation;
});
ipcMain.handle(IPC_CHANNELS.VALIDATE_INSTALLATION, async () => await installation.validate());
ipcMain.handle(IPC_CHANNELS.UV_INSTALL_REQUIREMENTS, () =>
installation.virtualEnvironment.reinstallRequirements(sendLogIpc)
);
ipcMain.handle(IPC_CHANNELS.UV_CLEAR_CACHE, async () => await installation.virtualEnvironment.clearUvCache());
ipcMain.handle(IPC_CHANNELS.UV_RESET_VENV, async (): Promise<boolean> => {
const venv = installation.virtualEnvironment;
const deleted = await venv.removeVenvDirectory();
if (!deleted) return false;

const created = await venv.createVenv(sendLogIpc);
if (!created) return false;

return await venv.upgradePip({ onStdout: sendLogIpc, onStderr: sendLogIpc });
});

// Replace the reinstall IPC handler.
ipcMain.removeHandler(IPC_CHANNELS.REINSTALL);
ipcMain.handle(IPC_CHANNELS.REINSTALL, async () => {
log.info('Reinstalling...');
await InstallationManager.reinstall(installation);
});
}

/**
Expand Down Expand Up @@ -116,7 +183,7 @@ export class InstallationManager {
useDesktopConfig().set('migrateCustomNodesFrom', installWizard.migrationSource);
}

const installation = new ComfyInstallation('installed', installWizard.basePath, device);
const installation = new ComfyInstallation('installed', installWizard.basePath, this.telemetry, device);
installation.setState('installed');
return installation;
}
Expand All @@ -135,40 +202,35 @@ export class InstallationManager {
return filePaths[0];
}

/** Notify user that the provided base apth is not valid. */
async #showInvalidBasePathMessage() {
await this.appWindow.showMessageBox({
title: 'Invalid base path',
message:
'ComfyUI needs a valid directory set as its base path. Inside, models, custom nodes, etc will be stored.\n\nClick OK, then selected a new base path.',
type: 'error',
});
}

/**
* Resolves any issues found during installation validation.
* @param installation The installation to resolve issues for
* @throws If the base path is invalid or cannot be saved
*/
async resolveIssues(installation: ComfyInstallation) {
const issues = [...installation.issues];
for (const issue of issues) {
switch (issue) {
// TODO: Other issues (uv mising, venv etc)
case 'invalidBasePath': {
// TODO: Add IPC listeners and proper UI for this
await this.#showInvalidBasePathMessage();

const path = await this.showBasePathPicker();
if (!path) return;

const success = await installation.updateBasePath(path);
if (!success) throw new Error('No base path selected or failed to save in config.');

installation.issues.delete('invalidBasePath');
break;
}
}
}
log.verbose('Resolving issues - awaiting user response:', installation.validation);

// Await user close window request, validate if any errors remain
const isValid = await new Promise<boolean>((resolve) => {
ipcMain.handleOnce(IPC_CHANNELS.COMPLETE_VALIDATION, async (): Promise<boolean> => {
log.verbose('Attempting to close validation window');
// Check if issues have been resolved externally
if (!installation.isValid) await installation.validate();

// Resolve main thread & renderer
const { isValid } = installation;
resolve(isValid);
return isValid;
});
});

log.verbose('Resolution complete:', installation.validation);
return isValid;
}

static async reinstall(installation: ComfyInstallation): Promise<void> {
await installation.uninstall();
app.relaunch();
app.quit();
}
}
17 changes: 5 additions & 12 deletions src/main-process/comfyDesktopApp.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@ import * as Sentry from '@sentry/electron/main';
import todesktop from '@todesktop/runtime';
import { Notification, type TitleBarOverlayOptions, app, dialog, ipcMain } from 'electron';
import log from 'electron-log/main';
import { rm } from 'node:fs/promises';
import path from 'node:path';
import { graphics } from 'systeminformation';

import { ComfyServerConfig } from '../config/comfyServerConfig';
import { ComfySettings } from '../config/comfySettings';
import { IPC_CHANNELS, ProgressStatus, ServerArgs } from '../constants';
import { InstallationManager } from '../install/installationManager';
import { DownloadManager } from '../models/DownloadManager';
import { type ElectronContextMenuOptions } from '../preload';
import { CmCli } from '../services/cmCli';
Expand Down Expand Up @@ -122,9 +122,12 @@ export class ComfyDesktopApp implements HasTelemetry {
ipcMain.handle(IPC_CHANNELS.IS_FIRST_TIME_SETUP, () => {
return !ComfyServerConfig.exists();
});

// Replace the reinstall IPC handler.
ipcMain.removeHandler(IPC_CHANNELS.REINSTALL);
ipcMain.handle(IPC_CHANNELS.REINSTALL, async () => {
log.info('Reinstalling...');
await this.reinstall();
await InstallationManager.reinstall(this.installation);
});
// Restart core
ipcMain.handle(IPC_CHANNELS.RESTART_CORE, async (): Promise<boolean> => {
Expand Down Expand Up @@ -203,16 +206,6 @@ export class ComfyDesktopApp implements HasTelemetry {
}
}

async uninstall(): Promise<void> {
await rm(ComfyServerConfig.configPath);
await useDesktopConfig().permanentlyDeleteConfigFile();
}

async reinstall(): Promise<void> {
await this.uninstall();
this.restart();
}

restart({ customMessage, delay }: { customMessage?: string; delay?: number } = {}): void {
function relaunchApplication(delay?: number) {
if (delay) {
Expand Down
Loading
Loading