Merge 257ae1f3c1 into 26515acd23

comfy_execution/asset_enrichment.py

@@ -0,0 +1,71 @@
+"""Enrich executed-node output entries with asset id."""
+import logging
+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.
+
+    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.
+                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
+
+                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 id: %s", entry.get("filename"), exc_info=True)
+            new_entries.append(entry)
+        enriched[key] = new_entries
+    return enriched

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 = []

tests-unit/execution_test/test_enrich_output.py

@@ -0,0 +1,177 @@
+"""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"):
+    ref = MagicMock()
+    ref.id = ref_id
+    return ref
+
+
+def _make_register_result(ref_id="ref-id-2"):
+    result = MagicMock()
+    result.ref.id = ref_id
+    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),
+        ),
+        "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_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")
+        # 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_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
+
+        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=None),
+            ),
+            "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()