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);