Merge branch 'main' into enh/Ask-user--Capture-error-in-Sentry-when-a…

…pp-crashes

.env_example

@@ -0,0 +1,8 @@
+# The host to use for the ComfyUI server.
+COMFY_HOST=127.0.0.1
+
+# The port to use for the ComfyUI server.
+COMFY_PORT=8188
+
+# Whether to use an external server instead of starting one locally.
+USE_EXTERNAL_SERVER=false

.env_test_example

@@ -0,0 +1,10 @@
+# Use this file to test with local ComfyUI_frontend dev server.
+
+# The host to use for the ComfyUI server.
+COMFY_HOST=localhost
+
+# The port to use for the ComfyUI server.
+COMFY_PORT=5173
+
+# Whether to use an external server instead of starting one locally.
+USE_EXTERNAL_SERVER=true

package.json

@@ -54,7 +54,6 @@
     "@types/react": "^18.3.4",
     "@types/react-dom": "^18.3.0",
     "@types/tar": "6.1.13",
-    "@types/ws": "^8.5.12",
     "@typescript-eslint/eslint-plugin": "^5.0.0",
     "@typescript-eslint/parser": "^5.0.0",
     "electron": "31.3.1",
@@ -90,8 +89,6 @@
     "react-dom": "^18.3.1",
     "systeminformation": "^5.23.5",
     "tar": "^7.4.3",
-    "update-electron-app": "^3.0.0",
-    "ws": "^8.18.0",
     "yaml": "^2.6.0"
   },
   "lint-staged": {

src/__tests__/unit/main.test.ts

@@ -81,14 +81,6 @@ jest.mock('electron-log/main', () => ({
   // Add other methods you might use from electron-log
 }));
 
-// Mock the update-electron-app module
-jest.mock('update-electron-app', () => ({
-  updateElectronApp: jest.fn(),
-  UpdateSourceType: {
-    StaticStorage: 'StaticStorage',
-  },
-}));
-
 describe('createWindow', () => {
   // it('should create a new BrowserWindow with correct options', async () => {
   //   const window = await createWindow('/');

src/constants.ts

@@ -10,8 +10,6 @@ export const IPC_CHANNELS = {
   FIRST_TIME_SETUP_COMPLETE: 'first-time-setup-complete',
   DEFAULT_INSTALL_LOCATION: 'default-install-location',
   GET_PRELOAD_SCRIPT: 'get-preload-script',
-  OPEN_DEVTOOLS: 'open-devtools',
-  OPEN_LOGS_FOLDER: 'open-logs-folder',
   DOWNLOAD_PROGRESS: 'download-progress',
   START_DOWNLOAD: 'start-download',
   PAUSE_DOWNLOAD: 'pause-download',
@@ -20,6 +18,12 @@ export const IPC_CHANNELS = {
   DELETE_MODEL: 'delete-model',
   GET_ALL_DOWNLOADS: 'get-all-downloads',
   SET_SEND_CRASH_REPORTS: 'set-send-crash-reports',
+  GET_ELECTRON_VERSION: 'get-electron-version',
+  SEND_ERROR_TO_SENTRY: 'send-error-to-sentry',
+  GET_BASE_PATH: 'get-base-path',
+  GET_MODEL_CONFIG_PATH: 'get-model-config-path',
+  OPEN_PATH: 'open-path',
+  OPEN_DEV_TOOLS: 'open-dev-tools',
 } as const;
 
 export const COMFY_ERROR_MESSAGE =

src/main.ts

@@ -17,24 +17,35 @@ import Store from 'electron-store';
 import * as net from 'net';
 import { graphics } from 'systeminformation';
 import { createModelConfigFiles, readBasePathFromConfig } from './config/extra_model_config';
-import { WebSocketServer } from 'ws';
 import { StoreType } from './store';
 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';
+
+dotenv.config();
 
 let comfyServerProcess: ChildProcess | null = null;
 let isRestarting: boolean = false; // Prevents double restarts TODO(robinhuang): Remove this once we have a better way to handle restarts. https://github.com/Comfy-Org/electron/issues/149
-const host = '127.0.0.1';
-let port = 8188;
+
+/** The host to use for the ComfyUI server. */
+const host = process.env.COMFY_HOST || '127.0.0.1';
+/** The port to use for the ComfyUI server. */
+let port = parseInt(process.env.COMFY_PORT || '-1');
+/**
+ * Whether to use an external server instead of starting one locally.
+ * Only effective if COMFY_PORT is set.
+ * Note: currently used for testing only.
+ */
+const useExternalServer = process.env.USE_EXTERNAL_SERVER === 'true';
+
 let mainWindow: BrowserWindow | null = null;
-let wss: WebSocketServer | null;
 let store: Store<StoreType> | null = null;
 const messageQueue: Array<any> = []; // Stores mesaages before renderer is ready.
 let downloadManager: DownloadManager;
-
+Sentry.captureMessage('Hello, world!');
 log.initialize();
 
 const comfySettings = new ComfySettings(app.getPath('documents'));
@@ -93,18 +104,15 @@ if (!gotTheLock) {
     }
   });
 
-  app.isPackaged &&
-    Sentry.init({
-      dsn: SENTRY_URL_ENDPOINT,
-      autoSessionTracking: false,
-      integrations: [
-        Sentry.childProcessIntegration({
-          breadcrumbs: ['abnormal-exit', 'killed', 'crashed', 'launch-failed', 'oom', 'integrity-failure'],
-          events: ['abnormal-exit', 'killed', 'crashed', 'launch-failed', 'oom', 'integrity-failure'],
-        }),
-      ],
-      beforeSend: async (event) => {
-        let sendCrashReport = comfySettings.sendCrashStatistics;
+  Sentry.init({
+    dsn: SENTRY_URL_ENDPOINT,
+    autoSessionTracking: false,
+    async beforeSend(event, hint) {
+      if (event.extra?.comfyUIExecutionError) {
+        return event;
+      }
+
+      let sendCrashReport = comfySettings.sendCrashStatistics;
 
         if (!sendCrashReport) {
           const { response } = await dialog.showMessageBox({
@@ -117,8 +125,14 @@ if (!gotTheLock) {
         }
 
         return sendCrashReport ? event : null;
-      },
-    });
+    },
+    integrations: [
+      Sentry.childProcessIntegration({
+        breadcrumbs: ['abnormal-exit', 'killed', 'crashed', 'launch-failed', 'oom', 'integrity-failure'],
+        events: ['abnormal-exit', 'killed', 'crashed', 'launch-failed', 'oom', 'integrity-failure'],
+      }),
+    ],
+  });
 
   graphics()
     .then((graphicsInfo) => {
@@ -184,8 +198,17 @@ if (!gotTheLock) {
           ...options,
         });
       });
-      ipcMain.on(IPC_CHANNELS.OPEN_LOGS_FOLDER, () => {
-        shell.openPath(app.getPath('logs'));
+      ipcMain.handle(IPC_CHANNELS.GET_BASE_PATH, () => {
+        return basePath;
+      });
+      ipcMain.handle(IPC_CHANNELS.GET_MODEL_CONFIG_PATH, () => {
+        return modelConfigPath;
+      });
+      ipcMain.on(IPC_CHANNELS.OPEN_PATH, (event, folderPath: string) => {
+        shell.openPath(folderPath);
+      });
+      ipcMain.on(IPC_CHANNELS.OPEN_DEV_TOOLS, () => {
+        mainWindow?.webContents.openDevTools();
       });
       ipcMain.on(IPC_CHANNELS.SET_SEND_CRASH_REPORTS, (_event, value) => {
         comfySettings.sendCrashStatistics = value;
@@ -200,28 +223,37 @@ if (!gotTheLock) {
         return;
       }
       downloadManager = DownloadManager.getInstance(mainWindow!, getModelsDirectory(basePath));
-      port = await findAvailablePort(8000, 9999).catch((err) => {
-        log.error(`ERROR: Failed to find available port: ${err}`);
-        throw err;
-      });
-
-      sendProgressUpdate('Setting up Python Environment...');
-      const pythonEnvironment = new PythonEnvironment(pythonInstallPath, appResourcesPath, spawnPythonAsync);
-      await pythonEnvironment.setup();
-
-      SetupTray(
-        mainWindow,
-        basePath,
-        modelConfigPath,
-        () => {
-          log.info('Resetting install location');
-          fs.rmSync(modelConfigPath);
-          restartApp();
-        },
-        pythonEnvironment
-      );
-      sendProgressUpdate('Starting Comfy Server...');
-      await launchPythonServer(pythonEnvironment.pythonInterpreterPath, appResourcesPath, modelConfigPath, basePath);
+      downloadManager.registerIpcHandlers();
+
+      port =
+        port !== -1
+          ? port
+          : await findAvailablePort(8000, 9999).catch((err) => {
+              log.error(`ERROR: Failed to find available port: ${err}`);
+              throw err;
+            });
+
+      if (!useExternalServer) {
+        sendProgressUpdate('Setting up Python Environment...');
+        const pythonEnvironment = new PythonEnvironment(pythonInstallPath, appResourcesPath, spawnPythonAsync);
+        await pythonEnvironment.setup();
+
+        // TODO: Make tray setup more flexible here as not all actions depend on the python environment.
+        SetupTray(
+          mainWindow,
+          () => {
+            log.info('Resetting install location');
+            fs.rmSync(modelConfigPath);
+            restartApp();
+          },
+          pythonEnvironment
+        );
+        sendProgressUpdate('Starting Comfy Server...');
+        await launchPythonServer(pythonEnvironment.pythonInterpreterPath, appResourcesPath, modelConfigPath, basePath);
+      } else {
+        sendProgressUpdate('Using external server at ' + host + ':' + port);
+        loadComfyIntoMainWindow();
+      }
     } catch (error) {
       log.error(error);
       sendProgressUpdate(COMFY_ERROR_MESSAGE);
@@ -238,6 +270,22 @@ if (!gotTheLock) {
         }
       }
     );
+
+    ipcMain.handle(IPC_CHANNELS.GET_ELECTRON_VERSION, () => {
+      return app.getVersion();
+    });
+
+    ipcMain.handle(IPC_CHANNELS.SEND_ERROR_TO_SENTRY, async (_event, { error, extras }): Promise<string | null> => {
+      try {
+        return Sentry.captureMessage(error, {
+          level: 'error',
+          extra: { ...extras, comfyUIExecutionError: true },
+        });
+      } catch (err) {
+        log.error('Failed to send error to Sentry:', err);
+        return null;
+      }
+    });
   });
 }
 

src/models/DownloadManager.ts

@@ -258,12 +258,11 @@ export class DownloadManager {
   public static getInstance(mainWindow: BrowserWindow, modelsDirectory: string): DownloadManager {
     if (!DownloadManager.instance) {
       DownloadManager.instance = new DownloadManager(mainWindow, modelsDirectory);
-      DownloadManager.instance.registerIpcHandlers();
     }
     return DownloadManager.instance;
   }
 
-  private registerIpcHandlers() {
+  public registerIpcHandlers() {
     ipcMain.handle(IPC_CHANNELS.START_DOWNLOAD, (event, { url, path, filename }) =>
       this.startDownload(url, path, filename)
     );