diff --git a/server.py b/server.py index 87abcc67..538f9075 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 @@ -88,6 +90,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() @@ -438,6 +442,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']) @@ -606,3 +611,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