mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-06-01 20:07:37 +08:00
refactor(server): spread envelope keys onto payload at top level
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:
parent
fd89498eac
commit
fc9820ebb9
@ -125,36 +125,40 @@ def inject_envelope(
|
|||||||
data: Any,
|
data: Any,
|
||||||
envelope_lookup: Callable[[str], Optional[dict]],
|
envelope_lookup: Callable[[str], Optional[dict]],
|
||||||
) -> Any:
|
) -> 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
|
``envelope_lookup`` is called with the payload's ``prompt_id`` and is
|
||||||
expected to return the registered envelope or ``None``. This keeps
|
expected to return the registered envelope or ``None``. This keeps
|
||||||
the function pure and avoids depending on any specific storage.
|
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:
|
Two payload shapes are handled:
|
||||||
|
|
||||||
- **dict** carrying ``prompt_id``. A shallow copy is returned with a
|
- **dict** carrying ``prompt_id``. A shallow copy is returned with
|
||||||
``metadata`` key set to the envelope.
|
the envelope's keys merged onto it.
|
||||||
- **(preview_image, metadata_dict) tuple** — the format used by
|
- **(preview_image, metadata_dict) tuple** — the format used by
|
||||||
``PREVIEW_IMAGE_WITH_METADATA``. Only the inner dict is augmented;
|
``PREVIEW_IMAGE_WITH_METADATA``. Only the inner dict is augmented;
|
||||||
the binary preview is passed through by reference.
|
the binary preview is passed through by reference.
|
||||||
|
|
||||||
No-op for payloads without a ``prompt_id``, payloads already
|
No-op for payloads without a ``prompt_id``, prompts with no
|
||||||
declaring their own ``metadata`` field, prompts with no registered
|
registered envelope, or any other payload shape.
|
||||||
envelope, or any other payload shape.
|
|
||||||
"""
|
"""
|
||||||
def inject(d: dict) -> dict:
|
def inject(d: dict) -> dict:
|
||||||
if not isinstance(d, dict):
|
if not isinstance(d, dict):
|
||||||
return d
|
return d
|
||||||
if "metadata" in d:
|
|
||||||
return d
|
|
||||||
prompt_id = d.get("prompt_id")
|
prompt_id = d.get("prompt_id")
|
||||||
if not prompt_id:
|
if not prompt_id:
|
||||||
return d
|
return d
|
||||||
envelope = envelope_lookup(prompt_id)
|
envelope = envelope_lookup(prompt_id)
|
||||||
if envelope is None:
|
if envelope is None:
|
||||||
return d
|
return d
|
||||||
return {**d, "metadata": envelope}
|
return {**envelope, **d}
|
||||||
|
|
||||||
if isinstance(data, dict):
|
if isinstance(data, dict):
|
||||||
return inject(data)
|
return inject(data)
|
||||||
|
|||||||
@ -166,12 +166,15 @@ class TestInjectEnvelope:
|
|||||||
def _lookup(table):
|
def _lookup(table):
|
||||||
return table.get
|
return table.get
|
||||||
|
|
||||||
def test_injects_envelope_on_dict_with_known_prompt_id(self):
|
def test_spreads_envelope_keys_onto_payload(self):
|
||||||
lookup = self._lookup({"p1": {"workflow_id": "wf-1"}})
|
"""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) == {
|
assert inject_envelope({"node": "5", "prompt_id": "p1"}, lookup) == {
|
||||||
"node": "5",
|
"node": "5",
|
||||||
"prompt_id": "p1",
|
"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):
|
def test_passthrough_when_prompt_id_not_registered(self):
|
||||||
@ -184,20 +187,28 @@ class TestInjectEnvelope:
|
|||||||
data = {"status": "ok"}
|
data = {"status": "ok"}
|
||||||
assert inject_envelope(data, lookup) == data
|
assert inject_envelope(data, lookup) == data
|
||||||
|
|
||||||
def test_passthrough_when_payload_already_has_metadata(self):
|
def test_server_keys_win_on_collision_with_envelope(self):
|
||||||
"""If a caller has already set a ``metadata`` field, the
|
"""A misbehaving client cannot shadow server-emitted fields by
|
||||||
function must not overwrite it."""
|
stamping the same key in their submission envelope."""
|
||||||
lookup = self._lookup({"p1": {"workflow_id": "wf-injected"}})
|
lookup = self._lookup({
|
||||||
data = {"prompt_id": "p1", "metadata": {"workflow_id": "wf-caller"}}
|
"p1": {"prompt_id": "client-claimed", "node": "spoofed", "workflow_id": "wf-1"}
|
||||||
result = inject_envelope(data, lookup)
|
})
|
||||||
assert result == data
|
result = inject_envelope({"prompt_id": "p1", "node": "5"}, lookup)
|
||||||
assert result["metadata"] == {"workflow_id": "wf-caller"}
|
assert result["prompt_id"] == "p1"
|
||||||
|
assert result["node"] == "5"
|
||||||
|
assert result["workflow_id"] == "wf-1"
|
||||||
|
|
||||||
def test_does_not_mutate_input_dict(self):
|
def test_does_not_mutate_input_dict(self):
|
||||||
lookup = self._lookup({"p1": {"workflow_id": "wf-1"}})
|
lookup = self._lookup({"p1": {"workflow_id": "wf-1"}})
|
||||||
original = {"node": "5", "prompt_id": "p1"}
|
original = {"node": "5", "prompt_id": "p1"}
|
||||||
inject_envelope(original, lookup)
|
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):
|
def test_injects_into_inner_dict_of_preview_metadata_tuple(self):
|
||||||
"""``PREVIEW_IMAGE_WITH_METADATA`` payloads arrive as
|
"""``PREVIEW_IMAGE_WITH_METADATA`` payloads arrive as
|
||||||
@ -212,9 +223,9 @@ class TestInjectEnvelope:
|
|||||||
assert result[1] == {
|
assert result[1] == {
|
||||||
"node_id": "5",
|
"node_id": "5",
|
||||||
"prompt_id": "p1",
|
"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):
|
def test_preview_tuple_passthrough_when_no_envelope_registered(self):
|
||||||
lookup = self._lookup({})
|
lookup = self._lookup({})
|
||||||
@ -223,13 +234,6 @@ class TestInjectEnvelope:
|
|||||||
result = inject_envelope((preview_image, inner), lookup)
|
result = inject_envelope((preview_image, inner), lookup)
|
||||||
assert result == (preview_image, inner)
|
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])
|
@pytest.mark.parametrize("payload", [b"raw-bytes", None, 42])
|
||||||
def test_non_dict_non_tuple_payloads_passthrough(self, payload):
|
def test_non_dict_non_tuple_payloads_passthrough(self, payload):
|
||||||
lookup = self._lookup({"p1": {"workflow_id": "wf-1"}})
|
lookup = self._lookup({"p1": {"workflow_id": "wf-1"}})
|
||||||
@ -251,9 +255,9 @@ class TestInjectEnvelope:
|
|||||||
second = inject_envelope({"prompt_id": "p1"}, store.get)
|
second = inject_envelope({"prompt_id": "p1"}, store.get)
|
||||||
del store["p1"]
|
del store["p1"]
|
||||||
third = inject_envelope({"prompt_id": "p1"}, store.get)
|
third = inject_envelope({"prompt_id": "p1"}, store.get)
|
||||||
assert first["metadata"] == {"workflow_id": "wf-1"}
|
assert first["workflow_id"] == "wf-1"
|
||||||
assert second["metadata"] == {"workflow_id": "wf-2"}
|
assert second["workflow_id"] == "wf-2"
|
||||||
assert "metadata" not in third
|
assert "workflow_id" not in third
|
||||||
|
|
||||||
|
|
||||||
class TestPromptMetadataStore:
|
class TestPromptMetadataStore:
|
||||||
@ -269,17 +273,18 @@ class TestPromptMetadataStore:
|
|||||||
assert injected == {
|
assert injected == {
|
||||||
"node": "5",
|
"node": "5",
|
||||||
"prompt_id": "p1",
|
"prompt_id": "p1",
|
||||||
"metadata": {"workflow_id": "wf-1"},
|
"workflow_id": "wf-1",
|
||||||
}
|
}
|
||||||
store.unregister("p1")
|
store.unregister("p1")
|
||||||
passthrough = store.inject({"node": "5", "prompt_id": "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):
|
def test_register_with_no_derivable_envelope_is_noop(self):
|
||||||
store = PromptMetadataStore()
|
store = PromptMetadataStore()
|
||||||
store.register("p1", {})
|
store.register("p1", {})
|
||||||
assert "p1" not in store
|
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):
|
def test_register_with_oversized_envelope_is_noop(self):
|
||||||
"""Sanitization rejection means nothing is registered — the
|
"""Sanitization rejection means nothing is registered — the
|
||||||
@ -310,11 +315,9 @@ class TestPromptMetadataStore:
|
|||||||
assert "p4" in store
|
assert "p4" in store
|
||||||
|
|
||||||
# The newer entries are still injectable.
|
# The newer entries are still injectable.
|
||||||
assert store.inject({"prompt_id": "p4"})["metadata"] == {
|
assert store.inject({"prompt_id": "p4"})["workflow_id"] == "wf-4"
|
||||||
"workflow_id": "wf-4"
|
|
||||||
}
|
|
||||||
# The evicted one is gone.
|
# 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):
|
def test_register_after_unregister_does_not_count_against_capacity(self):
|
||||||
"""Normal lifecycle: register, unregister, register many — the
|
"""Normal lifecycle: register, unregister, register many — the
|
||||||
@ -330,9 +333,7 @@ class TestPromptMetadataStore:
|
|||||||
store = PromptMetadataStore()
|
store = PromptMetadataStore()
|
||||||
store.register("p1", {"metadata": {"workflow_id": "wf-1"}})
|
store.register("p1", {"metadata": {"workflow_id": "wf-1"}})
|
||||||
store.register("p1", {"metadata": {"workflow_id": "wf-2"}})
|
store.register("p1", {"metadata": {"workflow_id": "wf-2"}})
|
||||||
assert store.inject({"prompt_id": "p1"})["metadata"] == {
|
assert store.inject({"prompt_id": "p1"})["workflow_id"] == "wf-2"
|
||||||
"workflow_id": "wf-2"
|
|
||||||
}
|
|
||||||
|
|
||||||
def test_inject_with_no_registrations_is_passthrough(self):
|
def test_inject_with_no_registrations_is_passthrough(self):
|
||||||
store = PromptMetadataStore()
|
store = PromptMetadataStore()
|
||||||
@ -345,5 +346,5 @@ class TestPromptMetadataStore:
|
|||||||
result = store.inject((b"image-bytes", {"prompt_id": "p1"}))
|
result = store.inject((b"image-bytes", {"prompt_id": "p1"}))
|
||||||
assert result == (b"image-bytes", {
|
assert result == (b"image-bytes", {
|
||||||
"prompt_id": "p1",
|
"prompt_id": "p1",
|
||||||
"metadata": {"workflow_id": "wf-1"},
|
"workflow_id": "wf-1",
|
||||||
})
|
})
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user