mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-06-12 01:07:30 +08:00
feat(assets): include asset id in executed WebSocket message (#13862)
Some checks are pending
Detect Unreviewed Merge / detect (push) Waiting to run
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Waiting to run
Execution Tests / test (macos-latest) (push) Waiting to run
Execution Tests / test (ubuntu-latest) (push) Waiting to run
Execution Tests / test (windows-latest) (push) Waiting to run
Test server launches without errors / test (push) Waiting to run
Unit Tests / test (windows-2022) (push) Waiting to run
Unit Tests / test (macos-latest) (push) Waiting to run
Unit Tests / test (ubuntu-latest) (push) Waiting to run
Some checks are pending
Detect Unreviewed Merge / detect (push) Waiting to run
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Waiting to run
Execution Tests / test (macos-latest) (push) Waiting to run
Execution Tests / test (ubuntu-latest) (push) Waiting to run
Execution Tests / test (windows-latest) (push) Waiting to run
Test server launches without errors / test (push) Waiting to run
Unit Tests / test (windows-2022) (push) Waiting to run
Unit Tests / test (macos-latest) (push) Waiting to run
Unit Tests / test (ubuntu-latest) (push) Waiting to run
* feat(assets): enrich executed WS message with asset metadata
When --enable-assets is set, each file-type output entry in the
`executed` WebSocket message now includes id, name, asset_hash, size,
and mime_type — matching the shape already returned by /upload/image.
The enrichment lives in comfy_execution/asset_enrichment.py (no torch
dependency) and is called from both send sites in execution.py: freshly
executed nodes register the file inline via register_file_in_place;
cached node re-sends look up the existing AssetReference by file path
to avoid re-hashing. Errors are caught per-entry so a failure never
blocks the WS message from sending.
* fix(assets): inject only id in executed WS message per Asset Identity RFC
Per the Asset Identity RFC, the executed WebSocket payload should carry
id alone — hash is already encoded in the filename, and name/preview_url/
size belong behind GET /api/assets/{id} rather than being pushed eagerly.
Simplifies the DB lookup path: we only need ref.id, so the asset.hash
null-check is no longer required as a fallback trigger.
* fix(assets): reject path traversal when resolving output abs_path
Subfolder/filename were joined and absolutized without containment check,
so '..' segments or an absolute filename could escape the type's base
directory and register an unrelated on-disk file as an asset.
Add commonpath-based containment check; skip enrichment (warn, leave
entry unchanged) when the resolved path escapes base. Catches ValueError
from cross-drive paths on Windows.
* docs(assets): drop Asset Identity RFC reference from docstring
* docs(assets): trim docstring to what enrichment does, not what it doesn't
* test(assets): use real platform paths so containment check works on Windows
The previous test setup patched os.path.abspath to identity and used a
POSIX-style '/output' base, which collided with Windows path separators
in os.path.commonpath. Drop the abspath/join patches and use a real
tempdir-rooted base so the containment check runs against actual
platform paths.
* refactor(assets): enrich at output-processing time, not in the WS send path
Per review: enrichment lived inside the client_id-guarded send sites, so a
headless run (no websocket client) never registered assets at all, and
ui_outputs/history stored the un-enriched entries.
Now output_ui is enriched once, right after the node produces it and before
it is stored in ui_outputs — so registration happens regardless of connected
clients, and the asset id flows into history and the execution cache for
free. _send_cached_ui re-sends the stored (already-enriched) dict verbatim,
which lets the DB-lookup-by-path fallback be deleted: every enrichment is
now a fresh output, and register_file_in_place re-hashes on upsert so an
overwritten path can never carry a stale id.
This commit is contained in:
parent
e5b7140dcc
commit
ce200c0850
66
comfy_execution/asset_enrichment.py
Normal file
66
comfy_execution/asset_enrichment.py
Normal file
@ -0,0 +1,66 @@
|
||||
"""Enrich executed-node output entries with asset id."""
|
||||
import logging
|
||||
import os
|
||||
|
||||
|
||||
def enrich_output_with_assets(output_ui: dict) -> dict:
|
||||
"""Register file-type output entries as assets and inject their ``id``.
|
||||
|
||||
Runs at output-processing time, once per produced output, when
|
||||
--enable-assets is set. Returns a new dict; entries without a resolvable
|
||||
on-disk file path are left unchanged. Errors are caught per-entry so a
|
||||
failure never blocks execution or the other entries.
|
||||
"""
|
||||
from comfy.cli_args import args
|
||||
if not args.enable_assets:
|
||||
return output_ui
|
||||
|
||||
import folder_paths
|
||||
from app.assets.services.ingest import register_file_in_place, DependencyMissingError
|
||||
|
||||
enriched = {}
|
||||
for key, entries in output_ui.items():
|
||||
if not isinstance(entries, list):
|
||||
enriched[key] = entries
|
||||
continue
|
||||
new_entries = []
|
||||
for entry in entries:
|
||||
if not isinstance(entry, dict) or "filename" not in entry or "type" not in entry:
|
||||
new_entries.append(entry)
|
||||
continue
|
||||
try:
|
||||
base = folder_paths.get_directory_by_type(entry["type"])
|
||||
if base is None:
|
||||
new_entries.append(entry)
|
||||
continue
|
||||
base_abs = os.path.abspath(base)
|
||||
abs_path = os.path.abspath(os.path.join(base_abs, entry.get("subfolder") or "", entry["filename"]))
|
||||
try:
|
||||
if os.path.commonpath([base_abs, abs_path]) != base_abs:
|
||||
raise ValueError("escapes base")
|
||||
except ValueError:
|
||||
logging.warning("Asset enrichment skipped (path escapes base): %s", entry.get("filename"))
|
||||
new_entries.append(entry)
|
||||
continue
|
||||
if not os.path.isfile(abs_path):
|
||||
new_entries.append(entry)
|
||||
continue
|
||||
|
||||
# Register unconditionally: the file was just produced, and
|
||||
# register_file_in_place re-hashes so an overwritten path can
|
||||
# never carry a stale id.
|
||||
result = register_file_in_place(
|
||||
abs_path=abs_path,
|
||||
name=entry["filename"],
|
||||
tags=[entry["type"]],
|
||||
)
|
||||
|
||||
entry = dict(entry)
|
||||
entry["id"] = result.ref.id
|
||||
except DependencyMissingError:
|
||||
logging.warning("Asset enrichment skipped (blake3 not available): %s", entry.get("filename"))
|
||||
except Exception:
|
||||
logging.warning("Failed to enrich output entry with asset id: %s", entry.get("filename"), exc_info=True)
|
||||
new_entries.append(entry)
|
||||
enriched[key] = new_entries
|
||||
return enriched
|
||||
@ -40,6 +40,7 @@ from comfy_execution.graph_utils import GraphBuilder, is_link
|
||||
from comfy_execution.validation import validate_node_input
|
||||
from comfy_execution.progress import get_progress_state, reset_progress_state, add_progress_handler, WebUIProgressHandler
|
||||
from comfy_execution.utils import CurrentNodeContext
|
||||
from comfy_execution.asset_enrichment import enrich_output_with_assets
|
||||
from comfy_api.internal import _ComfyNodeInternal, _NodeOutputInternal, first_real_override, is_class, make_locked_method_func
|
||||
from comfy_api.latest import io, _io
|
||||
from comfy_execution.cache_provider import _has_cache_providers, _get_cache_providers, _logger as _cache_logger
|
||||
@ -418,6 +419,7 @@ def _is_intermediate_output(dynprompt, node_id):
|
||||
class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
|
||||
return getattr(class_def, 'HAS_INTERMEDIATE_OUTPUT', False)
|
||||
|
||||
|
||||
def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs):
|
||||
if server.client_id is None:
|
||||
return
|
||||
@ -552,6 +554,10 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
asyncio.create_task(await_completion())
|
||||
return (ExecutionResult.PENDING, None, None)
|
||||
if len(output_ui) > 0:
|
||||
# Enrich at output-processing time (not in the send path) so assets
|
||||
# are registered even when no client is connected, and the asset id
|
||||
# flows into ui_outputs and the cache alongside the raw entries.
|
||||
output_ui = enrich_output_with_assets(output_ui)
|
||||
ui_outputs[unique_id] = {
|
||||
"meta": {
|
||||
"node_id": unique_id,
|
||||
|
||||
205
tests-unit/execution_test/test_enrich_output.py
Normal file
205
tests-unit/execution_test/test_enrich_output.py
Normal file
@ -0,0 +1,205 @@
|
||||
"""Tests for enrich_output_with_assets in comfy_execution/asset_enrichment.py."""
|
||||
import os
|
||||
import types
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
def _make_args(enable_assets: bool):
|
||||
a = types.SimpleNamespace()
|
||||
a.enable_assets = enable_assets
|
||||
return a
|
||||
|
||||
|
||||
def _make_register_result(ref_id="ref-id-2"):
|
||||
result = MagicMock()
|
||||
result.ref.id = ref_id
|
||||
return result
|
||||
|
||||
|
||||
# Platform-appropriate absolute base. tempfile.gettempdir() returns C:\... on
|
||||
# Windows and /tmp on POSIX, so containment via commonpath behaves naturally.
|
||||
_DEFAULT_BASE = os.path.join(__import__("tempfile").gettempdir(), "asset-enrichment-test-base")
|
||||
|
||||
|
||||
def _mocked_modules(*, enable_assets=True, register_file_in_place=None, directory=_DEFAULT_BASE):
|
||||
return {
|
||||
"comfy.cli_args": MagicMock(args=_make_args(enable_assets)),
|
||||
"folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=directory)),
|
||||
"app.assets.services.ingest": MagicMock(
|
||||
register_file_in_place=register_file_in_place or MagicMock(return_value=_make_register_result()),
|
||||
DependencyMissingError=type("DependencyMissingError", (Exception,), {}),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _call(output_ui, *, enable_assets=True, file_exists=True, register_result=None, directory=_DEFAULT_BASE):
|
||||
register_mock = MagicMock(return_value=register_result or _make_register_result())
|
||||
mocked = _mocked_modules(
|
||||
enable_assets=enable_assets,
|
||||
register_file_in_place=register_mock,
|
||||
directory=directory,
|
||||
)
|
||||
|
||||
# Only os.path.isfile is patched — abspath/join must run natively so the
|
||||
# containment check sees real platform paths.
|
||||
with patch.dict("sys.modules", mocked), \
|
||||
patch("os.path.isfile", return_value=file_exists):
|
||||
import importlib
|
||||
import comfy_execution.asset_enrichment as mod
|
||||
importlib.reload(mod)
|
||||
return mod.enrich_output_with_assets(output_ui)
|
||||
|
||||
|
||||
class TestEnrichOutputWithAssets(unittest.TestCase):
|
||||
|
||||
def test_disabled_returns_unchanged(self):
|
||||
output = {"images": [{"filename": "a.png", "subfolder": "", "type": "output"}]}
|
||||
result = _call(output, enable_assets=False)
|
||||
self.assertNotIn("id", result["images"][0])
|
||||
|
||||
def test_non_list_value_passed_through(self):
|
||||
output = {"text": "hello"}
|
||||
result = _call(output)
|
||||
self.assertEqual(result["text"], "hello")
|
||||
|
||||
def test_entry_without_filename_unchanged(self):
|
||||
output = {"latent": [{"subfolder": "", "type": "output"}]}
|
||||
result = _call(output)
|
||||
self.assertNotIn("id", result["latent"][0])
|
||||
|
||||
def test_entry_without_type_unchanged(self):
|
||||
output = {"data": [{"filename": "a.png", "subfolder": ""}]}
|
||||
result = _call(output)
|
||||
self.assertNotIn("id", result["data"][0])
|
||||
|
||||
def test_file_not_on_disk_unchanged(self):
|
||||
output = {"images": [{"filename": "missing.png", "subfolder": "", "type": "output"}]}
|
||||
result = _call(output, file_exists=False)
|
||||
self.assertNotIn("id", result["images"][0])
|
||||
|
||||
def test_unknown_type_returns_none_directory_unchanged(self):
|
||||
output = {"images": [{"filename": "a.png", "subfolder": "", "type": "unknown"}]}
|
||||
result = _call(output, directory=None)
|
||||
self.assertNotIn("id", result["images"][0])
|
||||
|
||||
def test_register_injects_only_id(self):
|
||||
reg = _make_register_result(ref_id="inline-ref")
|
||||
output = {"images": [{"filename": "new.png", "subfolder": "", "type": "output"}]}
|
||||
result = _call(output, register_result=reg)
|
||||
img = result["images"][0]
|
||||
self.assertEqual(img["id"], "inline-ref")
|
||||
# Only id is injected — no asset_hash, name, preview_url, size
|
||||
self.assertNotIn("asset_hash", img)
|
||||
self.assertNotIn("name", img)
|
||||
self.assertNotIn("preview_url", img)
|
||||
self.assertNotIn("size", img)
|
||||
|
||||
def test_register_called_per_entry(self):
|
||||
register_mock = MagicMock(return_value=_make_register_result())
|
||||
mocked = _mocked_modules(register_file_in_place=register_mock)
|
||||
output = {
|
||||
"images": [
|
||||
{"filename": "a.png", "subfolder": "", "type": "output"},
|
||||
{"filename": "b.png", "subfolder": "", "type": "output"},
|
||||
]
|
||||
}
|
||||
|
||||
with patch.dict("sys.modules", mocked), \
|
||||
patch("os.path.isfile", return_value=True):
|
||||
import importlib
|
||||
import comfy_execution.asset_enrichment as mod
|
||||
importlib.reload(mod)
|
||||
mod.enrich_output_with_assets(output)
|
||||
|
||||
self.assertEqual(register_mock.call_count, 2)
|
||||
|
||||
def test_original_entry_not_mutated(self):
|
||||
orig = {"filename": "a.png", "subfolder": "", "type": "output"}
|
||||
output = {"images": [orig]}
|
||||
_call(output)
|
||||
self.assertNotIn("id", orig)
|
||||
|
||||
def test_enrichment_error_does_not_block_sibling_entries(self):
|
||||
call_count = [0]
|
||||
good_reg = _make_register_result(ref_id="good-ref")
|
||||
|
||||
def register_side_effect(abs_path, name, tags):
|
||||
call_count[0] += 1
|
||||
if call_count[0] == 1:
|
||||
raise RuntimeError("boom")
|
||||
return good_reg
|
||||
|
||||
mocked = _mocked_modules(register_file_in_place=register_side_effect)
|
||||
|
||||
output = {
|
||||
"images": [
|
||||
{"filename": "bad.png", "subfolder": "", "type": "output"},
|
||||
{"filename": "good.png", "subfolder": "", "type": "output"},
|
||||
]
|
||||
}
|
||||
|
||||
with patch.dict("sys.modules", mocked), \
|
||||
patch("os.path.isfile", return_value=True):
|
||||
import importlib
|
||||
import comfy_execution.asset_enrichment as mod
|
||||
importlib.reload(mod)
|
||||
result = mod.enrich_output_with_assets(output)
|
||||
|
||||
imgs = result["images"]
|
||||
self.assertNotIn("id", imgs[0])
|
||||
self.assertEqual(imgs[1]["id"], "good-ref")
|
||||
|
||||
def test_multiple_output_keys_all_enriched(self):
|
||||
output = {
|
||||
"images": [{"filename": "a.png", "subfolder": "", "type": "output"}],
|
||||
"videos": [{"filename": "b.mp4", "subfolder": "", "type": "output"}],
|
||||
}
|
||||
result = _call(output)
|
||||
self.assertIn("id", result["images"][0])
|
||||
self.assertIn("id", result["videos"][0])
|
||||
|
||||
def test_none_entry_in_list_unchanged(self):
|
||||
output = {"images": [None, {"filename": "a.png", "subfolder": "", "type": "output"}]}
|
||||
result = _call(output)
|
||||
self.assertIsNone(result["images"][0])
|
||||
self.assertIn("id", result["images"][1])
|
||||
|
||||
def test_path_traversal_subfolder_skipped(self):
|
||||
register_mock = MagicMock(return_value=_make_register_result())
|
||||
mocked = _mocked_modules(register_file_in_place=register_mock)
|
||||
|
||||
output = {"images": [{"filename": "passwd", "subfolder": "../../etc", "type": "output"}]}
|
||||
|
||||
# Do NOT patch os.path.abspath — real resolution is required for the containment check.
|
||||
with patch.dict("sys.modules", mocked), \
|
||||
patch("os.path.isfile", return_value=True):
|
||||
import importlib
|
||||
import comfy_execution.asset_enrichment as mod
|
||||
importlib.reload(mod)
|
||||
result = mod.enrich_output_with_assets(output)
|
||||
|
||||
self.assertNotIn("id", result["images"][0])
|
||||
register_mock.assert_not_called()
|
||||
|
||||
def test_absolute_filename_skipped(self):
|
||||
register_mock = MagicMock(return_value=_make_register_result())
|
||||
mocked = _mocked_modules(register_file_in_place=register_mock)
|
||||
|
||||
# Absolute filename — os.path.join discards earlier components when a later one is absolute.
|
||||
absolute_filename = os.path.abspath(os.sep + "etc" + os.sep + "passwd")
|
||||
output = {"images": [{"filename": absolute_filename, "subfolder": "", "type": "output"}]}
|
||||
|
||||
with patch.dict("sys.modules", mocked), \
|
||||
patch("os.path.isfile", return_value=True):
|
||||
import importlib
|
||||
import comfy_execution.asset_enrichment as mod
|
||||
importlib.reload(mod)
|
||||
result = mod.enrich_output_with_assets(output)
|
||||
|
||||
self.assertNotIn("id", result["images"][0])
|
||||
register_mock.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Loading…
Reference in New Issue
Block a user