mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-06-17 21:39:45 +08:00
85a12d0a83
Adversarial review caught that a LIFO stack keyed by ``prompt_id`` still
mis-attributes events when queue execution order differs from registration
order: a second submission with the same ``prompt_id`` lands on top of the
stack, so the first prompt's events read the wrong workflow_id while it
runs, and the first's ``unregister`` then pops the second prompt's entry.
Replace the stack with an internal monotonic token. ``post_prompt``
registers metadata and stashes the returned token on ``extra_data`` under
``PROMPT_METADATA_TOKEN_KEY``. ``main.py``'s queue worker pulls the token
out, pins it on ``PromptServer.active_prompt_metadata_token`` for the
prompt's execution, and clears + unregisters in ``finally``. The merge in
``send_sync`` reads the active token, so each prompt's events are merged
with its own metadata regardless of ``prompt_id`` collisions.
- comfy_execution/metadata.py: ``merge_prompt_metadata`` now takes an
active token; registry is ``dict[int, PromptMetadata]``; new
``PROMPT_METADATA_TOKEN_KEY`` constant for the extra_data carrier.
- server.py: ``register_prompt_metadata`` returns a token (or ``None``
when no metadata applies); ``unregister`` takes a token;
``get_active_prompt_metadata`` snapshots the pinned entry.
- main.py: pops the token from extra_data, pins on the server, clears
after the terminal "executing: {node: None}" send.
- execution.py ``PromptQueue``: wipe_queue / delete_queue_item now
unregister by token extracted from each item's extra_data.
- comfy_execution/progress.py: reads workflow_id via
``get_active_prompt_metadata`` rather than per-prompt_id lookup.
- tests: unit tests updated for the token signature, plus a real E2E
test (test_prompt_metadata_e2e.py) that instantiates the actual
PromptServer and verifies same-prompt_id-different-workflow_id
submissions don't cross-attribute.
Verified end-to-end against a live ComfyUI server: two submissions with
identical client-supplied prompt_id but different workflow_id each emit
their full execution event stream (execution_start, execution_cached,
executing, executed, execution_success, progress_state, terminal executing)
with the correct workflow_id top-level. 68 / 68 tests pass.
87 lines
3.3 KiB
Python
87 lines
3.3 KiB
Python
"""Per-prompt metadata propagated alongside execution WebSocket events.
|
|
|
|
The execution layer (``execution.py``) is intentionally kept agnostic of
|
|
workflow-level concepts. ``PromptServer`` registers 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.
|
|
|
|
Identity model — registry is keyed by an internal monotonic token, NOT by
|
|
``prompt_id``. ``post_prompt`` accepts a client-supplied ``prompt_id``
|
|
verbatim, so two prompts can be queued with the same id and a registry keyed
|
|
only by ``prompt_id`` would misattribute events when queue order differs from
|
|
registration order. ``main.py``'s queue worker pins the active token on the
|
|
server for the duration of one prompt's execution and ``send_sync`` reads that
|
|
token — so each prompt's events get its own registered metadata regardless of
|
|
``prompt_id`` collisions.
|
|
"""
|
|
|
|
import threading
|
|
from typing import Optional, TypedDict
|
|
|
|
from comfy_execution.jobs import extract_workflow_id
|
|
|
|
|
|
PROMPT_METADATA_TOKEN_KEY = "__prompt_metadata_token"
|
|
|
|
|
|
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,
|
|
active_token: Optional[int],
|
|
data,
|
|
):
|
|
"""Return ``data`` with the metadata for the currently-active token merged
|
|
top-level. The event payload wins on conflict, and non-dict payloads (e.g.
|
|
the binary preview tuple) pass through untouched.
|
|
|
|
The active token is set by ``main.py``'s queue worker around each
|
|
``e.execute(...)``. The merge happens only when a token is currently active
|
|
*and* the payload carries a ``prompt_id`` — using ``prompt_id`` purely as a
|
|
marker that this is an execution event meant to receive metadata, not as
|
|
the lookup key. This makes the merge immune to ``prompt_id`` collisions.
|
|
"""
|
|
if not isinstance(data, dict):
|
|
return data
|
|
if not data.get("prompt_id"):
|
|
return data
|
|
if active_token is None:
|
|
return data
|
|
with lock:
|
|
meta = registry.get(active_token)
|
|
if not meta:
|
|
return data
|
|
return {**meta, **data}
|