From 8d0a142321f3a59303af04aa3c8ee1e9be4c19e1 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Fri, 24 Mar 2023 11:39:09 +0000
Subject: [PATCH] adds simple access to server from custom nodes

---
 server.py          |  2 ++
 web/scripts/api.js | 13 ++++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/server.py b/server.py
index 73429acc..2593e808 100644
--- a/server.py
+++ b/server.py
@@ -29,6 +29,8 @@ async def cache_control(request: web.Request, handler):
 
 class PromptServer():
     def __init__(self, loop):
+        PromptServer.instance = self
+        
         mimetypes.init(); 
         mimetypes.types_map['.js'] = 'application/javascript; charset=utf-8'
         self.prompt_queue = None
diff --git a/web/scripts/api.js b/web/scripts/api.js
index b90b1c65..2b90c2ab 100644
--- a/web/scripts/api.js
+++ b/web/scripts/api.js
@@ -1,8 +1,15 @@
 class ComfyApi extends EventTarget {
+	#registered = new Set();
+
 	constructor() {
 		super();
 	}
 
+	addEventListener(type, callback, options) {
+		super.addEventListener(type, callback, options);
+		this.#registered.add(type);
+	}
+
 	/**
 	 * Poll status  for colab and other things that don't support websockets.
 	 */
@@ -82,7 +89,11 @@ class ComfyApi extends EventTarget {
 						this.dispatchEvent(new CustomEvent("executed", { detail: msg.data }));
 						break;
 					default:
-						throw new Error("Unknown message type");
+						if (this.#registered.has(msg.type)) {
+							this.dispatchEvent(new CustomEvent(msg.type, { detail: msg.data }));
+						} else {
+							throw new Error("Unknown message type");
+						}
 				}
 			} catch (error) {
 				console.warn("Unhandled message:", event.data);