From 3225d7cb6f7c2153331b0c0c1c0578809d0e3cc9 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Tue, 18 Jun 2024 21:07:45 -0400
Subject: [PATCH] Make dev server ready for use (#43)

* basic dev server

* nit

* Return empty extensions list for dev server

* Update readme
---
 README.md               |  2 --
 src/scripts/api.ts      |  6 +++++-
 src/scripts/app.ts      |  2 +-
 tests-ui/utils/setup.ts |  5 +++--
 vite.config.mts         | 23 ++++++++++++++++++-----
 5 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 407ebad8fa..c736b67a18 100644
--- a/README.md
+++ b/README.md
@@ -18,8 +18,6 @@ in ComfyUI's front-end code.
 
 ## Development
 
-Currently the dev server does not work as the ws runs on root path '/', and all api endpoints are all defined on '/'. There might need to be some API changes before dev server can work.
-
 - Run `npm install` to install the necessary packages
 - Start local ComfyUI backend at `localhost:8188`
 - Run `npm run dev` to start the dev server
diff --git a/src/scripts/api.ts b/src/scripts/api.ts
index 2dece13a48..6b5e4b972d 100644
--- a/src/scripts/api.ts
+++ b/src/scripts/api.ts
@@ -31,7 +31,11 @@ class ComfyApi extends EventTarget {
 		this.initialClientId = sessionStorage.getItem("clientId");
 	}
 
-	apiURL(route) {
+	apiURL(route: string): string {
+		return this.api_base + "/api" + route;
+	}
+
+	fileURL(route: string): string {
 		return this.api_base + route;
 	}
 
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 1cc4210677..ecfda6145d 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -1510,7 +1510,7 @@ export class ComfyApp {
 			.filter(extension => !extension.includes("extensions/core"))
 			.map(async ext => {
 				try {
-					await import(/* @vite-ignore */api.apiURL(ext));
+					await import(/* @vite-ignore */api.fileURL(ext));
 				} catch (error) {
 					console.error("Error loading extension", ext, error);
 				}
diff --git a/tests-ui/utils/setup.ts b/tests-ui/utils/setup.ts
index c830a459bd..ccb39839d6 100644
--- a/tests-ui/utils/setup.ts
+++ b/tests-ui/utils/setup.ts
@@ -26,8 +26,8 @@ export interface APIConfig {
  */
 
 /**
- * @param {{ 
- * 	mockExtensions?: string[], 
+ * @param {{
+ * 	mockExtensions?: string[],
  * 	mockNodeDefs?: Record<string, ComfyObjectInfo>,
 * 	settings?: Record<string, string>
 * 	userConfig?: {storage: "server" | "browser", users?: Record<string, any>, migrated?: boolean },
@@ -59,6 +59,7 @@ export function mockApi(config: APIConfig = {}) {
 		getNodeDefs: jest.fn(() => mockNodeDefs),
 		init: jest.fn(),
 		apiURL: jest.fn((x) => "src/" + x),
+		fileURL: jest.fn((x) => "src/" + x),
 		createUser: jest.fn((username) => {
 			// @ts-ignore
 			if(username in userConfig.users) {
diff --git a/vite.config.mts b/vite.config.mts
index ede9783755..438b5e1c32 100644
--- a/vite.config.mts
+++ b/vite.config.mts
@@ -2,6 +2,7 @@ import { defineConfig, Plugin } from 'vite';
 import { viteStaticCopy } from 'vite-plugin-static-copy'
 import path from 'path';
 
+const IS_DEV = process.env.NODE_ENV === 'development';
 
 interface ShimResult {
 	code: string;
@@ -12,6 +13,9 @@ function comfyAPIPlugin(): Plugin {
 	return {
 		name: 'comfy-api-plugin',
 		transform(code: string, id: string) {
+			if (IS_DEV)
+				return null;
+
 			// TODO: Remove second condition after all js files are converted to ts
 			if (id.endsWith('.ts') || (id.endsWith('.js') && id.includes("extensions/core"))) {
 				const result = transformExports(code, id);
@@ -74,20 +78,29 @@ function getModuleName(id: string): string {
 
 export default defineConfig({
 	server: {
-		open: true,
 		proxy: {
-			// Proxy websocket requests to the server
-			'/': {
+			'/api': {
+				target: 'http://127.0.0.1:8188',
+				// Return empty array for extensions API as these modules
+				// are not on vite's dev server.
+				bypass: (req, res, options) => {
+					if (req.url === '/api/extensions') {
+						res.end(JSON.stringify([]));
+					}
+					return null;
+				},
+			},
+			'/ws': {
 				target: 'ws://127.0.0.1:8188',
 				ws: true,
-			}
+			},
 		}
 	},
 	plugins: [
 		comfyAPIPlugin(),
 		viteStaticCopy({
 			targets: [
-				{src: "src/lib/*", dest: "lib/"},
+				{ src: "src/lib/*", dest: "lib/" },
 			],
 		}),
 	],