From 2f88b5d719f5818b75f022f3fca13126c41ee983 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Tue, 12 May 2026 18:05:34 -0700 Subject: [PATCH 1/7] feat(assets): enrich executed WS message with asset metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- comfy_execution/asset_enrichment.py | 75 +++++++ execution.py | 9 +- .../execution_test/test_enrich_output.py | 184 ++++++++++++++++++ 3 files changed, 266 insertions(+), 2 deletions(-) create mode 100644 comfy_execution/asset_enrichment.py create mode 100644 tests-unit/execution_test/test_enrich_output.py diff --git a/comfy_execution/asset_enrichment.py b/comfy_execution/asset_enrichment.py new file mode 100644 index 000000000..44c5cb1f1 --- /dev/null +++ b/comfy_execution/asset_enrichment.py @@ -0,0 +1,75 @@ +"""Enrich executed-node output entries with asset metadata.""" +import logging +import os + + +def enrich_output_with_assets(output_ui: dict) -> dict: + """Inject asset metadata into file-type output entries 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 the WS message + from sending. + """ + 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 + from app.assets.database.queries.asset_reference import get_reference_by_file_path + from app.database.db import create_session + + 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 + abs_path = os.path.abspath(os.path.join(base, entry.get("subfolder", ""), entry["filename"])) + if not os.path.isfile(abs_path): + new_entries.append(entry) + continue + + # Try DB lookup first (cached node re-send); fall back to registering inline. + ref = asset = None + with create_session() as session: + db_ref = get_reference_by_file_path(session, abs_path) + if db_ref is not None and db_ref.asset is not None and db_ref.asset.hash is not None: + ref = db_ref + asset = db_ref.asset + + if ref is None: + result = register_file_in_place( + abs_path=abs_path, + name=entry["filename"], + tags=[entry["type"]], + ) + entry = dict(entry) + entry["id"] = result.ref.id + entry["name"] = result.ref.name + entry["asset_hash"] = result.asset.hash + entry["size"] = result.asset.size_bytes + entry["mime_type"] = result.asset.mime_type + else: + entry = dict(entry) + entry["id"] = ref.id + entry["name"] = ref.name + entry["asset_hash"] = asset.hash + entry["size"] = asset.size_bytes + entry["mime_type"] = asset.mime_type + 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 metadata: %s", entry.get("filename"), exc_info=True) + new_entries.append(entry) + enriched[key] = new_entries + return enriched diff --git a/execution.py b/execution.py index 4c7de2e84..73a454751 100644 --- a/execution.py +++ b/execution.py @@ -39,6 +39,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 @@ -417,11 +418,15 @@ 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 cached_ui = cached.ui or {} - server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, server.client_id) + output = cached_ui.get("output", None) + if output: + output = enrich_output_with_assets(output) + server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": output, "prompt_id": prompt_id }, server.client_id) if cached.ui is not None: ui_outputs[node_id] = cached.ui @@ -561,7 +566,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed, "output": output_ui } if server.client_id is not None: - server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": output_ui, "prompt_id": prompt_id }, server.client_id) + server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": enrich_output_with_assets(output_ui), "prompt_id": prompt_id }, server.client_id) if has_subgraph: cached_outputs = [] new_node_ids = [] diff --git a/tests-unit/execution_test/test_enrich_output.py b/tests-unit/execution_test/test_enrich_output.py new file mode 100644 index 000000000..6583143eb --- /dev/null +++ b/tests-unit/execution_test/test_enrich_output.py @@ -0,0 +1,184 @@ +"""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_db_ref(ref_id="ref-id-1", name="a.png", asset_hash="blake3:abc123", size=1024, mime="image/png"): + ref = MagicMock() + ref.id = ref_id + ref.name = name + ref.asset.hash = asset_hash + ref.asset.size_bytes = size + ref.asset.mime_type = mime + return ref + + +def _make_register_result(ref_id="ref-id-2", name="b.png", asset_hash="blake3:def456", size=2048, mime="image/png"): + result = MagicMock() + result.ref.id = ref_id + result.ref.name = name + result.asset.hash = asset_hash + result.asset.size_bytes = size + result.asset.mime_type = mime + return result + + +def _call(output_ui, *, enable_assets=True, file_exists=True, db_ref=None, register_result=None, directory="/output"): + fake_session_cm = MagicMock() + fake_session_cm.__enter__ = MagicMock(return_value=MagicMock()) + fake_session_cm.__exit__ = MagicMock(return_value=False) + + mocked_modules = { + "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=MagicMock(return_value=register_result or _make_register_result()), + DependencyMissingError=type("DependencyMissingError", (Exception,), {}), + ), + "app.assets.database.queries.asset_reference": MagicMock( + get_reference_by_file_path=MagicMock(return_value=db_ref or _make_db_ref()), + ), + "app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)), + } + + with patch.dict("sys.modules", mocked_modules), \ + patch("os.path.abspath", side_effect=lambda p: p), \ + patch("os.path.isfile", return_value=file_exists), \ + patch("os.path.join", side_effect=os.path.join): + 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_db_hit_injects_from_db(self): + db_ref = _make_db_ref(ref_id="db-ref", name="from-db.png", asset_hash="blake3:fromdb", size=512) + output = {"images": [{"filename": "a.png", "subfolder": "", "type": "output"}]} + result = _call(output, db_ref=db_ref) + img = result["images"][0] + self.assertEqual(img["id"], "db-ref") + self.assertEqual(img["asset_hash"], "blake3:fromdb") + self.assertEqual(img["size"], 512) + + def test_db_miss_falls_back_to_register(self): + no_hash_ref = _make_db_ref(asset_hash=None) + reg = _make_register_result(ref_id="inline-ref", asset_hash="blake3:inline", size=999) + output = {"images": [{"filename": "new.png", "subfolder": "", "type": "output"}]} + result = _call(output, db_ref=no_hash_ref, register_result=reg) + img = result["images"][0] + self.assertEqual(img["id"], "inline-ref") + self.assertEqual(img["asset_hash"], "blake3:inline") + self.assertEqual(img["size"], 999) + + 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", asset_hash="blake3:good") + no_hash_ref = _make_db_ref(asset_hash=None) + + def register_side_effect(abs_path, name, tags): + call_count[0] += 1 + if call_count[0] == 1: + raise RuntimeError("boom") + return good_reg + + fake_session_cm = MagicMock() + fake_session_cm.__enter__ = MagicMock(return_value=MagicMock()) + fake_session_cm.__exit__ = MagicMock(return_value=False) + + mocked_modules = { + "comfy.cli_args": MagicMock(args=_make_args(True)), + "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value="/output")), + "app.assets.services.ingest": MagicMock( + register_file_in_place=register_side_effect, + DependencyMissingError=type("DependencyMissingError", (Exception,), {}), + ), + "app.assets.database.queries.asset_reference": MagicMock( + get_reference_by_file_path=MagicMock(return_value=no_hash_ref), + ), + "app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)), + } + + output = { + "images": [ + {"filename": "bad.png", "subfolder": "", "type": "output"}, + {"filename": "good.png", "subfolder": "", "type": "output"}, + ] + } + + with patch.dict("sys.modules", mocked_modules), \ + patch("os.path.abspath", side_effect=lambda p: p), \ + patch("os.path.isfile", return_value=True), \ + patch("os.path.join", side_effect=os.path.join): + 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]) + + +if __name__ == "__main__": + unittest.main() From 29556e60986d4aa27ede0490ada3172f74b28eb5 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 13 May 2026 17:00:28 -0700 Subject: [PATCH 2/7] fix(assets): inject only id in executed WS message per Asset Identity RFC MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- comfy_execution/asset_enrichment.py | 42 +++++++++---------- .../execution_test/test_enrich_output.py | 39 +++++++---------- 2 files changed, 35 insertions(+), 46 deletions(-) diff --git a/comfy_execution/asset_enrichment.py b/comfy_execution/asset_enrichment.py index 44c5cb1f1..cc794fd79 100644 --- a/comfy_execution/asset_enrichment.py +++ b/comfy_execution/asset_enrichment.py @@ -1,14 +1,20 @@ -"""Enrich executed-node output entries with asset metadata.""" +"""Enrich executed-node output entries with asset id.""" import logging import os def enrich_output_with_assets(output_ui: dict) -> dict: - """Inject asset metadata into file-type output entries when --enable-assets is set. + """Inject asset id into file-type output entries when --enable-assets is set. + + Only ``id`` is added — per the Asset Identity RFC the WebSocket payload + carries just enough for the client to fetch the full asset via + GET /api/assets/{id}. hash, name, preview_url, and size are intentionally + omitted: hash is already encoded in the filename; the rest require an + explicit API call. 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 the WS message - from sending. + unchanged. Errors are caught per-entry so a failure never blocks the WS + message from sending. """ from comfy.cli_args import args if not args.enable_assets: @@ -40,36 +46,26 @@ def enrich_output_with_assets(output_ui: dict) -> dict: continue # Try DB lookup first (cached node re-send); fall back to registering inline. - ref = asset = None + asset_id = None with create_session() as session: db_ref = get_reference_by_file_path(session, abs_path) - if db_ref is not None and db_ref.asset is not None and db_ref.asset.hash is not None: - ref = db_ref - asset = db_ref.asset + if db_ref is not None: + asset_id = db_ref.id - if ref is None: + if asset_id is None: result = register_file_in_place( abs_path=abs_path, name=entry["filename"], tags=[entry["type"]], ) - entry = dict(entry) - entry["id"] = result.ref.id - entry["name"] = result.ref.name - entry["asset_hash"] = result.asset.hash - entry["size"] = result.asset.size_bytes - entry["mime_type"] = result.asset.mime_type - else: - entry = dict(entry) - entry["id"] = ref.id - entry["name"] = ref.name - entry["asset_hash"] = asset.hash - entry["size"] = asset.size_bytes - entry["mime_type"] = asset.mime_type + asset_id = result.ref.id + + entry = dict(entry) + entry["id"] = asset_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 metadata: %s", entry.get("filename"), exc_info=True) + 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 diff --git a/tests-unit/execution_test/test_enrich_output.py b/tests-unit/execution_test/test_enrich_output.py index 6583143eb..9cac44f5c 100644 --- a/tests-unit/execution_test/test_enrich_output.py +++ b/tests-unit/execution_test/test_enrich_output.py @@ -11,23 +11,15 @@ def _make_args(enable_assets: bool): return a -def _make_db_ref(ref_id="ref-id-1", name="a.png", asset_hash="blake3:abc123", size=1024, mime="image/png"): +def _make_db_ref(ref_id="ref-id-1"): ref = MagicMock() ref.id = ref_id - ref.name = name - ref.asset.hash = asset_hash - ref.asset.size_bytes = size - ref.asset.mime_type = mime return ref -def _make_register_result(ref_id="ref-id-2", name="b.png", asset_hash="blake3:def456", size=2048, mime="image/png"): +def _make_register_result(ref_id="ref-id-2"): result = MagicMock() result.ref.id = ref_id - result.ref.name = name - result.asset.hash = asset_hash - result.asset.size_bytes = size - result.asset.mime_type = mime return result @@ -44,7 +36,7 @@ def _call(output_ui, *, enable_assets=True, file_exists=True, db_ref=None, regis DependencyMissingError=type("DependencyMissingError", (Exception,), {}), ), "app.assets.database.queries.asset_reference": MagicMock( - get_reference_by_file_path=MagicMock(return_value=db_ref or _make_db_ref()), + get_reference_by_file_path=MagicMock(return_value=db_ref), ), "app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)), } @@ -91,24 +83,26 @@ class TestEnrichOutputWithAssets(unittest.TestCase): result = _call(output, directory=None) self.assertNotIn("id", result["images"][0]) - def test_db_hit_injects_from_db(self): - db_ref = _make_db_ref(ref_id="db-ref", name="from-db.png", asset_hash="blake3:fromdb", size=512) + def test_db_hit_injects_id(self): + db_ref = _make_db_ref(ref_id="db-ref") output = {"images": [{"filename": "a.png", "subfolder": "", "type": "output"}]} result = _call(output, db_ref=db_ref) img = result["images"][0] self.assertEqual(img["id"], "db-ref") - self.assertEqual(img["asset_hash"], "blake3:fromdb") - self.assertEqual(img["size"], 512) + # 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_db_miss_falls_back_to_register(self): - no_hash_ref = _make_db_ref(asset_hash=None) - reg = _make_register_result(ref_id="inline-ref", asset_hash="blake3:inline", size=999) + reg = _make_register_result(ref_id="inline-ref") output = {"images": [{"filename": "new.png", "subfolder": "", "type": "output"}]} - result = _call(output, db_ref=no_hash_ref, register_result=reg) + result = _call(output, db_ref=None, register_result=reg) img = result["images"][0] self.assertEqual(img["id"], "inline-ref") - self.assertEqual(img["asset_hash"], "blake3:inline") - self.assertEqual(img["size"], 999) + self.assertNotIn("asset_hash", img) + self.assertNotIn("name", img) def test_original_entry_not_mutated(self): orig = {"filename": "a.png", "subfolder": "", "type": "output"} @@ -118,8 +112,7 @@ class TestEnrichOutputWithAssets(unittest.TestCase): def test_enrichment_error_does_not_block_sibling_entries(self): call_count = [0] - good_reg = _make_register_result(ref_id="good-ref", asset_hash="blake3:good") - no_hash_ref = _make_db_ref(asset_hash=None) + good_reg = _make_register_result(ref_id="good-ref") def register_side_effect(abs_path, name, tags): call_count[0] += 1 @@ -139,7 +132,7 @@ class TestEnrichOutputWithAssets(unittest.TestCase): DependencyMissingError=type("DependencyMissingError", (Exception,), {}), ), "app.assets.database.queries.asset_reference": MagicMock( - get_reference_by_file_path=MagicMock(return_value=no_hash_ref), + get_reference_by_file_path=MagicMock(return_value=None), ), "app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)), } From f91e3416b6ca9b41ca3be3afb3343dc29be919c6 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 20 May 2026 10:57:40 -0700 Subject: [PATCH 3/7] 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. --- comfy_execution/asset_enrichment.py | 10 ++- .../execution_test/test_enrich_output.py | 64 +++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/comfy_execution/asset_enrichment.py b/comfy_execution/asset_enrichment.py index cc794fd79..637140bda 100644 --- a/comfy_execution/asset_enrichment.py +++ b/comfy_execution/asset_enrichment.py @@ -40,7 +40,15 @@ def enrich_output_with_assets(output_ui: dict) -> dict: if base is None: new_entries.append(entry) continue - abs_path = os.path.abspath(os.path.join(base, entry.get("subfolder", ""), entry["filename"])) + 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 diff --git a/tests-unit/execution_test/test_enrich_output.py b/tests-unit/execution_test/test_enrich_output.py index 9cac44f5c..2312207a5 100644 --- a/tests-unit/execution_test/test_enrich_output.py +++ b/tests-unit/execution_test/test_enrich_output.py @@ -172,6 +172,70 @@ class TestEnrichOutputWithAssets(unittest.TestCase): self.assertIsNone(result["images"][0]) self.assertIn("id", result["images"][1]) + def test_path_traversal_subfolder_skipped(self): + fake_session_cm = MagicMock() + fake_session_cm.__enter__ = MagicMock(return_value=MagicMock()) + fake_session_cm.__exit__ = MagicMock(return_value=False) + + register_mock = MagicMock(return_value=_make_register_result()) + mocked_modules = { + "comfy.cli_args": MagicMock(args=_make_args(True)), + "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value="/output")), + "app.assets.services.ingest": MagicMock( + register_file_in_place=register_mock, + DependencyMissingError=type("DependencyMissingError", (Exception,), {}), + ), + "app.assets.database.queries.asset_reference": MagicMock( + get_reference_by_file_path=MagicMock(return_value=None), + ), + "app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)), + } + + 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_modules), \ + 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): + fake_session_cm = MagicMock() + fake_session_cm.__enter__ = MagicMock(return_value=MagicMock()) + fake_session_cm.__exit__ = MagicMock(return_value=False) + + register_mock = MagicMock(return_value=_make_register_result()) + mocked_modules = { + "comfy.cli_args": MagicMock(args=_make_args(True)), + "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value="/output")), + "app.assets.services.ingest": MagicMock( + register_file_in_place=register_mock, + DependencyMissingError=type("DependencyMissingError", (Exception,), {}), + ), + "app.assets.database.queries.asset_reference": MagicMock( + get_reference_by_file_path=MagicMock(return_value=None), + ), + "app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)), + } + + # Absolute filename — os.path.join discards earlier components when a later one is absolute. + output = {"images": [{"filename": "/etc/passwd", "subfolder": "", "type": "output"}]} + + with patch.dict("sys.modules", mocked_modules), \ + 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() From 1bdd0d3ff20258a67206b520ec5e3ec7579b1894 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 20 May 2026 10:59:20 -0700 Subject: [PATCH 4/7] docs(assets): drop Asset Identity RFC reference from docstring --- comfy_execution/asset_enrichment.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/comfy_execution/asset_enrichment.py b/comfy_execution/asset_enrichment.py index 637140bda..1b45fd6c4 100644 --- a/comfy_execution/asset_enrichment.py +++ b/comfy_execution/asset_enrichment.py @@ -6,11 +6,10 @@ import os def enrich_output_with_assets(output_ui: dict) -> dict: """Inject asset id into file-type output entries when --enable-assets is set. - Only ``id`` is added — per the Asset Identity RFC the WebSocket payload - carries just enough for the client to fetch the full asset via - GET /api/assets/{id}. hash, name, preview_url, and size are intentionally - omitted: hash is already encoded in the filename; the rest require an - explicit API call. + Only ``id`` is added — the WebSocket payload carries just enough for the + client to fetch the full asset via GET /api/assets/{id}. hash, name, + preview_url, and size are intentionally omitted: hash is already encoded + in the filename; the rest require an explicit API call. 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 the WS From 4922875fe076c6a59a96e34e0ca6f8216b22f97e Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 20 May 2026 11:00:29 -0700 Subject: [PATCH 5/7] docs(assets): trim docstring to what enrichment does, not what it doesn't --- comfy_execution/asset_enrichment.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/comfy_execution/asset_enrichment.py b/comfy_execution/asset_enrichment.py index 1b45fd6c4..3a050a7bf 100644 --- a/comfy_execution/asset_enrichment.py +++ b/comfy_execution/asset_enrichment.py @@ -4,12 +4,7 @@ import os def enrich_output_with_assets(output_ui: dict) -> dict: - """Inject asset id into file-type output entries when --enable-assets is set. - - Only ``id`` is added — the WebSocket payload carries just enough for the - client to fetch the full asset via GET /api/assets/{id}. hash, name, - preview_url, and size are intentionally omitted: hash is already encoded - in the filename; the rest require an explicit API call. + """Inject asset ``id`` into file-type output entries 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 the WS From 2eb15b62373d156cf50626bddd0dfc0e401c2076 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 20 May 2026 11:11:26 -0700 Subject: [PATCH 6/7] spec(assets): document optional id field on ResultItem ExecutedWsMessage's per-entry shape (ResultItem) now optionally carries the asset reference UUID when --enable-assets is set. --- openapi.yaml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/openapi.yaml b/openapi.yaml index 2658b9b86..a086a1a1b 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -6569,6 +6569,13 @@ components: enum: [input, output, temp] display_name: type: string + id: + type: string + format: uuid + description: | + Asset reference UUID. Present only when the server is started with + `--enable-assets` and the file resolves to a registered asset. + Fetch the full asset via `GET /api/assets/{id}`. NodeOutputs: type: object From 1ebbee4522dfcf1cd0c2b3360fb26069aa4611fb Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 20 May 2026 11:17:59 -0700 Subject: [PATCH 7/7] 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. --- .../execution_test/test_enrich_output.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/tests-unit/execution_test/test_enrich_output.py b/tests-unit/execution_test/test_enrich_output.py index 2312207a5..5580113ba 100644 --- a/tests-unit/execution_test/test_enrich_output.py +++ b/tests-unit/execution_test/test_enrich_output.py @@ -23,7 +23,12 @@ def _make_register_result(ref_id="ref-id-2"): return result -def _call(output_ui, *, enable_assets=True, file_exists=True, db_ref=None, register_result=None, directory="/output"): +# 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 _call(output_ui, *, enable_assets=True, file_exists=True, db_ref=None, register_result=None, directory=_DEFAULT_BASE): fake_session_cm = MagicMock() fake_session_cm.__enter__ = MagicMock(return_value=MagicMock()) fake_session_cm.__exit__ = MagicMock(return_value=False) @@ -41,10 +46,10 @@ def _call(output_ui, *, enable_assets=True, file_exists=True, db_ref=None, regis "app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)), } + # 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_modules), \ - patch("os.path.abspath", side_effect=lambda p: p), \ - patch("os.path.isfile", return_value=file_exists), \ - patch("os.path.join", side_effect=os.path.join): + patch("os.path.isfile", return_value=file_exists): import importlib import comfy_execution.asset_enrichment as mod importlib.reload(mod) @@ -126,7 +131,7 @@ class TestEnrichOutputWithAssets(unittest.TestCase): mocked_modules = { "comfy.cli_args": MagicMock(args=_make_args(True)), - "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value="/output")), + "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=_DEFAULT_BASE)), "app.assets.services.ingest": MagicMock( register_file_in_place=register_side_effect, DependencyMissingError=type("DependencyMissingError", (Exception,), {}), @@ -145,9 +150,7 @@ class TestEnrichOutputWithAssets(unittest.TestCase): } with patch.dict("sys.modules", mocked_modules), \ - patch("os.path.abspath", side_effect=lambda p: p), \ - patch("os.path.isfile", return_value=True), \ - patch("os.path.join", side_effect=os.path.join): + patch("os.path.isfile", return_value=True): import importlib import comfy_execution.asset_enrichment as mod importlib.reload(mod) @@ -180,7 +183,7 @@ class TestEnrichOutputWithAssets(unittest.TestCase): register_mock = MagicMock(return_value=_make_register_result()) mocked_modules = { "comfy.cli_args": MagicMock(args=_make_args(True)), - "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value="/output")), + "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=_DEFAULT_BASE)), "app.assets.services.ingest": MagicMock( register_file_in_place=register_mock, DependencyMissingError=type("DependencyMissingError", (Exception,), {}), @@ -212,7 +215,7 @@ class TestEnrichOutputWithAssets(unittest.TestCase): register_mock = MagicMock(return_value=_make_register_result()) mocked_modules = { "comfy.cli_args": MagicMock(args=_make_args(True)), - "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value="/output")), + "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=_DEFAULT_BASE)), "app.assets.services.ingest": MagicMock( register_file_in_place=register_mock, DependencyMissingError=type("DependencyMissingError", (Exception,), {}), @@ -224,7 +227,8 @@ class TestEnrichOutputWithAssets(unittest.TestCase): } # Absolute filename — os.path.join discards earlier components when a later one is absolute. - output = {"images": [{"filename": "/etc/passwd", "subfolder": "", "type": "output"}]} + 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_modules), \ patch("os.path.isfile", return_value=True):