From 257ae1f3c12467657835686874023d4f1cb58390 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 13 May 2026 17:00:28 -0700 Subject: [PATCH] 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)), }