Draft: Refactor installation state.

src/config/comfyConfigManager.ts

@@ -4,6 +4,9 @@ import log from 'electron-log/main';
 
 export type DirectoryStructure = (string | DirectoryStructure)[];
 
+/**
+ * Responsible for creating the ComfyUI directory structure.
+ */
 export class ComfyConfigManager {
   private static readonly DEFAULT_DIRECTORIES: DirectoryStructure = [
     'custom_nodes',

src/constants.ts

@@ -9,6 +9,10 @@ export const IPC_CHANNELS = {
   OPEN_DIALOG: 'open-dialog',
   FIRST_TIME_SETUP_COMPLETE: 'first-time-setup-complete',
   DEFAULT_INSTALL_LOCATION: 'default-install-location',
+  // Installation state
+  INSTALLATION_STATE_CHANGED: 'installation-state-changed',
+  GET_INSTALLATION_STATE: 'get-installation-state',
+  // End Installation state
   DOWNLOAD_PROGRESS: 'download-progress',
   START_DOWNLOAD: 'start-download',
   PAUSE_DOWNLOAD: 'pause-download',

src/install/install.ts

@@ -0,0 +1,124 @@
+import { app, ipcMain } from 'electron';
+import log from 'electron-log/main';
+import { IPC_CHANNELS } from '../constants';
+import { createModelConfigFiles, getModelConfigPath } from '../config/extra_model_config';
+import fs from 'fs';
+import { AppWindow } from '../main-process/appWindow';
+import { ComfyConfigManager } from '../config/comfyConfigManager';
+
+type InstallationState =
+  | { status: 'NOT_INSTALLED' }
+  | { status: 'ASKING_FOR_DIRECTORY'; defaultLocation: string; validationErrorMessage?: string }
+  | {
+      status: 'INSTALLING';
+      location: string;
+    }
+  | {
+      status: 'READY';
+      location: string; // Must exist!
+    }
+  | {
+      status: 'ERROR';
+      error: Error; // Must exist!
+    };
+
+/**
+ * This class is responsible for handling the installation of ComfyUI.
+ * It registers an IPC handler for the current installation state.
+ * It also sends any changes to the state to the renderer process.
+ */
+export class ComfyUIInstall {
+  private static instance: ComfyUIInstall;
+  private installLocation: string;
+  private defaultInstallLocation: string;
+  private state: InstallationState = { status: 'NOT_INSTALLED' };
+  private appWindow: AppWindow;
+
+  private constructor(appWindow: AppWindow) {
+    const defaultInstallLocation = app.getPath('documents');
+    ipcMain.handle(IPC_CHANNELS.DEFAULT_INSTALL_LOCATION, () => defaultInstallLocation); // TODO: Remove after migration.
+    ipcMain.handle(IPC_CHANNELS.GET_INSTALLATION_STATE, () => this.state);
+    this.installLocation = defaultInstallLocation;
+    this.defaultInstallLocation = defaultInstallLocation;
+    this.appWindow = appWindow;
+  }
+
+  static get(appWindow: AppWindow): ComfyUIInstall {
+    if (!ComfyUIInstall.instance) {
+      ComfyUIInstall.instance = new ComfyUIInstall(appWindow);
+    }
+    return ComfyUIInstall.instance;
+  }
+
+  setState(state: InstallationState): void {
+    log.info('Setting installation state:', state);
+    this.state = state;
+    this.appWindow.send(IPC_CHANNELS.INSTALLATION_STATE_CHANGED, state);
+  }
+
+  public async install(): Promise<void> {
+    const firstTimeSetup = await this.isFirstTimeSetup();
+    log.info('First time setup:', firstTimeSetup);
+    if (!firstTimeSetup) {
+      this.appWindow.send(IPC_CHANNELS.FIRST_TIME_SETUP_COMPLETE, null); // TODO:Remove after migration.
+      return;
+    }
+    this.setState({ status: 'ASKING_FOR_DIRECTORY', defaultLocation: this.defaultInstallLocation });
+    this.appWindow.send(IPC_CHANNELS.SHOW_SELECT_DIRECTORY, null); // TODO:Remove after migration.
+    while (this.state.status === 'ASKING_FOR_DIRECTORY') {
+      const selectedDirectory = await this.selectInstallDirectory();
+      const { valid, errorMessage } = await this.isValidComfyDirectory(selectedDirectory);
+      if (!valid) {
+        this.setState({ ...this.state, validationErrorMessage: errorMessage });
+        return;
+      } else {
+        this.setState({ status: 'INSTALLING', location: selectedDirectory });
+        this.installLocation = selectedDirectory;
+      }
+    }
+
+    const actualComfyDirectory = ComfyConfigManager.setUpComfyUI(this.installLocation);
+    const modelConfigPath = getModelConfigPath();
+    await createModelConfigFiles(modelConfigPath, actualComfyDirectory);
+
+    this.setState({ status: 'READY', location: this.installLocation });
+  }
+
+  /**
+   * Check if the user has completed the first time setup wizard.
+   * This means the extra_models_config.yaml file exists in the user's data directory.
+   */
+  private async isFirstTimeSetup(): Promise<boolean> {
+    const extraModelsConfigPath = getModelConfigPath();
+    return !fs.existsSync(extraModelsConfigPath);
+  }
+
+  private async selectInstallDirectory(): Promise<string> {
+    return new Promise((resolve, reject) => {
+      ipcMain.on(IPC_CHANNELS.SELECTED_DIRECTORY, (_event, value: string) => {
+        log.info('Directory selected:', value);
+        resolve(value);
+      });
+    });
+  }
+
+  public async isValidComfyDirectory(selectedDirectory: string): Promise<{ valid: boolean; errorMessage?: string }> {
+    try {
+      const files = await fs.promises.readdir(selectedDirectory);
+
+      if (files.includes('ComfyUI')) {
+        return {
+          valid: false,
+          errorMessage: 'A ComfyUI installation already exists in this directory. Please choose another location.',
+        };
+      }
+
+      return { valid: true };
+    } catch (error) {
+      return {
+        valid: false,
+        errorMessage: 'Unable to access the selected directory. Please ensure you have proper permissions.',
+      };
+    }
+  }
+}

src/main.ts

@@ -4,24 +4,24 @@ import axios from 'axios';
 import path from 'node:path';
 import { SetupTray } from './tray';
 import { IPC_CHANNELS, SENTRY_URL_ENDPOINT, ProgressStatus } from './constants';
-import { app, BrowserWindow, dialog, ipcMain, shell } from 'electron';
+import { app, dialog, ipcMain } from 'electron';
 import log from 'electron-log/main';
 import * as Sentry from '@sentry/electron/main';
 import * as net from 'net';
 import { graphics } from 'systeminformation';
-import { createModelConfigFiles, getModelConfigPath, readBasePathFromConfig } from './config/extra_model_config';
+import { getModelConfigPath } from './config/extra_model_config';
 import todesktop from '@todesktop/runtime';
 import { PythonEnvironment } from './pythonEnvironment';
 import { DownloadManager } from './models/DownloadManager';
 import { getModelsDirectory } from './utils';
 import { ComfySettings } from './config/comfySettings';
 import dotenv from 'dotenv';
 import { buildMenu } from './menu/menu';
-import { ComfyConfigManager } from './config/comfyConfigManager';
 import { AppWindow } from './main-process/appWindow';
 import { getAppResourcesPath, getBasePath, getPythonInstallPath } from './install/resourcePaths';
 import { PathHandlers } from './handlers/pathHandlers';
 import { AppInfoHandlers } from './handlers/appInfoHandlers';
+import { ComfyUIInstall } from './install/install';
 
 dotenv.config();
 
@@ -157,13 +157,8 @@ if (!gotTheLock) {
         });
       });
 
-      ipcMain.on(IPC_CHANNELS.OPEN_DEV_TOOLS, () => {
-        appWindow.openDevTools();
-      });
-      ipcMain.handle(IPC_CHANNELS.IS_FIRST_TIME_SETUP, () => {
-        return isFirstTimeSetup();
-      });
-      await handleFirstTimeSetup();
+      const install = ComfyUIInstall.get(appWindow);
+      install.install();
       const basePath = await getBasePath();
       const pythonInstallPath = await getPythonInstallPath();
       if (!basePath || !pythonInstallPath) {
@@ -551,38 +546,6 @@ function findAvailablePort(startPort: number, endPort: number): Promise<number>
     tryPort(startPort);
   });
 }
-/**
- * Check if the user has completed the first time setup wizard.
- * This means the extra_models_config.yaml file exists in the user's data directory.
- */
-function isFirstTimeSetup(): boolean {
-  const extraModelsConfigPath = getModelConfigPath();
-  return !fs.existsSync(extraModelsConfigPath);
-}
-
-async function selectedInstallDirectory(): Promise<string> {
-  return new Promise((resolve, reject) => {
-    ipcMain.on(IPC_CHANNELS.SELECTED_DIRECTORY, (_event, value) => {
-      log.info('User selected to install ComfyUI in:', value);
-      resolve(value);
-    });
-  });
-}
-
-async function handleFirstTimeSetup() {
-  const firstTimeSetup = isFirstTimeSetup();
-  log.info('First time setup:', firstTimeSetup);
-  if (firstTimeSetup) {
-    appWindow.send(IPC_CHANNELS.SHOW_SELECT_DIRECTORY, null);
-    const selectedDirectory = await selectedInstallDirectory();
-    const actualComfyDirectory = ComfyConfigManager.setUpComfyUI(selectedDirectory);
-
-    const modelConfigPath = getModelConfigPath();
-    await createModelConfigFiles(modelConfigPath, actualComfyDirectory);
-  } else {
-    appWindow.send(IPC_CHANNELS.FIRST_TIME_SETUP_COMPLETE, null);
-  }
-}
 
 /**
  * Rotate old log files by adding a timestamp to the end of the file.