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.
This commit is contained in:
Matt Miller 2026-05-12 18:05:34 -07:00
parent e35348aa53
commit cba7d83657
3 changed files with 266 additions and 2 deletions

View File

@ -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

View File

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

View File

@ -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()