From 2eaa664089cb036fc6123b5276714d427aa739ce Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Thu, 2 Mar 2023 21:33:50 +0000
Subject: [PATCH] Refactored sockets

---
 web/scripts/api.js | 81 +++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 80 insertions(+), 1 deletion(-)

diff --git a/web/scripts/api.js b/web/scripts/api.js
index 896c15cad..322cd95f6 100644
--- a/web/scripts/api.js
+++ b/web/scripts/api.js
@@ -1,4 +1,83 @@
-class ComfyApi {
+class ComfyApi extends EventTarget {
+	constructor() {
+      super();
+	}
+
+	#pollQueue() {
+		setInterval(async () => {
+			try {
+				const resp = await fetch("/prompt");
+				const status = await resp.json();
+				this.dispatchEvent(new CustomEvent("status", { detail: status }));
+			} catch (error) {
+				this.dispatchEvent(new CustomEvent("status", { detail: null }));
+			}
+		}, 1000);
+	}
+
+	#createSocket(isReconnect) {
+		if (this.socket) {
+			return;
+		}
+
+		let opened = false;
+		this.socket = new WebSocket(`ws${window.location.protocol === "https:" ? "s" : ""}://${location.host}/ws`);
+
+		this.socket.addEventListener("open", () => {
+			opened = true;
+			if (isReconnect) {
+				this.dispatchEvent(new CustomEvent("reconnected"));
+			}
+		});
+
+		this.socket.addEventListener("error", () => {
+			if (this.socket) this.socket.close();
+			this.#pollQueue();
+		});
+
+		this.socket.addEventListener("close", () => {
+			setTimeout(() => {
+				this.socket = null;
+				this.#createSocket(true);
+			}, 300);
+			if (opened) {
+				this.dispatchEvent(new CustomEvent("status", { detail: null }));
+				this.dispatchEvent(new CustomEvent("reconnecting"));
+			}
+		});
+
+		this.socket.addEventListener("message", (event) => {
+			try {
+				const msg = JSON.parse(event.data);
+				switch (msg.type) {
+					case "status":
+						if (msg.data.sid) {
+							this.clientId = msg.data.sid;
+						}
+						this.dispatchEvent(new CustomEvent("status", { detail: msg.data.status }));
+						break;
+					case "progress":
+						this.dispatchEvent(new CustomEvent("progress", { detail: msg.data }));
+						break;
+					case "executing":
+						this.dispatchEvent(new CustomEvent("executing", { detail: msg.data.node }));
+						break;
+					case "executed":
+						this.dispatchEvent(new CustomEvent("executed", { detail: msg.data }));
+						break;
+					default:
+						throw new Error("Unknown message type");
+				}
+			} catch (error) {
+				console.warn("Unhandled message:", event.data);
+			}
+		});
+	}
+
+   init() {
+      this.#createSocket();
+   }
+
 	async getNodeDefs() {
 		const resp = await fetch("object_info", { cache: "no-store" });
 		return await resp.json();