Add optional webhook after job finished

main.py

@@ -7,7 +7,8 @@
 import time
 from comfy.cli_args import args
 from app.logger import setup_logger
-
+import aiohttp
+import asyncio
 
 setup_logger(log_level=args.verbose)
 
@@ -104,6 +105,18 @@ def cuda_malloc_warning():
  if cuda_malloc_warning:
  logging.warning("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n")
 
+async def send_webhook(server, prompt_id, data):
+ webhook_url = server.webhooks.pop(prompt_id, None)
+ if webhook_url:
+ try:
+ async with aiohttp.ClientSession() as session:
+ logging.info(f"Sending webhook for prompt {prompt_id}")
+ async with session.post(webhook_url, json=data) as response:
+ if response.status != 200:
+ logging.warning(f"Webhook delivery failed for prompt {prompt_id}. Status: {response.status}")
+ except Exception as e:
+ logging.error(f"Error sending webhook for prompt {prompt_id}: {str(e)}")
+
 def prompt_worker(q, server):
  e = execution.PromptExecutor(server, lru_size=args.cache_lru)
  last_gc_collect = 0
@@ -137,6 +150,15 @@ def prompt_worker(q, server):
  execution_time = current_time - execution_start_time
  logging.info("Prompt executed in {:.2f} seconds".format(execution_time))
 
+ # Send webhook after execution is complete
+ webhook_data = {
+ "prompt_id": prompt_id,
+ "execution_time": execution_time,
+ "status": "success" if e.success else "error",
+ "result": e.history_result
+ }
+ asyncio.run_coroutine_threadsafe(send_webhook(server, prompt_id, webhook_data), server.loop)
+
  flags = q.get_flags()
  free_memory = flags.get("free_memory", False)
 

requirements.txt

@@ -14,7 +14,11 @@ scipy
 tqdm
 psutil
 
+#webhook handling
+aiohttp
+
 #non essential dependencies:
 kornia>=0.7.1
 spandrel
 soundfile
+

server.py

@@ -32,6 +32,10 @@
 from model_filemanager import download_model, DownloadModelStatus
 from typing import Optional
 from api_server.routes.internal.internal_routes import InternalRoutes
+import requests
+import json
+from urllib.parse import urlparse, parse_qs, urlencode
+
 
 class BinaryEventTypes:
  PREVIEW_IMAGE = 1
@@ -159,6 +163,7 @@ def __init__(self, loop):
  self.messages = asyncio.Queue()
  self.client_session:Optional[aiohttp.ClientSession] = None
  self.number = 0
+ self.webhooks = {}
 
  middlewares = [cache_control]
  if args.enable_cors_header:
@@ -623,8 +628,15 @@ async def post_prompt(request):
  if "client_id" in json_data:
  extra_data["client_id"] = json_data["client_id"]
  if valid[0]:
- prompt_id = str(uuid.uuid4())
+ # allow to accept prompt_id from api caller to reference it in webhook handler if needed
+ prompt_id = json_data.get("prompt_id", str(uuid.uuid4()))
  outputs_to_execute = valid[2]
+
+ # Add webhook URL to the webhooks dict if provided
+ webhook_url = json_data.get("webhook_url")
+ if webhook_url:
+ self.webhooks[prompt_id] = webhook_url
+
  self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute))
  response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]}
  return web.json_response(response)
@@ -864,4 +876,4 @@ def trigger_on_prompt(self, json_data):
  logging.warning(f"[ERROR] An error occurred during the on_prompt_handler processing")
  logging.warning(traceback.format_exc())
 
- return json_data
+ return json_data