diff --git a/app/assets/scanner.py b/app/assets/scanner.py index 42c4c1e9d26..0db7588de19 100644 --- a/app/assets/scanner.py +++ b/app/assets/scanner.py @@ -57,10 +57,11 @@ class _AssetAccumulator(TypedDict): refs: list[_RefInfo] +# Temp is deliberately absent: it is wiped before every scan, so walking it finds nothing. RootType = Literal["models", "input", "output"] -def get_prefixes_for_root(root: RootType) -> list[str]: +def get_scan_prefixes_for_root(root: RootType) -> list[str]: if root == "models": bases: list[str] = [] for _bucket, paths, _exts in get_comfy_models_folders(): @@ -73,10 +74,15 @@ def get_prefixes_for_root(root: RootType) -> list[str]: return [] -def get_all_known_prefixes() -> list[str]: - """Get all known asset prefixes across all root types.""" - all_roots: tuple[RootType, ...] = ("models", "input", "output") - return [p for root in all_roots for p in get_prefixes_for_root(root)] +def get_owned_prefixes() -> list[str]: + """Every directory an asset may live in; references outside these are marked missing.""" + scan_roots: tuple[RootType, ...] = ("models", "input", "output") + prefixes = [p for root in scan_roots for p in get_scan_prefixes_for_root(root)] + return prefixes + get_temp_prefixes() + + +def get_temp_prefixes() -> list[str]: + return [os.path.abspath(folder_paths.get_temp_directory())] def collect_models_files() -> list[str]: @@ -107,7 +113,21 @@ def sync_references_with_filesystem( collect_existing_paths: bool = False, update_missing_tags: bool = False, ) -> set[str] | None: - """Reconcile asset references with filesystem for a root. + return sync_prefixes_with_filesystem( + session, + get_scan_prefixes_for_root(root), + collect_existing_paths=collect_existing_paths, + update_missing_tags=update_missing_tags, + ) + + +def sync_prefixes_with_filesystem( + session, + prefixes: list[str], + collect_existing_paths: bool = False, + update_missing_tags: bool = False, +) -> set[str] | None: + """Reconcile asset references with filesystem under the given prefixes. - Toggle needs_verify per reference using mtime/size stat check - For hashed assets with at least one stat-unchanged ref: delete stale missing refs @@ -117,14 +137,13 @@ def sync_references_with_filesystem( Args: session: Database session - root: Root type to scan + prefixes: Absolute directory prefixes whose references to reconcile collect_existing_paths: If True, return set of surviving file paths update_missing_tags: If True, update 'missing' tags based on file status Returns: Set of surviving absolute paths if collect_existing_paths=True, else None """ - prefixes = get_prefixes_for_root(root) if not prefixes: return set() if collect_existing_paths else None @@ -251,6 +270,16 @@ def sync_root_safely(root: RootType) -> set[str]: return set() +def sync_temp_references_safely() -> None: + """Retire temp references whose file is gone; temp is never scanned, so nothing else stats them.""" + try: + with create_session() as sess: + sync_prefixes_with_filesystem(sess, get_temp_prefixes()) + sess.commit() + except Exception as e: + logging.exception("temp reference sync failed: %s", e) + + def mark_missing_outside_prefixes_safely(prefixes: list[str]) -> int: """Mark references as missing when outside the given prefixes. @@ -384,7 +413,7 @@ def get_unenriched_assets_for_roots( """ prefixes: list[str] = [] for root in roots: - prefixes.extend(get_prefixes_for_root(root)) + prefixes.extend(get_scan_prefixes_for_root(root)) if not prefixes: return [] diff --git a/app/assets/seeder.py b/app/assets/seeder.py index 2262928e599..134fc98a844 100644 --- a/app/assets/seeder.py +++ b/app/assets/seeder.py @@ -15,12 +15,13 @@ build_asset_specs, collect_paths_for_roots, enrich_assets_batch, - get_all_known_prefixes, - get_prefixes_for_root, + get_owned_prefixes, + get_scan_prefixes_for_root, get_unenriched_assets_for_roots, insert_asset_specs, mark_missing_outside_prefixes_safely, sync_root_safely, + sync_temp_references_safely, ) from app.database.db import dependencies_available @@ -413,7 +414,7 @@ def mark_missing_outside_prefixes(self) -> int: ) return 0 - all_prefixes = get_all_known_prefixes() + all_prefixes = get_owned_prefixes() marked = mark_missing_outside_prefixes_safely(all_prefixes) if marked > 0: logging.info("Marked %d references as missing", marked) @@ -523,7 +524,7 @@ def _log_scan_config(self, roots: tuple[RootType, ...]) -> None: os.path.abspath(folder_paths.models_dir), ) else: - prefixes = get_prefixes_for_root(root) + prefixes = get_scan_prefixes_for_root(root) if prefixes: logging.info("Asset scan [%s] directories: %s", root, prefixes) @@ -548,10 +549,11 @@ def _run_scan(self) -> None: return if self._prune_first: - all_prefixes = get_all_known_prefixes() + all_prefixes = get_owned_prefixes() marked = mark_missing_outside_prefixes_safely(all_prefixes) if marked > 0: logging.info("Marked %d refs as missing before scan", marked) + sync_temp_references_safely() if self._check_pause_and_cancel(): logging.info("Asset scan cancelled after pruning phase") diff --git a/tests-unit/assets_test/test_sync_references.py b/tests-unit/assets_test/test_sync_references.py index 2e85076e021..39c19ed8a41 100644 --- a/tests-unit/assets_test/test_sync_references.py +++ b/tests-unit/assets_test/test_sync_references.py @@ -150,7 +150,7 @@ def test_needs_verify_toggling(session, temp_dir, case): ) session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): sync_references_with_filesystem(session, "models") session.commit() @@ -185,7 +185,7 @@ def test_is_missing_flag(session, temp_dir, case): _make_asset(session, "a1", fp, "r1", asset_hash="blake3:abc", mtime_ns=mtime) session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): sync_references_with_filesystem(session, "models") session.commit() @@ -200,7 +200,7 @@ def test_seed_asset_all_missing_deletes_asset(session, temp_dir): _make_asset(session, "seed1", fp, "r1", asset_hash=None, mtime_ns=999) session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): sync_references_with_filesystem(session, "models") session.commit() @@ -215,7 +215,7 @@ def test_seed_asset_some_exist_returns_survivors(session, temp_dir): _make_asset(session, "seed1", fp, "r1", asset_hash=None, mtime_ns=mtime) session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): survivors = sync_references_with_filesystem( session, "models", collect_existing_paths=True, ) @@ -240,7 +240,7 @@ def test_hashed_asset_prunes_missing_refs_when_one_is_ok(session, temp_dir): session.add(ref_gone) session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): sync_references_with_filesystem(session, "models") session.commit() @@ -255,7 +255,7 @@ def test_hashed_asset_all_missing_keeps_refs(session, temp_dir): _make_asset(session, "h1", fp, "r1", asset_hash="blake3:aaa", mtime_ns=999) session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): sync_references_with_filesystem(session, "models") session.commit() @@ -272,7 +272,7 @@ def test_missing_tag_added_when_all_refs_gone(session, temp_dir): _make_asset(session, "h1", fp, "r1", asset_hash="blake3:aaa", mtime_ns=999) session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): sync_references_with_filesystem( session, "models", update_missing_tags=True, ) @@ -295,7 +295,7 @@ def test_missing_tag_removed_when_ref_ok(session, temp_dir): )) session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): sync_references_with_filesystem( session, "models", update_missing_tags=True, ) @@ -313,7 +313,7 @@ def test_missing_tags_not_touched_when_flag_false(session, temp_dir): _make_asset(session, "h1", fp, "r1", asset_hash="blake3:aaa", mtime_ns=999) session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): sync_references_with_filesystem( session, "models", update_missing_tags=False, ) @@ -329,7 +329,7 @@ def test_returns_none_when_collect_false(session, temp_dir): _make_asset(session, "a1", fp, "r1", asset_hash="blake3:abc", mtime_ns=mtime) session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): result = sync_references_with_filesystem( session, "models", collect_existing_paths=False, ) @@ -338,7 +338,7 @@ def test_returns_none_when_collect_false(session, temp_dir): def test_returns_empty_set_for_no_prefixes(session): - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[]): result = sync_references_with_filesystem( session, "models", collect_existing_paths=True, ) @@ -348,7 +348,7 @@ def test_returns_empty_set_for_no_prefixes(session): def test_no_references_is_noop(session, temp_dir): """No crash and no side effects when there are no references.""" - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): survivors = sync_references_with_filesystem( session, "models", collect_existing_paths=True, ) @@ -388,7 +388,7 @@ def test_sync_does_not_resurrect_soft_deleted_ref(session, temp_dir): _soft_delete_ref(session, "r1") session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): sync_references_with_filesystem(session, "models") session.commit() @@ -472,7 +472,7 @@ def test_sync_ignores_soft_deleted_seed_asset(session, temp_dir): _soft_delete_ref(session, "r1") session.commit() - with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]): + with patch("app.assets.scanner.get_scan_prefixes_for_root", return_value=[str(temp_dir)]): sync_references_with_filesystem(session, "models") session.commit() diff --git a/tests-unit/assets_test/test_temp_assets.py b/tests-unit/assets_test/test_temp_assets.py new file mode 100644 index 00000000000..8a5f77d5232 --- /dev/null +++ b/tests-unit/assets_test/test_temp_assets.py @@ -0,0 +1,180 @@ + +import os +import tempfile +from pathlib import Path +from unittest.mock import patch + +import pytest +from sqlalchemy import create_engine +from sqlalchemy.orm import Session + +from app.assets.database.models import Asset, AssetReference, Base +from app.assets.database.queries.asset_reference import ( + mark_references_missing_outside_prefixes, +) +from app.assets.scanner import ( + collect_paths_for_roots, + get_owned_prefixes, + get_temp_prefixes, + sync_prefixes_with_filesystem, +) +from app.assets.services.file_utils import get_mtime_ns + + +@pytest.fixture(autouse=True) +def autoclean_unit_test_assets(): + """Override parent autouse fixture - temp asset tests don't need server cleanup.""" + yield + + +@pytest.fixture +def session(): + engine = create_engine("sqlite:///:memory:") + Base.metadata.create_all(engine) + with Session(engine) as sess: + yield sess + + +@pytest.fixture +def comfy_dirs(): + with tempfile.TemporaryDirectory() as base: + dirs = { + name: Path(base) / name + for name in ("models", "input", "output", "temp", "elsewhere") + } + for d in dirs.values(): + d.mkdir() + with ( + patch("folder_paths.get_input_directory", return_value=str(dirs["input"])), + patch("folder_paths.get_output_directory", return_value=str(dirs["output"])), + patch("folder_paths.get_temp_directory", return_value=str(dirs["temp"])), + patch( + "app.assets.scanner.get_comfy_models_folders", + return_value=[("checkpoints", [str(dirs["models"])], set())], + ), + ): + yield dirs + + +def _write(directory: Path, name: str) -> str: + p = directory / name + p.write_bytes(b"\x00" * 100) + return str(p) + + +def _register( + session: Session, + file_path: str, + ref_id: str, + *, + mtime_ns: int, + asset_hash: str | None = "", +) -> None: + if asset_hash == "": + asset_hash = f"blake3:{ref_id}" + session.add(Asset(id=f"asset-{ref_id}", hash=asset_hash, size_bytes=100)) + session.flush() + session.add( + AssetReference( + id=ref_id, + asset_id=f"asset-{ref_id}", + name=os.path.basename(file_path), + owner_id="", + file_path=file_path, + mtime_ns=mtime_ns, + ) + ) + session.flush() + + +def _mtime(path: str) -> int: + return get_mtime_ns(os.stat(path, follow_symlinks=True)) + + +def test_owned_prefixes_include_temp(comfy_dirs): + owned = get_owned_prefixes() + assert str(comfy_dirs["temp"]) in owned, ( + "temp must be owned, or the prune disowns assets whose files are present" + ) + for name in ("models", "input", "output"): + assert str(comfy_dirs[name]) in owned, f"{name} must stay owned" + + +def test_discovery_does_not_walk_temp(comfy_dirs): + temp_file = _write(comfy_dirs["temp"], "preview.png") + output_file = _write(comfy_dirs["output"], "render.png") + + with patch("app.assets.scanner.collect_models_files", return_value=[]): + paths = collect_paths_for_roots(("models", "input", "output")) + + assert output_file in paths, "scan roots must still be walked" + assert temp_file not in paths, ( + "temp is wiped before every scan, so walking it only ever finds nothing" + ) + + +def test_prune_keeps_live_temp_reference(session, comfy_dirs): + temp_file = _write(comfy_dirs["temp"], "preview.png") + stray_file = _write(comfy_dirs["elsewhere"], "stray.png") + _register(session, temp_file, "temp-ref", mtime_ns=_mtime(temp_file)) + _register(session, stray_file, "stray-ref", mtime_ns=_mtime(stray_file)) + session.commit() + + marked = mark_references_missing_outside_prefixes(session, get_owned_prefixes()) + session.commit() + + session.expire_all() + assert marked == 1, "only the reference outside every owned directory is disowned" + assert session.get(AssetReference, "temp-ref").is_missing is False, ( + "a temp file on disk is not missing, however often the prune runs" + ) + assert session.get(AssetReference, "stray-ref").is_missing is True, ( + "owning temp must not stop the prune disowning files elsewhere" + ) + + +def test_temp_sync_marks_deleted_file_missing(session, comfy_dirs): + temp_file = _write(comfy_dirs["temp"], "preview.png") + _register(session, temp_file, "temp-ref", mtime_ns=_mtime(temp_file)) + session.commit() + os.remove(temp_file) + + sync_prefixes_with_filesystem(session, get_temp_prefixes()) + session.commit() + + session.expire_all() + assert session.get(AssetReference, "temp-ref").is_missing is True, ( + "nothing else stats temp, so this pass is what retires a wiped file" + ) + + +def test_temp_sync_drops_unhashed_asset_whose_file_is_gone(session, comfy_dirs): + temp_file = _write(comfy_dirs["temp"], "preview.png") + _register(session, temp_file, "temp-ref", mtime_ns=_mtime(temp_file), asset_hash=None) + session.commit() + os.remove(temp_file) + + sync_prefixes_with_filesystem(session, get_temp_prefixes()) + session.commit() + + session.expire_all() + assert session.get(AssetReference, "temp-ref") is None, ( + "an unhashed asset with no surviving reference is retired, not kept as missing" + ) + assert session.get(Asset, "asset-temp-ref") is None, ( + "the orphaned asset row goes with its last reference" + ) + + +def test_temp_sync_keeps_live_file(session, comfy_dirs): + temp_file = _write(comfy_dirs["temp"], "preview.png") + _register(session, temp_file, "temp-ref", mtime_ns=_mtime(temp_file)) + session.commit() + + sync_prefixes_with_filesystem(session, get_temp_prefixes()) + session.commit() + + session.expire_all() + ref = session.get(AssetReference, "temp-ref") + assert ref.is_missing is False, "the file is still there" + assert ref.needs_verify is False, "an unchanged file needs no re-verification" diff --git a/tests-unit/seeder_test/test_seeder.py b/tests-unit/seeder_test/test_seeder.py index 6aed6d6f36c..55f81358529 100644 --- a/tests-unit/seeder_test/test_seeder.py +++ b/tests-unit/seeder_test/test_seeder.py @@ -399,7 +399,7 @@ def test_mark_missing_when_idle(self, fresh_seeder: _AssetSeeder): with ( patch("app.assets.seeder.dependencies_available", return_value=True), patch( - "app.assets.seeder.get_all_known_prefixes", + "app.assets.seeder.get_owned_prefixes", return_value=["/models", "/input", "/output"], ), patch( @@ -454,8 +454,9 @@ def track_sync(root): with ( patch("app.assets.seeder.dependencies_available", return_value=True), - patch("app.assets.seeder.get_all_known_prefixes", return_value=["/models"]), + patch("app.assets.seeder.get_owned_prefixes", return_value=["/models"]), patch("app.assets.seeder.mark_missing_outside_prefixes_safely", side_effect=track_mark), + patch("app.assets.seeder.sync_temp_references_safely"), patch("app.assets.seeder.sync_root_safely", side_effect=track_sync), patch("app.assets.seeder.collect_paths_for_roots", return_value=[]), patch("app.assets.seeder.build_asset_specs", return_value=([], set(), 0)), @@ -469,6 +470,29 @@ def track_sync(root): assert call_order[0] == "mark_missing" assert "sync_models" in call_order + def test_prune_first_flag_reconciles_temp_references( + self, fresh_seeder: _AssetSeeder + ): + with ( + patch("app.assets.seeder.dependencies_available", return_value=True), + patch("app.assets.seeder.get_owned_prefixes", return_value=["/models"]), + patch("app.assets.seeder.mark_missing_outside_prefixes_safely", return_value=0), + patch("app.assets.seeder.sync_temp_references_safely") as sync_temp, + patch("app.assets.seeder.sync_root_safely", return_value=set()), + patch("app.assets.seeder.collect_paths_for_roots", return_value=[]), + patch("app.assets.seeder.build_asset_specs", return_value=([], set(), 0)), + patch("app.assets.seeder.insert_asset_specs", return_value=0), + patch("app.assets.seeder.get_unenriched_assets_for_roots", return_value=[]), + patch("app.assets.seeder.enrich_assets_batch", return_value=(0, 0)), + ): + fresh_seeder.start(roots=("models",), prune_first=True) + fresh_seeder.wait(timeout=5.0) + + assert sync_temp.called, ( + "temp is not a scan root, so the scan must reconcile it explicitly " + "or files wiped at startup stay listed" + ) + class TestSeederPhases: """Test phased scanning behavior.""" diff --git a/tests/test_asset_seeder.py b/tests/test_asset_seeder.py index 4274dab8e12..2c6727abe99 100644 --- a/tests/test_asset_seeder.py +++ b/tests/test_asset_seeder.py @@ -162,7 +162,7 @@ class TestPendingEnrichDrain: """Verify that _run_scan drains _pending_enrich via start_enrich.""" @patch("app.assets.seeder.dependencies_available", return_value=True) - @patch("app.assets.seeder.get_all_known_prefixes", return_value=[]) + @patch("app.assets.seeder.get_owned_prefixes", return_value=[]) @patch("app.assets.seeder.sync_root_safely", return_value=set()) @patch("app.assets.seeder.collect_paths_for_roots", return_value=[]) @patch("app.assets.seeder.build_asset_specs", return_value=([], {}, 0)) @@ -187,7 +187,7 @@ def test_pending_enrich_starts_after_scan(self, *_mocks): assert seeder._pending_enrich is None @patch("app.assets.seeder.dependencies_available", return_value=True) - @patch("app.assets.seeder.get_all_known_prefixes", return_value=[]) + @patch("app.assets.seeder.get_owned_prefixes", return_value=[]) @patch("app.assets.seeder.sync_root_safely", return_value=set()) @patch("app.assets.seeder.collect_paths_for_roots", return_value=[]) @patch("app.assets.seeder.build_asset_specs", return_value=([], {}, 0)) @@ -206,7 +206,7 @@ def test_pending_cleared_even_when_start_fails(self, *_mocks): assert seeder._pending_enrich is None @patch("app.assets.seeder.dependencies_available", return_value=True) - @patch("app.assets.seeder.get_all_known_prefixes", return_value=[]) + @patch("app.assets.seeder.get_owned_prefixes", return_value=[]) @patch("app.assets.seeder.sync_root_safely", return_value=set()) @patch("app.assets.seeder.collect_paths_for_roots", return_value=[]) @patch("app.assets.seeder.build_asset_specs", return_value=([], {}, 0))