feat(server): per-prompt metadata envelope on websocket events

Replaces the workflow_id-on-every-event approach (#13684, reverted in
#13901) with a generic metadata envelope captured at submission and
injected at the server-side send chokepoint.

- POST /prompt accepts an opaque ``extra_data.metadata`` dict (falls
  back to synthesizing ``{"workflow_id": <id>}`` from
  ``extra_pnginfo.workflow.id`` so existing frontends keep working).
- ``PromptServer`` owns a ``prompt_id -> metadata`` map populated at
  submission, drained when the prompt finishes. ``send_sync`` injects
  the envelope into any outbound payload that carries a ``prompt_id``,
  including the ``(preview_image, metadata_dict)`` tuple used by
  ``PREVIEW_IMAGE_WITH_METADATA``. WS reconnect path carries it too.
- Pure helpers live in ``app/prompt_metadata.py`` so the execution
  layer never depends on workflow concepts and the helpers can be
  unit-tested without torch.

Execution layer (``execution.py``, ``comfy_execution/*``) and the jobs
API are unchanged. Backward compatible: existing fields and shapes are
preserved, only an additional ``metadata`` field is attached when
present.
This commit is contained in:
Deep Mehta
2026-05-14 20:47:00 -07:00
parent 77e2ed5e01
commit 74cfcaa318
4 changed files with 310 additions and 1 deletions

View File

@@ -44,6 +44,10 @@ from app.model_manager import ModelFileManager
from app.custom_node_manager import CustomNodeManager
from app.subgraph_manager import SubgraphManager
from app.node_replace_manager import NodeReplaceManager
from app.prompt_metadata import (
extract_envelope_from_extra_data,
inject_envelope,
)
from typing import Optional, Union
from api_server.routes.internal.internal_routes import InternalRoutes
from protocol import BinaryEventTypes
@@ -250,8 +254,14 @@ class PromptServer():
routes = web.RouteTableDef()
self.routes = routes
self.last_node_id = None
self.last_prompt_id = None
self.client_id = None
# prompt_id -> metadata envelope captured at submission and injected
# into outbound execution events. Keeps the workflow scope (and any
# other client-supplied tags) out of the execution layer.
self._prompt_metadata: dict[str, dict] = {}
self.on_prompt_handlers = []
@routes.get('/ws')
@@ -275,7 +285,10 @@ class PromptServer():
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", self._inject_prompt_metadata({
"node": self.last_node_id,
"prompt_id": self.last_prompt_id,
}), sid)
# Flag to track if we've received the first message
first_message = True
@@ -955,6 +968,7 @@ class PromptServer():
if sensitive_val in extra_data:
sensitive[sensitive_val] = extra_data.pop(sensitive_val)
extra_data["create_time"] = int(time.time() * 1000) # timestamp in milliseconds
self.register_prompt_metadata(prompt_id, extra_data)
self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute, sensitive))
response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]}
return web.json_response(response)
@@ -1216,7 +1230,23 @@ class PromptServer():
elif sid in self.sockets:
await send_socket_catch_exception(self.sockets[sid].send_json, message)
def register_prompt_metadata(self, prompt_id: str, extra_data) -> None:
"""Record per-prompt metadata for injection into outbound execution
events. Called at submission, before the prompt is queued."""
envelope = extract_envelope_from_extra_data(extra_data)
if envelope is not None:
self._prompt_metadata[prompt_id] = envelope
def unregister_prompt_metadata(self, prompt_id: str) -> None:
"""Drop the per-prompt metadata envelope. Called after the prompt
has finished executing and its terminal events have been queued."""
self._prompt_metadata.pop(prompt_id, None)
def _inject_prompt_metadata(self, data):
return inject_envelope(data, self._prompt_metadata.get)
def send_sync(self, event, data, sid=None):
data = self._inject_prompt_metadata(data)
self.loop.call_soon_threadsafe(
self.messages.put_nowait, (event, data, sid))