Include workflow_id in all execution WebSocket messages (CORE-198) (Comfy-Org#13684)

christian-byrne · web-flow · commit 4f6018982dcd · 2026-05-14T15:11:34.000-07:00
diff --git a/comfy_execution/jobs.py b/comfy_execution/jobs.py
@@ -93,15 +93,35 @@ def _create_text_preview(value: str) -> dict:
     }
 
 
+def extract_workflow_id(extra_data: Optional[dict]) -> Optional[str]:
+    """Extract the workflow id from a prompt's ``extra_data``.
+
+    The frontend stores the id at ``extra_data["extra_pnginfo"]["workflow"]["id"]``
+    when a prompt is queued. Any value that is not a non-empty string is treated as
+    missing so callers can rely on the return being either ``None`` or a string.
+    """
+    if not isinstance(extra_data, dict):
+        return None
+    extra_pnginfo = extra_data.get('extra_pnginfo')
+    if not isinstance(extra_pnginfo, dict):
+        return None
+    workflow = extra_pnginfo.get('workflow')
+    if not isinstance(workflow, dict):
+        return None
+    workflow_id = workflow.get('id')
+    if isinstance(workflow_id, str) and workflow_id:
+        return workflow_id
+    return None
+
+
 def _extract_job_metadata(extra_data: dict) -> tuple[Optional[int], Optional[str]]:
     """Extract create_time and workflow_id from extra_data.
 
     Returns:
         tuple: (create_time, workflow_id)
     """
     create_time = extra_data.get('create_time')
-    extra_pnginfo = extra_data.get('extra_pnginfo', {})
-    workflow_id = extra_pnginfo.get('workflow', {}).get('id')
+    workflow_id = extract_workflow_id(extra_data)
     return create_time, workflow_id
 
 
diff --git a/comfy_execution/progress.py b/comfy_execution/progress.py
@@ -164,6 +164,8 @@ def _send_progress_state(self, prompt_id: str, nodes: Dict[str, NodeProgressStat
         if self.server_instance is None:
             return
 
+        workflow_id = self.registry.workflow_id if self.registry else None
+
         # Only send info for non-pending nodes
         active_nodes = {
             node_id: {
@@ -172,6 +174,7 @@ def _send_progress_state(self, prompt_id: str, nodes: Dict[str, NodeProgressStat
                 "state": state["state"].value,
                 "node_id": node_id,
                 "prompt_id": prompt_id,
+                "workflow_id": workflow_id,
                 "display_node_id": self.registry.dynprompt.get_display_node_id(node_id),
                 "parent_node_id": self.registry.dynprompt.get_parent_node_id(node_id),
                 "real_node_id": self.registry.dynprompt.get_real_node_id(node_id),
@@ -183,7 +186,7 @@ def _send_progress_state(self, prompt_id: str, nodes: Dict[str, NodeProgressStat
         # Send a combined progress_state message with all node states
         # Include client_id to ensure message is only sent to the initiating client
         self.server_instance.send_sync(
-            "progress_state", {"prompt_id": prompt_id, "nodes": active_nodes}, self.server_instance.client_id
+            "progress_state", {"prompt_id": prompt_id, "workflow_id": workflow_id, "nodes": active_nodes}, self.server_instance.client_id
         )
 
     @override
@@ -215,6 +218,7 @@ def update_handler(
                 metadata = {
                     "node_id": node_id,
                     "prompt_id": prompt_id,
+                    "workflow_id": self.registry.workflow_id if self.registry else None,
                     "display_node_id": self.registry.dynprompt.get_display_node_id(
                         node_id
                     ),
@@ -240,9 +244,10 @@ class ProgressRegistry:
     Registry that maintains node progress state and notifies registered handlers.
     """
 
-    def __init__(self, prompt_id: str, dynprompt: "DynamicPrompt"):
+    def __init__(self, prompt_id: str, dynprompt: "DynamicPrompt", workflow_id: Optional[str] = None):
         self.prompt_id = prompt_id
         self.dynprompt = dynprompt
+        self.workflow_id = workflow_id
         self.nodes: Dict[str, NodeProgressState] = {}
         self.handlers: Dict[str, ProgressHandler] = {}
 
@@ -322,15 +327,15 @@ def reset_handlers(self) -> None:
 # Global registry instance
 global_progress_registry: ProgressRegistry | None = None
 
-def reset_progress_state(prompt_id: str, dynprompt: "DynamicPrompt") -> None:
+def reset_progress_state(prompt_id: str, dynprompt: "DynamicPrompt", workflow_id: Optional[str] = None) -> None:
     global global_progress_registry
 
     # Reset existing handlers if registry exists
     if global_progress_registry is not None:
         global_progress_registry.reset_handlers()
 
     # Create new registry
-    global_progress_registry = ProgressRegistry(prompt_id, dynprompt)
+    global_progress_registry = ProgressRegistry(prompt_id, dynprompt, workflow_id)
 
 
 def add_progress_handler(handler: ProgressHandler) -> None:
diff --git a/execution.py b/execution.py
@@ -38,6 +38,7 @@
 from comfy_execution.graph_utils import GraphBuilder, is_link
 from comfy_execution.validation import validate_node_input
 from comfy_execution.progress import get_progress_state, reset_progress_state, add_progress_handler, WebUIProgressHandler
+from comfy_execution.jobs import extract_workflow_id
 from comfy_execution.utils import CurrentNodeContext
 from comfy_api.internal import _ComfyNodeInternal, _NodeOutputInternal, first_real_override, is_class, make_locked_method_func
 from comfy_api.latest import io, _io
@@ -417,15 +418,15 @@ def _is_intermediate_output(dynprompt, node_id):
     class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
     return getattr(class_def, 'HAS_INTERMEDIATE_OUTPUT', False)
 
-def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs):
+def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, workflow_id, ui_outputs):
     if server.client_id is None:
         return
     cached_ui = cached.ui or {}
-    server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, server.client_id)
+    server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id, "workflow_id": workflow_id }, server.client_id)
     if cached.ui is not None:
         ui_outputs[node_id] = cached.ui
 
-async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs):
+async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, workflow_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs):
     unique_id = current_item
     real_node_id = dynprompt.get_real_node_id(unique_id)
     display_node_id = dynprompt.get_display_node_id(unique_id)
@@ -435,7 +436,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
     class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
     cached = await caches.outputs.get(unique_id)
     if cached is not None:
-        _send_cached_ui(server, unique_id, display_node_id, cached, prompt_id, ui_outputs)
+        _send_cached_ui(server, unique_id, display_node_id, cached, prompt_id, workflow_id, ui_outputs)
         get_progress_state().finish_progress(unique_id)
         execution_list.cache_update(unique_id, cached)
         return (ExecutionResult.SUCCESS, None, None)
@@ -483,7 +484,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
             input_data_all, missing_keys, v3_data = get_input_data(inputs, class_def, unique_id, execution_list, dynprompt, extra_data)
             if server.client_id is not None:
                 server.last_node_id = display_node_id
-                server.send_sync("executing", { "node": unique_id, "display_node": display_node_id, "prompt_id": prompt_id }, server.client_id)
+                server.send_sync("executing", { "node": unique_id, "display_node": display_node_id, "prompt_id": prompt_id, "workflow_id": workflow_id }, server.client_id)
 
             obj = await caches.objects.get(unique_id)
             if obj is None:
@@ -513,6 +514,7 @@ def execution_block_cb(block):
                 if block.message is not None:
                     mes = {
                         "prompt_id": prompt_id,
+                        "workflow_id": workflow_id,
                         "node_id": unique_id,
                         "node_type": class_type,
                         "executed": list(executed),
@@ -561,7 +563,7 @@ async def await_completion():
                 "output": output_ui
             }
             if server.client_id is not None:
-                server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": output_ui, "prompt_id": prompt_id }, server.client_id)
+                server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": output_ui, "prompt_id": prompt_id, "workflow_id": workflow_id }, server.client_id)
         if has_subgraph:
             cached_outputs = []
             new_node_ids = []
@@ -658,6 +660,7 @@ def reset(self):
         self.caches = CacheSet(cache_type=self.cache_type, cache_args=self.cache_args)
         self.status_messages = []
         self.success = True
+        self.workflow_id = None
 
     def add_message(self, event, data: dict, broadcast: bool):
         data = {
@@ -677,6 +680,7 @@ def handle_execution_error(self, prompt_id, prompt, current_outputs, executed, e
         if isinstance(ex, comfy.model_management.InterruptProcessingException):
             mes = {
                 "prompt_id": prompt_id,
+                "workflow_id": self.workflow_id,
                 "node_id": node_id,
                 "node_type": class_type,
                 "executed": list(executed),
@@ -685,6 +689,7 @@ def handle_execution_error(self, prompt_id, prompt, current_outputs, executed, e
         else:
             mes = {
                 "prompt_id": prompt_id,
+                "workflow_id": self.workflow_id,
                 "node_id": node_id,
                 "node_type": class_type,
                 "executed": list(executed),
@@ -723,7 +728,9 @@ async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=
             self.server.client_id = None
 
         self.status_messages = []
-        self.add_message("execution_start", { "prompt_id": prompt_id}, broadcast=False)
+        self.workflow_id = extract_workflow_id(extra_data)
+        self.server.last_workflow_id = self.workflow_id
+        self.add_message("execution_start", { "prompt_id": prompt_id, "workflow_id": self.workflow_id }, broadcast=False)
 
         self._notify_prompt_lifecycle("start", prompt_id)
         ram_headroom = int(self.cache_args["ram"] * (1024 ** 3))
@@ -733,7 +740,7 @@ async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=
         try:
             with torch.inference_mode():
                 dynamic_prompt = DynamicPrompt(prompt)
-                reset_progress_state(prompt_id, dynamic_prompt)
+                reset_progress_state(prompt_id, dynamic_prompt, self.workflow_id)
                 add_progress_handler(WebUIProgressHandler(self.server))
                 is_changed_cache = IsChangedCache(prompt_id, dynamic_prompt, self.caches.outputs)
                 for cache in self.caches.all:
@@ -751,7 +758,7 @@ async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=
 
                 comfy.model_management.cleanup_models_gc()
                 self.add_message("execution_cached",
-                              { "nodes": cached_nodes, "prompt_id": prompt_id},
+                              { "nodes": cached_nodes, "prompt_id": prompt_id, "workflow_id": self.workflow_id },
                               broadcast=False)
                 pending_subgraph_results = {}
                 pending_async_nodes = {} # TODO - Unify this with pending_subgraph_results
@@ -769,7 +776,7 @@ async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=
                         break
 
                     assert node_id is not None, "Node ID should not be None at this point"
-                    result, error, ex = await execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_node_outputs)
+                    result, error, ex = await execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, self.workflow_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_node_outputs)
                     self.success = result != ExecutionResult.FAILURE
                     if result == ExecutionResult.FAILURE:
                         self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
@@ -793,8 +800,8 @@ async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=
                         cached = await self.caches.outputs.get(node_id)
                         if cached is not None:
                             display_node_id = dynamic_prompt.get_display_node_id(node_id)
-                            _send_cached_ui(self.server, node_id, display_node_id, cached, prompt_id, ui_node_outputs)
-                    self.add_message("execution_success", { "prompt_id": prompt_id }, broadcast=False)
+                            _send_cached_ui(self.server, node_id, display_node_id, cached, prompt_id, self.workflow_id, ui_node_outputs)
+                    self.add_message("execution_success", { "prompt_id": prompt_id, "workflow_id": self.workflow_id }, broadcast=False)
 
                 ui_outputs = {}
                 meta_outputs = {}
@@ -811,6 +818,8 @@ async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=
         finally:
             comfy.memory_management.set_ram_cache_release_state(None, 0)
             self._notify_prompt_lifecycle("end", prompt_id)
+            self.server.last_workflow_id = None
+            self.workflow_id = None
 
 
 async def validate_inputs(prompt_id, prompt, item, validated, visiting=None):
diff --git a/main.py b/main.py
@@ -29,6 +29,7 @@
 import sys
 from comfy_execution.progress import get_progress_state
 from comfy_execution.utils import get_executing_context
+from comfy_execution.jobs import extract_workflow_id
 from comfy_api import feature_flags
 from app.database.db import init_db, dependencies_available
 
@@ -317,6 +318,12 @@ def prompt_worker(q, server_instance):
             for k in sensitive:
                 extra_data[k] = sensitive[k]
 
+            # Capture the workflow id for this prompt before execution: the
+            # executor clears server.last_workflow_id in its finally block, so
+            # reading it after e.execute() returns would emit workflow_id=None
+            # on the terminal "executing" reset below.
+            workflow_id = extract_workflow_id(extra_data)
+
             asset_seeder.pause()
             e.execute(item[2], prompt_id, extra_data, item[4])
 
@@ -330,7 +337,7 @@ def prompt_worker(q, server_instance):
                             completed=e.success,
                             messages=e.status_messages), process_item=remove_sensitive)
             if server_instance.client_id is not None:
-                server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server_instance.client_id)
+                server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id, "workflow_id": workflow_id}, server_instance.client_id)
 
             current_time = time.perf_counter()
             execution_time = current_time - execution_start_time
@@ -393,7 +400,7 @@ def hook(value, total, preview_image, prompt_id=None, node_id=None):
             prompt_id = server_instance.last_prompt_id
         if node_id is None:
             node_id = server_instance.last_node_id
-        progress = {"value": value, "max": total, "prompt_id": prompt_id, "node": node_id}
+        progress = {"value": value, "max": total, "prompt_id": prompt_id, "workflow_id": getattr(server_instance, 'last_workflow_id', None), "node": node_id}
         get_progress_state().update_progress(node_id, value, total, preview_image)
 
         server_instance.send_sync("progress", progress, server_instance.client_id)
diff --git a/server.py b/server.py
@@ -275,7 +275,11 @@ async def websocket_handler(request):
                 await self.send("status", {"status": self.get_queue_info(), "sid": sid}, sid)
                 # On reconnect if we are the currently executing client send the current node
                 if self.client_id == sid and self.last_node_id is not None:
-                    await self.send("executing", { "node": self.last_node_id }, sid)
+                    await self.send("executing", {
+                        "node": self.last_node_id,
+                        "prompt_id": getattr(self, "last_prompt_id", None),
+                        "workflow_id": getattr(self, "last_workflow_id", None),
+                    }, sid)
 
                 # Flag to track if we've received the first message
                 first_message = True
diff --git a/tests-unit/execution_test/test_workflow_id_in_ws_messages.py b/tests-unit/execution_test/test_workflow_id_in_ws_messages.py
diff --git a/tests/execution/test_jobs.py b/tests/execution/test_jobs.py
