wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-31 03:17:23 +08:00
Code Issues Actions 14 Packages Projects Releases Wiki Activity

refactor(server): spread envelope keys onto payload at top level

Browse Source 
Switch the wire shape from nested ``metadata: {workflow_id: ...}`` to
spreading the envelope's keys directly onto each event payload. The
contract on the websocket is now identical to the prior workflow-id-on-
events work — consumers read ``event.workflow_id`` directly — but the
core executor still has no concept of workflow scope; the envelope is
captured at submission and decorated at the server transport layer.

Server-emitted fields always win on collision (``{**envelope, **d}``):
a misbehaving client cannot shadow ``prompt_id``, ``node``, etc. by
stamping the same key in their submission envelope.
This commit is contained in:
Deep Mehta 2026-05-14 21:18:43 -07:00
parent fd89498eac
commit fc9820ebb9
2 changed files with 49 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
app/prompt_metadata.py
View File

@ -125,36 +125,40 @@ def inject_envelope(
data: Any,
envelope_lookup: Callable[[str], Optional[dict]],
) -> Any:
"""Return ``data`` with a per-prompt ``metadata`` envelope attached.
"""Return ``data`` with the per-prompt envelope's keys spread onto it.
``envelope_lookup`` is called with the payload's ``prompt_id`` and is
expected to return the registered envelope or ``None``. This keeps
the function pure and avoids depending on any specific storage.
The envelope's keys are merged onto the payload at the top level so
consumers can read them directly (e.g. ``event.workflow_id``) —
matching the wire shape of the prior workflow-id-on-events work and
avoiding an extra nesting hop for clients. Server-emitted fields on
the payload always win on collision (``{**envelope, **d}``); a
misbehaving client cannot shadow ``prompt_id``, ``node``, etc.
Two payload shapes are handled:
- **dict** carrying ``prompt_id``. A shallow copy is returned with a
``metadata`` key set to the envelope.
- **dict** carrying ``prompt_id``. A shallow copy is returned with
the envelope's keys merged onto it.
- **(preview_image, metadata_dict) tuple** — the format used by
``PREVIEW_IMAGE_WITH_METADATA``. Only the inner dict is augmented;
the binary preview is passed through by reference.
No-op for payloads without a ``prompt_id``, payloads already
declaring their own ``metadata`` field, prompts with no registered
envelope, or any other payload shape.
No-op for payloads without a ``prompt_id``, prompts with no
registered envelope, or any other payload shape.
"""
def inject(d: dict) -> dict:
if not isinstance(d, dict):
return d
if "metadata" in d:
return d
prompt_id = d.get("prompt_id")
if not prompt_id:
return d
envelope = envelope_lookup(prompt_id)
if envelope is None:
return d
return {**d, "metadata": envelope}
return {**envelope, **d}
if isinstance(data, dict):
return inject(data)

71
tests-unit/app_test/test_prompt_metadata.py
View File

@ -166,12 +166,15 @@ class TestInjectEnvelope:
def _lookup(table):
return table.get
def test_injects_envelope_on_dict_with_known_prompt_id(self):
lookup = self._lookup({"p1": {"workflow_id": "wf-1"}})
def test_spreads_envelope_keys_onto_payload(self):
"""Envelope keys are merged at the top level so consumers can
read them directly (e.g. ``event.workflow_id``)."""
lookup = self._lookup({"p1": {"workflow_id": "wf-1", "trace_id": "t-9"}})
assert inject_envelope({"node": "5", "prompt_id": "p1"}, lookup) == {
"node": "5",
"prompt_id": "p1",
"metadata": {"workflow_id": "wf-1"},
"workflow_id": "wf-1",
"trace_id": "t-9",
}
def test_passthrough_when_prompt_id_not_registered(self):
@ -184,20 +187,28 @@ class TestInjectEnvelope:
data = {"status": "ok"}
assert inject_envelope(data, lookup) == data
def test_passthrough_when_payload_already_has_metadata(self):
"""If a caller has already set a ``metadata`` field, the
function must not overwrite it."""
lookup = self._lookup({"p1": {"workflow_id": "wf-injected"}})
data = {"prompt_id": "p1", "metadata": {"workflow_id": "wf-caller"}}
result = inject_envelope(data, lookup)
assert result == data
assert result["metadata"] == {"workflow_id": "wf-caller"}
def test_server_keys_win_on_collision_with_envelope(self):
"""A misbehaving client cannot shadow server-emitted fields by
stamping the same key in their submission envelope."""
lookup = self._lookup({
"p1": {"prompt_id": "client-claimed", "node": "spoofed", "workflow_id": "wf-1"}
})
result = inject_envelope({"prompt_id": "p1", "node": "5"}, lookup)
assert result["prompt_id"] == "p1"
assert result["node"] == "5"
assert result["workflow_id"] == "wf-1"
def test_does_not_mutate_input_dict(self):
lookup = self._lookup({"p1": {"workflow_id": "wf-1"}})
original = {"node": "5", "prompt_id": "p1"}
inject_envelope(original, lookup)
assert "metadata" not in original
assert "workflow_id" not in original
def test_does_not_mutate_envelope_dict(self):
envelope = {"workflow_id": "wf-1"}
lookup = self._lookup({"p1": envelope})
inject_envelope({"prompt_id": "p1", "node": "5"}, lookup)
assert envelope == {"workflow_id": "wf-1"}
def test_injects_into_inner_dict_of_preview_metadata_tuple(self):
"""``PREVIEW_IMAGE_WITH_METADATA`` payloads arrive as
@ -212,9 +223,9 @@ class TestInjectEnvelope:
assert result[1] == {
"node_id": "5",
"prompt_id": "p1",
"metadata": {"workflow_id": "wf-1"},
"workflow_id": "wf-1",
}
assert "metadata" not in inner
assert "workflow_id" not in inner
def test_preview_tuple_passthrough_when_no_envelope_registered(self):
lookup = self._lookup({})
@ -223,13 +234,6 @@ class TestInjectEnvelope:
result = inject_envelope((preview_image, inner), lookup)
assert result == (preview_image, inner)
def test_preview_tuple_passthrough_when_inner_already_has_metadata(self):
lookup = self._lookup({"p1": {"workflow_id": "wf-injected"}})
preview_image = ("PNG", object(), 256)
inner = {"node_id": "5", "prompt_id": "p1", "metadata": {"x": "1"}}
result = inject_envelope((preview_image, inner), lookup)
assert result == (preview_image, inner)
@pytest.mark.parametrize("payload", [b"raw-bytes", None, 42])
def test_non_dict_non_tuple_payloads_passthrough(self, payload):
lookup = self._lookup({"p1": {"workflow_id": "wf-1"}})
@ -251,9 +255,9 @@ class TestInjectEnvelope:
second = inject_envelope({"prompt_id": "p1"}, store.get)
del store["p1"]
third = inject_envelope({"prompt_id": "p1"}, store.get)
assert first["metadata"] == {"workflow_id": "wf-1"}
assert second["metadata"] == {"workflow_id": "wf-2"}
assert "metadata" not in third
assert first["workflow_id"] == "wf-1"
assert second["workflow_id"] == "wf-2"
assert "workflow_id" not in third
class TestPromptMetadataStore:
@ -269,17 +273,18 @@ class TestPromptMetadataStore:
assert injected == {
"node": "5",
"prompt_id": "p1",
"metadata": {"workflow_id": "wf-1"},
"workflow_id": "wf-1",
}
store.unregister("p1")
passthrough = store.inject({"node": "5", "prompt_id": "p1"})
assert "metadata" not in passthrough
assert "workflow_id" not in passthrough
def test_register_with_no_derivable_envelope_is_noop(self):
store = PromptMetadataStore()
store.register("p1", {})
assert "p1" not in store
assert store.inject({"prompt_id": "p1"}) == {"prompt_id": "p1"}
data = {"prompt_id": "p1"}
assert store.inject(data) == data
def test_register_with_oversized_envelope_is_noop(self):
"""Sanitization rejection means nothing is registered — the
@ -310,11 +315,9 @@ class TestPromptMetadataStore:
assert "p4" in store
# The newer entries are still injectable.
assert store.inject({"prompt_id": "p4"})["metadata"] == {
"workflow_id": "wf-4"
}
assert store.inject({"prompt_id": "p4"})["workflow_id"] == "wf-4"
# The evicted one is gone.
assert "metadata" not in store.inject({"prompt_id": "p1"})
assert "workflow_id" not in store.inject({"prompt_id": "p1"})
def test_register_after_unregister_does_not_count_against_capacity(self):
"""Normal lifecycle: register, unregister, register many — the
@ -330,9 +333,7 @@ class TestPromptMetadataStore:
store = PromptMetadataStore()
store.register("p1", {"metadata": {"workflow_id": "wf-1"}})
store.register("p1", {"metadata": {"workflow_id": "wf-2"}})
assert store.inject({"prompt_id": "p1"})["metadata"] == {
"workflow_id": "wf-2"
}
assert store.inject({"prompt_id": "p1"})["workflow_id"] == "wf-2"
def test_inject_with_no_registrations_is_passthrough(self):
store = PromptMetadataStore()
@ -345,5 +346,5 @@ class TestPromptMetadataStore:
result = store.inject((b"image-bytes", {"prompt_id": "p1"}))
assert result == (b"image-bytes", {
"prompt_id": "p1",
"metadata": {"workflow_id": "wf-1"},
"workflow_id": "wf-1",
})