From 283ccd884b057992bd6c6465bc2a9cd48c9bd5e8 Mon Sep 17 00:00:00 2001 From: magicwang1111 Date: Fri, 17 Jul 2026 11:41:12 +0800 Subject: [PATCH] =?UTF-8?q?=E6=8C=81=E4=B9=85=E5=8C=96=E5=8E=86=E5=8F=B2?= =?UTF-8?q?=E4=B8=AD=E7=9A=84=E4=B8=B4=E6=97=B6=E5=AA=92=E4=BD=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 2 + tests/test_history_assets.py | 106 +++++++++++++++++++++++++++++++++++ utils/access_control.py | 19 +++++++ utils/history_assets.py | 81 ++++++++++++++++++++++++++ 4 files changed, 208 insertions(+) create mode 100644 tests/test_history_assets.py create mode 100644 utils/history_assets.py diff --git a/README.md b/README.md index 99b5789..7be258d 100644 --- a/README.md +++ b/README.md @@ -50,6 +50,8 @@ Edit `config.json` before starting ComfyUI. Completed generation history is persisted in `history.sqlite3` and restored when ComfyUI restarts. The history database keeps the same 10,000-item limit as ComfyUI's in-memory history. +Temporary preview images, videos, and audio referenced by completed history are linked into the user's persistent output folder before the history entry is saved. This keeps history thumbnails and media available after ComfyUI clears its temp folder during a restart. + ComfyUI's optional Assets panel is separate from queue history and requires starting ComfyUI with `--enable-assets`. ## API Access diff --git a/tests/test_history_assets.py b/tests/test_history_assets.py new file mode 100644 index 0000000..aea4c06 --- /dev/null +++ b/tests/test_history_assets.py @@ -0,0 +1,106 @@ +import importlib.util +import os +import shutil +import tempfile +import unittest +from pathlib import Path + + +module_path = Path(__file__).parents[1] / "utils" / "history_assets.py" +spec = importlib.util.spec_from_file_location("history_assets", module_path) +history_assets = importlib.util.module_from_spec(spec) +spec.loader.exec_module(history_assets) +persist_temp_assets = history_assets.persist_temp_assets + + +class HistoryAssetsTests(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.root = Path(self.temp_dir.name) + self.temp_root = self.root / "temp" + self.output_root = self.root / "output" + self.source_dir = self.temp_root / "alice" / "2026-07-17" / "alice" + self.source_dir.mkdir(parents=True) + + def tearDown(self): + self.temp_dir.cleanup() + + def test_temp_reference_is_made_durable_and_rewritten(self): + source = self.source_dir / "preview.png" + source.write_bytes(b"preview-data") + item = { + "outputs": { + "1": { + "images": [ + { + "filename": "preview.png", + "subfolder": "2026-07-17/alice", + "type": "temp", + } + ] + } + } + } + + count = persist_temp_assets( + item, + str(self.temp_root), + str(self.output_root), + "alice", + "2026-07-17/alice/history_assets/prompt-1", + ) + + reference = item["outputs"]["1"]["images"][0] + destination = self.output_root / reference["subfolder"] / reference["filename"] + self.assertEqual(1, count) + self.assertEqual("output", reference["type"]) + self.assertTrue(destination.is_file()) + self.assertEqual(b"preview-data", destination.read_bytes()) + + shutil.rmtree(self.temp_root) + self.assertEqual(b"preview-data", destination.read_bytes()) + + def test_missing_temp_file_keeps_original_reference(self): + reference = { + "filename": "missing.mp4", + "subfolder": "2026-07-17/alice", + "type": "temp", + } + item = {"outputs": {"1": {"videos": [reference]}}} + + count = persist_temp_assets( + item, + str(self.temp_root), + str(self.output_root), + "alice", + "2026-07-17/alice/history_assets/prompt-2", + ) + + self.assertEqual(0, count) + self.assertEqual("temp", reference["type"]) + + def test_unowned_base_temp_path_is_not_copied(self): + foreign_dir = self.temp_root / "2026-07-17" / "bob" + foreign_dir.mkdir(parents=True) + (foreign_dir / "foreign.png").write_bytes(b"foreign") + reference = { + "filename": "foreign.png", + "subfolder": "2026-07-17/bob", + "type": "temp", + } + item = {"outputs": {"1": {"images": [reference]}}} + + count = persist_temp_assets( + item, + str(self.temp_root), + str(self.output_root), + "alice", + "2026-07-17/alice/history_assets/prompt-3", + ) + + self.assertEqual(0, count) + self.assertEqual("temp", reference["type"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/access_control.py b/utils/access_control.py index 7a8d220..908c3fa 100644 --- a/utils/access_control.py +++ b/utils/access_control.py @@ -14,6 +14,7 @@ from server import PromptServer from execution import PromptQueue, MAXIMUM_HISTORY_SIZE from .users_db import UsersDB +from .history_assets import persist_temp_assets from .history_store import HistoryStore @@ -380,9 +381,27 @@ class AccessControl: "user_id": user_id, } self.__prompt_queue.history[prompt[1]].update(history_result) + self._persist_history_assets(prompt[1], user_id) self._save_history_item(prompt[1]) self.server.queue_updated() + def _persist_history_assets(self, prompt_id: str, user_id: str) -> None: + if not user_id: + return + try: + destination_subfolder = ( + f"{self._get_user_path_prefix(user_id)}/history_assets/{prompt_id}" + ) + persist_temp_assets( + self.__prompt_queue.history[prompt_id], + self.__get_temp_directory(), + self.__get_output_directory(), + self._get_user_slug(user_id), + destination_subfolder, + ) + except Exception: + logger.exception("Failed to persist temporary assets for %s", prompt_id) + def _save_history_item(self, prompt_id: str) -> None: if not self.__history_store: return diff --git a/utils/history_assets.py b/utils/history_assets.py new file mode 100644 index 0000000..fb74d10 --- /dev/null +++ b/utils/history_assets.py @@ -0,0 +1,81 @@ +import os +import shutil + + +def _path_parts(path: str): + normalized = str(path or "").replace("\\", "/").strip("/") + parts = [part for part in normalized.split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return None + return parts + + +def _iter_temp_references(value): + if isinstance(value, dict): + if value.get("type") == "temp" and value.get("filename"): + yield value + for nested in value.values(): + yield from _iter_temp_references(nested) + elif isinstance(value, list): + for nested in value: + yield from _iter_temp_references(nested) + + +def persist_temp_assets( + history_item: dict, + temp_root: str, + output_root: str, + owner_slug: str, + destination_subfolder: str, +) -> int: + """Persist history-only temp references onto durable output paths.""" + owner_parts = _path_parts(owner_slug) + destination_parts = _path_parts(destination_subfolder) + if not owner_parts or len(owner_parts) != 1 or not destination_parts: + return 0 + + owner_slug = owner_parts[0] + destination_dir = os.path.join(output_root, *destination_parts) + persisted = 0 + + for reference in _iter_temp_references(history_item): + filename = str(reference.get("filename") or "") + if not filename or filename != os.path.basename(filename): + continue + + subfolder_parts = _path_parts(reference.get("subfolder")) + if not subfolder_parts: + continue + + owned_subfolder = ( + subfolder_parts[0] == owner_slug + or (len(subfolder_parts) >= 2 and subfolder_parts[1] == owner_slug) + ) + source_candidates = [ + os.path.join(temp_root, owner_slug, *subfolder_parts, filename) + ] + if owned_subfolder: + source_candidates.append( + os.path.join(temp_root, *subfolder_parts, filename) + ) + + source = next( + (candidate for candidate in source_candidates if os.path.isfile(candidate)), + None, + ) + if source is None: + continue + + os.makedirs(destination_dir, exist_ok=True) + destination = os.path.join(destination_dir, filename) + if not os.path.exists(destination): + try: + os.link(source, destination) + except OSError: + shutil.copy2(source, destination) + + reference["type"] = "output" + reference["subfolder"] = "/".join(destination_parts) + persisted += 1 + + return persisted