ComfyUI/app/assets/lifecycle.py

"""Runs the asset system's startup and shutdown maintenance: clearing temp rows
and files, recording a hash-mode transition, and handing the filesystem scan to
the background seeder. Startup only enqueues transition work — draining it
belongs to the seeder, so a large library cannot stall the server before it
accepts requests. It also settles, before any of that, whether the database
dependencies exist at all, disabling the asset routes when they do not.
"""

from __future__ import annotations

import logging
import os
import shutil

import folder_paths
from sqlalchemy import select

from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries.records import delete_record
from app.assets.helpers import sql_path_under_prefix
from app.assets.services.hash_mode_state import enqueue_transition_work
from app.assets.services.hash_mode_state import record_transition_intent
from app.database.db import can_create_session, create_session
from comfy.cli_args import args

_excluded_scan_roots: set[str] = set()
_hash_mode_transition: str | None = None


def get_excluded_scan_roots() -> frozenset[str]:
    return frozenset(_excluded_scan_roots)


def record_hash_mode_transition_intent() -> None:
    global _hash_mode_transition

    with create_session() as session:
        _hash_mode_transition = record_transition_intent(session)
        session.commit()


def enqueue_mode_transition_work() -> None:
    with create_session() as session:
        enqueue_transition_work(session, _hash_mode_transition)
        session.commit()


def wipe_temp_db_rows(session) -> tuple[int, int]:
    try:
        temp_root = os.path.abspath(folder_paths.get_temp_directory())
    except OSError:
        return 0, 0
    # These rows are hard-deleted, so the predicate must stay case-SENSITIVE: admitting a
    # case-different persistent directory destroys user assets.
    under_temp = sql_path_under_prefix(AssetContent.path, temp_root)

    temp_record_ids = list(
        session.scalars(
            select(Asset.id)
            .join(AssetContent, Asset.content_id == AssetContent.id)
            .where(under_temp)
        )
    )

    records_deleted = 0
    for record_id in temp_record_ids:
        delete_record(session, record_id)
        records_deleted += 1

    contents_deleted = 0
    for content in session.scalars(select(AssetContent).where(under_temp)).all():
        session.delete(content)
        contents_deleted += 1

    session.flush()
    return records_deleted, contents_deleted


def cleanup_temp_filesystem() -> bool:
    temp_dir = os.path.abspath(folder_paths.get_temp_directory())
    if not os.path.exists(temp_dir):
        return True
    try:
        shutil.rmtree(temp_dir)
        return True
    except OSError as exc:
        logging.warning(
            "Failed to remove temp directory %s: %s — excluding from scan for this process",
            temp_dir,
            exc,
        )
        _excluded_scan_roots.add(temp_dir)
        return False


def start_asset_seeder() -> bool:
    from app.assets.seeder import asset_seeder

    started = asset_seeder.start(
        roots=("models", "input", "output"),
        prune_first=True,
        compute_hashes=args.enable_asset_hashing,
    )
    if started:
        logging.info("Background asset scan initiated for models, input, output")
    return started


def run_asset_startup() -> None:
    try:
        with create_session() as session:
            wipe_temp_db_rows(session)
            session.commit()
    except Exception:
        logging.exception("Temp DB row wipe failed; skipping filesystem cleanup")
        enqueue_mode_transition_work()
        start_asset_seeder()
        return
    cleanup_temp_filesystem()
    enqueue_mode_transition_work()
    start_asset_seeder()


def run_startup(*, enable_assets: bool) -> None:
    try:
        if enable_assets:
            run_asset_startup()
        else:
            cleanup_temp_filesystem()
    except Exception:
        logging.exception("Asset startup maintenance failed")


def run_asset_shutdown_cleanup() -> None:
    try:
        with create_session() as session:
            wipe_temp_db_rows(session)
            session.commit()
    except Exception:
        logging.exception("Temp DB row wipe failed during shutdown")
    finally:
        cleanup_temp_filesystem()


def run_shutdown() -> None:
    try:
        if can_create_session():
            run_asset_shutdown_cleanup()
        else:
            cleanup_temp_filesystem()
    except Exception:
        logging.exception("Asset shutdown cleanup failed")