Add workflow_id to all websocket messages

This commit addresses BE-672 by ensuring all execution-related websocket
messages include the workflow_id field when available.

Changes:
- Added extract_workflow_id() helper function in comfy_execution/jobs.py
  to extract workflow_id from extra_data
- Updated execution.py to include workflow_id in all websocket messages:
  - execution_start
  - execution_cached
  - execution_success
  - execution_error
  - execution_interrupted
  - executing
  - executed (including cached UI)
- Updated main.py to include workflow_id in:
  - progress messages (via hijack_progress hook)
  - final executing message (node=None)
- Updated comfy_execution/progress.py to include workflow_id in:
  - progress_state messages
  - preview image metadata

The workflow_id is extracted from extra_data['extra_pnginfo']['workflow']['id']
and is conditionally included in messages only when present, maintaining
backward compatibility with workflows that don't have this field.

Fixes: BE-672

Co-authored-by: Luke Mino-Altherr <luke-mino-altherr@users.noreply.github.com>
This commit is contained in:
Cursor Agent
2026-05-01 00:35:43 +00:00
parent e9c311b245
commit 7578d1989f
4 changed files with 80 additions and 22 deletions

View File

@@ -182,8 +182,11 @@ class WebUIProgressHandler(ProgressHandler):
# Send a combined progress_state message with all node states
# Include client_id to ensure message is only sent to the initiating client
message = {"prompt_id": prompt_id, "nodes": active_nodes}
if self.registry.workflow_id is not None:
message["workflow_id"] = self.registry.workflow_id
self.server_instance.send_sync(
"progress_state", {"prompt_id": prompt_id, "nodes": active_nodes}, self.server_instance.client_id
"progress_state", message, self.server_instance.client_id
)
@override
@@ -223,6 +226,8 @@ class WebUIProgressHandler(ProgressHandler):
),
"real_node_id": self.registry.dynprompt.get_real_node_id(node_id),
}
if self.registry.workflow_id is not None:
metadata["workflow_id"] = self.registry.workflow_id
self.server_instance.send_sync(
BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA,
(image, metadata),
@@ -240,9 +245,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,7 +328,7 @@ class ProgressRegistry:
# 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
@@ -330,7 +336,7 @@ def reset_progress_state(prompt_id: str, dynprompt: "DynamicPrompt") -> 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: