From 7cd47564bdf559f8206ed99aea23475df85f5515 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <hcl@comfy.org>
Date: Sun, 17 Nov 2024 16:48:39 -0500
Subject: [PATCH] Replace custom wait logic with wait-on package (#278)

---
 package.json |  2 ++
 src/main.ts  | 71 ++++++++++++--------------------------
 yarn.lock    | 96 +++++++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 119 insertions(+), 50 deletions(-)

diff --git a/package.json b/package.json
index ca7d1240..f301b533 100644
--- a/package.json
+++ b/package.json
@@ -64,6 +64,7 @@
     "@types/jest": "^29.5.13",
     "@types/node": "^22.5.0",
     "@types/tar": "6.1.13",
+    "@types/wait-on": "^5.3.4",
     "@typescript-eslint/eslint-plugin": "^5.0.0",
     "@typescript-eslint/parser": "^5.0.0",
     "electron": "31.3.1",
@@ -98,6 +99,7 @@
     "jest": "^29.7.0",
     "systeminformation": "^5.23.5",
     "tar": "^7.4.3",
+    "wait-on": "^8.0.1",
     "yaml": "^2.6.0"
   },
   "lint-staged": {
diff --git a/src/main.ts b/src/main.ts
index 58f49750..744f9864 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -1,6 +1,5 @@
 import { ChildProcess } from 'node:child_process';
 import fs from 'fs';
-import axios from 'axios';
 import path from 'node:path';
 import { IPC_CHANNELS, SENTRY_URL_ENDPOINT, ProgressStatus } from './constants';
 import { app, dialog, ipcMain } from 'electron';
@@ -20,6 +19,7 @@ import { PathHandlers } from './handlers/pathHandlers';
 import { AppInfoHandlers } from './handlers/appInfoHandlers';
 import { InstallOptions } from './preload';
 import { VirtualEnvironment } from './virtualEnvironment';
+import waitOn from 'wait-on';
 
 dotenv.config();
 
@@ -259,35 +259,20 @@ function restartApp({ customMessage, delay }: { customMessage?: string; delay?:
 }
 
 const isComfyServerReady = async (host: string, port: number): Promise<boolean> => {
-  const url = `http://${host}:${port}/queue`;
-
   try {
-    const response = await axios.get(url, {
+    await waitOn({
+      resources: [`http://${host}:${port}/queue`],
       timeout: 5000, // 5 seconds timeout
+      interval: 1000, // Check every second
     });
-
-    if (response.status >= 200 && response.status < 300) {
-      log.info(`Server responded with status ${response.status} at ${url}`);
-      return true;
-    } else {
-      log.warn(`Server responded with status ${response.status} at ${url}`);
-      return false;
-    }
+    log.info(`Server is ready at http://${host}:${port}/queue`);
+    return true;
   } catch (error) {
-    if (axios.isAxiosError(error)) {
-      log.error(`Failed to connect to server at ${url}: ${error.message}`);
-    } else {
-      log.error(`Unexpected error when checking server at ${url}: ${error}`);
-    }
+    log.error(`Server not ready at http://${host}:${port}/queue: ${error}`);
     return false;
   }
 };
 
-// Launch Python Server Variables
-const maxFailWait: number = 120 * 1000; // 120seconds
-let currentWaitTime = 0;
-let spawnServerTimeout: NodeJS.Timeout | null = null;
-
 const launchPythonServer = async (
   virtualEnvironment: VirtualEnvironment,
   appResourcesPath: string,
@@ -352,35 +337,23 @@ const launchPythonServer = async (
       }
     });
 
-    const checkInterval = 1000; // Check every 1 second
+    try {
+      await waitOn({
+        resources: [`http://${host}:${port}/queue`],
+        timeout: 120000, // 120 seconds timeout
+        interval: 1000, // Check every second
+      });
 
-    const checkServerReady = async (): Promise<void> => {
-      currentWaitTime += 1000;
-      if (currentWaitTime > maxFailWait) {
-        //Something has gone wrong and we need to backout.
-        if (spawnServerTimeout) {
-          clearTimeout(spawnServerTimeout);
-        }
-        reject('Python Server Failed To Start Within Timeout.');
-      }
-      const isReady = await isComfyServerReady(host, port);
-      if (isReady) {
-        sendProgressUpdate(ProgressStatus.READY);
-        log.info('Python server is ready');
-
-        //For now just replace the source of the main window to the python server
-        setTimeout(() => loadComfyIntoMainWindow(), 1000);
-        if (spawnServerTimeout) {
-          clearTimeout(spawnServerTimeout);
-        }
-        return resolve();
-      } else {
-        log.info('Ping failed. Retrying...');
-        spawnServerTimeout = setTimeout(checkServerReady, checkInterval);
-      }
-    };
+      sendProgressUpdate(ProgressStatus.READY);
+      log.info('Python server is ready');
 
-    checkServerReady();
+      //For now just replace the source of the main window to the python server
+      setTimeout(() => loadComfyIntoMainWindow(), 1000);
+      resolve();
+    } catch (error) {
+      log.error('Server failed to start:', error);
+      reject('Python Server Failed To Start Within Timeout.');
+    }
   });
 };
 
diff --git a/yarn.lock b/yarn.lock
index 0f8117b5..bf0c2ae1 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -482,6 +482,7 @@ __metadata:
     "@types/jest": "npm:^29.5.13"
     "@types/node": "npm:^22.5.0"
     "@types/tar": "npm:6.1.13"
+    "@types/wait-on": "npm:^5.3.4"
     "@typescript-eslint/eslint-plugin": "npm:^5.0.0"
     "@typescript-eslint/parser": "npm:^5.0.0"
     adm-zip: "npm:^0.5.15"
@@ -505,6 +506,7 @@ __metadata:
     typescript: "npm:~5.5.4"
     vite: "npm:^5.0.12"
     vite-plugin-dts: "npm:^4.3.0"
+    wait-on: "npm:^8.0.1"
     yaml: "npm:^2.6.0"
   languageName: unknown
   linkType: soft
@@ -1221,6 +1223,22 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@hapi/hoek@npm:^9.0.0, @hapi/hoek@npm:^9.3.0":
+  version: 9.3.0
+  resolution: "@hapi/hoek@npm:9.3.0"
+  checksum: 10c0/a096063805051fb8bba4c947e293c664b05a32b47e13bc654c0dd43813a1cec993bdd8f29ceb838020299e1d0f89f68dc0d62a603c13c9cc8541963f0beca055
+  languageName: node
+  linkType: hard
+
+"@hapi/topo@npm:^5.1.0":
+  version: 5.1.0
+  resolution: "@hapi/topo@npm:5.1.0"
+  dependencies:
+    "@hapi/hoek": "npm:^9.0.0"
+  checksum: 10c0/b16b06d9357947149e032bdf10151eb71aea8057c79c4046bf32393cb89d0d0f7ca501c40c0f7534a5ceca078de0700d2257ac855c15e59fe4e00bba2f25c86f
+  languageName: node
+  linkType: hard
+
 "@humanwhocodes/config-array@npm:^0.13.0":
   version: 0.13.0
   resolution: "@humanwhocodes/config-array@npm:0.13.0"
@@ -2959,6 +2977,29 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@sideway/address@npm:^4.1.5":
+  version: 4.1.5
+  resolution: "@sideway/address@npm:4.1.5"
+  dependencies:
+    "@hapi/hoek": "npm:^9.0.0"
+  checksum: 10c0/638eb6f7e7dba209053dd6c8da74d7cc995e2b791b97644d0303a7dd3119263bcb7225a4f6804d4db2bc4f96e5a9d262975a014f58eae4d1753c27cbc96ef959
+  languageName: node
+  linkType: hard
+
+"@sideway/formula@npm:^3.0.1":
+  version: 3.0.1
+  resolution: "@sideway/formula@npm:3.0.1"
+  checksum: 10c0/3fe81fa9662efc076bf41612b060eb9b02e846ea4bea5bd114f1662b7f1541e9dedcf98aff0d24400bcb92f113964a50e0290b86e284edbdf6346fa9b7e2bf2c
+  languageName: node
+  linkType: hard
+
+"@sideway/pinpoint@npm:^2.0.0":
+  version: 2.0.0
+  resolution: "@sideway/pinpoint@npm:2.0.0"
+  checksum: 10c0/d2ca75dacaf69b8fc0bb8916a204e01def3105ee44d8be16c355e5f58189eb94039e15ce831f3d544f229889ccfa35562a0ce2516179f3a7ee1bbe0b71e55b36
+  languageName: node
+  linkType: hard
+
 "@sinclair/typebox@npm:^0.27.8":
   version: 0.27.8
   resolution: "@sinclair/typebox@npm:0.27.8"
@@ -3455,6 +3496,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@types/wait-on@npm:^5.3.4":
+  version: 5.3.4
+  resolution: "@types/wait-on@npm:5.3.4"
+  dependencies:
+    "@types/node": "npm:*"
+  checksum: 10c0/e366fbfa78fbed4a033aa03072291ba869328dc1a1b715109540af3a328f8f023ec868219f6d2148d5a2ea21f1ce0f12d29e42411f3255ba155da2af978319ee
+  languageName: node
+  linkType: hard
+
 "@types/yargs-parser@npm:*":
   version: 21.0.3
   resolution: "@types/yargs-parser@npm:21.0.3"
@@ -8709,6 +8759,19 @@ __metadata:
   languageName: node
   linkType: hard
 
+"joi@npm:^17.13.3":
+  version: 17.13.3
+  resolution: "joi@npm:17.13.3"
+  dependencies:
+    "@hapi/hoek": "npm:^9.3.0"
+    "@hapi/topo": "npm:^5.1.0"
+    "@sideway/address": "npm:^4.1.5"
+    "@sideway/formula": "npm:^3.0.1"
+    "@sideway/pinpoint": "npm:^2.0.0"
+  checksum: 10c0/9262aef1da3f1bec5b03caf50c46368899fe03b8ff26cbe3d53af4584dd1049079fc97230bbf1500b6149db7cc765b9ee45f0deb24bb6fc3fa06229d7148c17f
+  languageName: node
+  linkType: hard
+
 "join-component@npm:^1.1.0":
   version: 1.1.0
   resolution: "join-component@npm:1.1.0"
@@ -9139,7 +9202,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"lodash@npm:^4.17.12, lodash@npm:^4.17.15, lodash@npm:^4.17.20, lodash@npm:~4.17.15":
+"lodash@npm:^4.17.12, lodash@npm:^4.17.15, lodash@npm:^4.17.20, lodash@npm:^4.17.21, lodash@npm:~4.17.15":
   version: 4.17.21
   resolution: "lodash@npm:4.17.21"
   checksum: 10c0/d8cbea072bb08655bb4c989da418994b073a608dffa608b09ac04b43a791b12aeae7cd7ad919aa4c925f33b48490b5cfe6c1f71d827956071dae2e7bb3a6b74c
@@ -11355,6 +11418,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"rxjs@npm:^7.8.1":
+  version: 7.8.1
+  resolution: "rxjs@npm:7.8.1"
+  dependencies:
+    tslib: "npm:^2.1.0"
+  checksum: 10c0/3c49c1ecd66170b175c9cacf5cef67f8914dcbc7cd0162855538d365c83fea631167cacb644b3ce533b2ea0e9a4d0b12175186985f89d75abe73dbd8f7f06f68
+  languageName: node
+  linkType: hard
+
 "safe-array-concat@npm:^1.1.2":
   version: 1.1.2
   resolution: "safe-array-concat@npm:1.1.2"
@@ -12328,6 +12400,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"tslib@npm:^2.1.0":
+  version: 2.8.1
+  resolution: "tslib@npm:2.8.1"
+  checksum: 10c0/9c4759110a19c53f992d9aae23aac5ced636e99887b51b9e61def52611732872ff7668757d4e4c61f19691e36f4da981cd9485e869b4a7408d689f6bf1f14e62
+  languageName: node
+  linkType: hard
+
 "tsutils@npm:^3.21.0":
   version: 3.21.0
   resolution: "tsutils@npm:3.21.0"
@@ -12781,6 +12860,21 @@ __metadata:
   languageName: node
   linkType: hard
 
+"wait-on@npm:^8.0.1":
+  version: 8.0.1
+  resolution: "wait-on@npm:8.0.1"
+  dependencies:
+    axios: "npm:^1.7.7"
+    joi: "npm:^17.13.3"
+    lodash: "npm:^4.17.21"
+    minimist: "npm:^1.2.8"
+    rxjs: "npm:^7.8.1"
+  bin:
+    wait-on: bin/wait-on
+  checksum: 10c0/e8fe78084465a5fa95b17095136dcb3942480c7d9e47a0cfb72363238e2411883630dc9ced02fb372f4cc7ec0ec2484c92c67060ea2f81d0e97e41ef21df8006
+  languageName: node
+  linkType: hard
+
 "walker@npm:^1.0.8":
   version: 1.0.8
   resolution: "walker@npm:1.0.8"