mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-06-17 21:39:45 +08:00
b0c05af67f
Two issues raised against the per-prompt metadata registry: 1. Client-supplied prompt_id can collide (post_prompt accepts the id verbatim). With a flat dict-keyed registry, the second submission clobbered the first and a single unregister could erase metadata still needed by the other prompt. Now stored as a LIFO stack per prompt_id — most recent registration wins on merge, unregister pops one entry, the key is dropped only when the stack drains. 2. BinaryEventTypes.TEXT (send_progress_text) bypasses the metadata merge because the payload is bytes, and the wire format has no prompt_id / workflow_id field. The merge can't fix this without a wire-format change + frontend feature flag, which is out of scope for FE-745. Inside scope: default the sid to PromptServer.client_id so other clients no longer silently receive untagged text frames. Cross-tab isolation inside a single client still depends on the wire-format follow-up. - comfy_execution/metadata.py: registry is dict[str, list[PromptMetadata]]; merge_prompt_metadata reads stack[-1]; new resolve_progress_text_sid helper extracted so the routing default is unit-testable without the full server import chain. - server.py: register_prompt_metadata appends to the stack; unregister_prompt_metadata pops; get_prompt_metadata returns a copy of the top entry; send_progress_text routes through resolve_progress_text_sid. - tests: collision LIFO behavior, sid resolution default, and the existing merge tests updated to the stack shape. 16 new assertions in this file, 104/104 pass overall.
72 lines
2.5 KiB
Python
72 lines
2.5 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 resolve_progress_text_sid(sid, default_sid):
|
|
"""Pick the recipient for a ``send_progress_text`` binary frame.
|
|
|
|
Returns ``default_sid`` (typically ``PromptServer.client_id`` — the client
|
|
that submitted the active prompt) when the caller didn't pin a specific
|
|
socket. This narrows the audience for text status updates from "every
|
|
connected client" to "the client running this prompt", matching the
|
|
cross-client isolation other execution events already have.
|
|
|
|
Splitting this out keeps the unit test independent of the full ``server``
|
|
import chain.
|
|
"""
|
|
return default_sid if sid is None else sid
|
|
|
|
|
|
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.
|
|
|
|
The registry is a stack per ``prompt_id`` (``dict[str, list[PromptMetadata]]``)
|
|
so duplicate submissions of the same ``prompt_id`` don't clobber each
|
|
other's metadata; the most recently registered entry wins.
|
|
"""
|
|
if not isinstance(data, dict):
|
|
return data
|
|
prompt_id = data.get("prompt_id")
|
|
if not prompt_id:
|
|
return data
|
|
with lock:
|
|
stack = registry.get(prompt_id)
|
|
meta = stack[-1] if stack else None
|
|
if not meta:
|
|
return data
|
|
return {**meta, **data}
|