From 95172b2c3d4935e737815e07425411d6b5442e5f Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Tue, 9 Jun 2026 21:59:42 -0700 Subject: [PATCH] refactor(assets): enrich at output-processing time, not in the WS send path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- comfy_execution/asset_enrichment.py | 35 ++--- execution.py | 11 +- .../execution_test/test_enrich_output.py | 124 ++++++------------ 3 files changed, 62 insertions(+), 108 deletions(-) diff --git a/comfy_execution/asset_enrichment.py b/comfy_execution/asset_enrichment.py index 3a050a7bf..38e9496a8 100644 --- a/comfy_execution/asset_enrichment.py +++ b/comfy_execution/asset_enrichment.py @@ -4,11 +4,12 @@ import os def enrich_output_with_assets(output_ui: dict) -> dict: - """Inject asset ``id`` into file-type output entries when --enable-assets is set. + """Register file-type output entries as assets and inject their ``id``. - 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. + 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: @@ -16,8 +17,6 @@ def enrich_output_with_assets(output_ui: dict) -> dict: 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(): @@ -47,23 +46,17 @@ def enrich_output_with_assets(output_ui: dict) -> dict: new_entries.append(entry) continue - # Try DB lookup first (cached node re-send); fall back to registering inline. - asset_id = None - with create_session() as session: - db_ref = get_reference_by_file_path(session, abs_path) - if db_ref is not None: - asset_id = db_ref.id - - if asset_id is None: - result = register_file_in_place( - abs_path=abs_path, - name=entry["filename"], - tags=[entry["type"]], - ) - asset_id = result.ref.id + # 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"] = asset_id + entry["id"] = result.ref.id except DependencyMissingError: logging.warning("Asset enrichment skipped (blake3 not available): %s", entry.get("filename")) except Exception: diff --git a/execution.py b/execution.py index f6934ea6a..e6c6f39d6 100644 --- a/execution.py +++ b/execution.py @@ -424,10 +424,7 @@ def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outp if server.client_id is None: return cached_ui = cached.ui or {} - 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) + 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) if cached.ui is not None: ui_outputs[node_id] = cached.ui @@ -557,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, @@ -567,7 +568,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": enrich_output_with_assets(output_ui), "prompt_id": prompt_id }, server.client_id) + server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": 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 index 5580113ba..61490c49e 100644 --- a/tests-unit/execution_test/test_enrich_output.py +++ b/tests-unit/execution_test/test_enrich_output.py @@ -11,12 +11,6 @@ def _make_args(enable_assets: bool): return a -def _make_db_ref(ref_id="ref-id-1"): - ref = MagicMock() - ref.id = ref_id - return ref - - def _make_register_result(ref_id="ref-id-2"): result = MagicMock() result.ref.id = ref_id @@ -28,27 +22,28 @@ def _make_register_result(ref_id="ref-id-2"): _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) - - mocked_modules = { +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=MagicMock(return_value=register_result or _make_register_result()), + register_file_in_place=register_file_in_place or MagicMock(return_value=_make_register_result()), DependencyMissingError=type("DependencyMissingError", (Exception,), {}), ), - "app.assets.database.queries.asset_reference": MagicMock( - get_reference_by_file_path=MagicMock(return_value=db_ref), - ), - "app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)), } + +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_modules), \ + with patch.dict("sys.modules", mocked), \ patch("os.path.isfile", return_value=file_exists): import importlib import comfy_execution.asset_enrichment as mod @@ -88,26 +83,36 @@ class TestEnrichOutputWithAssets(unittest.TestCase): result = _call(output, directory=None) self.assertNotIn("id", result["images"][0]) - 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) + 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"], "db-ref") + 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_db_miss_falls_back_to_register(self): - reg = _make_register_result(ref_id="inline-ref") - output = {"images": [{"filename": "new.png", "subfolder": "", "type": "output"}]} - result = _call(output, db_ref=None, register_result=reg) - img = result["images"][0] - self.assertEqual(img["id"], "inline-ref") - self.assertNotIn("asset_hash", img) - self.assertNotIn("name", 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"} @@ -125,22 +130,7 @@ class TestEnrichOutputWithAssets(unittest.TestCase): 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=_DEFAULT_BASE)), - "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=None), - ), - "app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)), - } + mocked = _mocked_modules(register_file_in_place=register_side_effect) output = { "images": [ @@ -149,7 +139,7 @@ class TestEnrichOutputWithAssets(unittest.TestCase): ] } - with patch.dict("sys.modules", mocked_modules), \ + with patch.dict("sys.modules", mocked), \ patch("os.path.isfile", return_value=True): import importlib import comfy_execution.asset_enrichment as mod @@ -176,28 +166,13 @@ class TestEnrichOutputWithAssets(unittest.TestCase): 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=_DEFAULT_BASE)), - "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)), - } + 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_modules), \ + with patch.dict("sys.modules", mocked), \ patch("os.path.isfile", return_value=True): import importlib import comfy_execution.asset_enrichment as mod @@ -208,29 +183,14 @@ class TestEnrichOutputWithAssets(unittest.TestCase): 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=_DEFAULT_BASE)), - "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)), - } + 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_modules), \ + with patch.dict("sys.modules", mocked), \ patch("os.path.isfile", return_value=True): import importlib import comfy_execution.asset_enrichment as mod