mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-06-18 05:49:41 +08:00
5396b4fe67
PR #13684 added workflow_id directly to ~9 dict literals across execution.py, progress.py and main.py, along with executor.workflow_id and server.last_workflow_id state. It was reverted because the execution layer should not know about workflow concepts and because a finally-clear race emitted workflow_id=None on the terminal "executing" frame. Instead, register per-prompt metadata on PromptServer at submission time and merge it onto outbound WebSocket payloads inside send_sync. The merge keys off prompt_id (already present on every execution event), so execution.py stays workflow-agnostic. Metadata is unregistered in main.py's queue loop AFTER the terminal executing send, which structurally removes the race. - New comfy_execution/metadata.py: PromptMetadata TypedDict + build_prompt_metadata + merge_prompt_metadata helpers. - PromptServer: prompt_metadata registry (lock-protected), register on post_prompt, merge in send_sync, expose get_prompt_metadata. - jobs.py: extracted extract_workflow_id with strict isinstance guards; _extract_job_metadata delegates. - main.py: try/finally around the queue iteration; unregister after the terminal "executing: {node: None}" send. - execution.py PromptQueue: drop registry entries on wipe_queue / delete_queue_item so cancellations don't leak. - progress.py: look up workflow_id from the server registry for the per-node nested copies and the binary preview metadata, matching #13684's wire shape so the frontend needs no changes. - Tests: tests-unit/server_test/test_prompt_metadata.py covers the merge, the passthrough cases (no prompt_id, unknown prompt_id, binary payloads), and the terminal-frame race regression.
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
"""Per-prompt metadata propagated alongside execution WebSocket events.
|
|
|
|
The execution layer (``execution.py``) is intentionally kept agnostic of
|
|
workflow-level concepts. Instead, ``PromptServer`` registers per-``prompt_id``
|
|
metadata at submission time and merges it onto outgoing WebSocket payloads in
|
|
``send_sync``. Today only ``workflow_id`` is propagated; the structure is a
|
|
``TypedDict`` so additional keys can be added without churn at the call sites.
|
|
"""
|
|
|
|
import threading
|
|
from typing import Optional, TypedDict
|
|
|
|
from comfy_execution.jobs import extract_workflow_id
|
|
|
|
|
|
class PromptMetadata(TypedDict, total=False):
|
|
workflow_id: Optional[str]
|
|
|
|
|
|
def build_prompt_metadata(extra_data: Optional[dict]) -> PromptMetadata:
|
|
"""Build a ``PromptMetadata`` snapshot from a prompt's ``extra_data``.
|
|
|
|
Returns an empty dict when no recognized metadata is present so callers can
|
|
skip registering anything in that case.
|
|
"""
|
|
meta: PromptMetadata = {}
|
|
workflow_id = extract_workflow_id(extra_data)
|
|
if workflow_id is not None:
|
|
meta["workflow_id"] = workflow_id
|
|
return meta
|
|
|
|
|
|
def merge_prompt_metadata(
|
|
registry: dict,
|
|
lock: threading.Lock,
|
|
data,
|
|
):
|
|
"""Return ``data`` with the registered metadata for its ``prompt_id`` merged
|
|
top-level. The event payload wins on conflict, and non-dict payloads (e.g.
|
|
the binary preview tuple) pass through untouched.
|
|
"""
|
|
if not isinstance(data, dict):
|
|
return data
|
|
prompt_id = data.get("prompt_id")
|
|
if not prompt_id:
|
|
return data
|
|
with lock:
|
|
meta = registry.get(prompt_id)
|
|
if not meta:
|
|
return data
|
|
return {**meta, **data}
|