From 9d41a9d53308763dbe39bc8acda7e42caddeb353 Mon Sep 17 00:00:00 2001
From: "Lt.Dr.Data" <lt.dr.data@gmail.com>
Date: Tue, 13 Jun 2023 17:41:13 +0900
Subject: [PATCH] support on prompt event handler

---
 server.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/server.py b/server.py
index 174d38af1..8babdc278 100644
--- a/server.py
+++ b/server.py
@@ -1,6 +1,8 @@
 import os
 import sys
 import asyncio
+import traceback
+
 import nodes
 import folder_paths
 import execution
@@ -79,6 +81,8 @@ class PromptServer():
         self.last_node_id = None
         self.client_id = None
 
+        self.on_prompt_handlers = []
+
         @routes.get('/ws')
         async def websocket_handler(request):
             ws = web.WebSocketResponse()
@@ -386,6 +390,7 @@ class PromptServer():
             resp_code = 200
             out_string = ""
             json_data =  await request.json()
+            json_data = self.trigger_on_prompt(json_data)
 
             if "number" in json_data:
                 number = float(json_data['number'])
@@ -521,3 +526,15 @@ class PromptServer():
         if call_on_start is not None:
             call_on_start(address, port)
 
+    def add_on_prompt_handler(self, handler):
+        self.on_prompt_handlers.append(handler)
+
+    def trigger_on_prompt(self, json_data):
+        for handler in self.on_prompt_handlers:
+            try:
+                json_data = handler(json_data)
+            except Exception as e:
+                print(f"[ERROR] An error occurred during the on_prompt_handler processing")
+                traceback.print_exc()
+
+        return json_data