ComfyUI/comfy_execution/asset_enrichment.py
Matt Miller 257ae1f3c1
Some checks are pending
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
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.
2026-05-13 17:00:28 -07:00

72 lines
2.9 KiB
Python

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