ComfyAssets-ComfyUI_PromptManager/py/api/jobs.py

"""The session job registry: every long operation as one tracked job.

A job runs as an asyncio task that outlives the HTTP request that started it,
so an accidental tab close never abandons half-finished work. Each job keeps
its counters, a capped list of file outcomes per outcome group and a capped
log, and fans its progress out to SSE subscribers: the request that started
it, a reconnecting page, a second tab. A subscriber receives the latest event
as a snapshot and then every event published after it, until the job ends.

Jobs live only in server memory and are gone after a restart. The registry
keeps the running jobs and this session's last ``MAX_FINISHED_JOBS`` finished
ones. One job of a kind runs at a time, and each job has an admission class
(``ADMISSIONS``) that decides what may run beside it:

- ``records`` changes library records or files (the metadata scan, bulk image
  tagging, and the legacy restore, maintenance, cleanup, duplicate delete,
  LoRA import and batch auto-tag routes). It is refused while any
  ``records`` or ``guarded`` job runs.
- ``guarded`` must not overlap record changes but may overlap other guarded
  work (thumbnail generation and clearing, later backup). It is refused while
  a ``records`` job runs.
- ``read`` only reads; it is refused only by a running job of its own kind.

A short legacy-page operation takes a ``hold``: a running, non-cancellable
job that lasts as long as the request, so the Activity page lists it and a
workspace job is refused while it runs (and the other way round).
"""

import asyncio
import itertools
import logging
import re
import time
import uuid
from pathlib import PurePosixPath, PureWindowsPath

logger = logging.getLogger(__name__)
# Start and finish order: wall-clock times can tie (coarse Windows clocks).
_sequence = itertools.count()

JOB_STATES = ("running", "finished", "cancelled", "failed")
FINAL_STATES = JOB_STATES[1:]
MAX_OUTCOME_ROWS = 200  # rows listed per outcome group; counters stay exact
MAX_LOG_LINES = 2000
LOG_CAP_LINE = "[Log capped at 2,000 lines]"
MAX_FINISHED_JOBS = 20  # finished jobs this session keeps

# The metadata scan's outcome groups, in counter-tile order of the spec's
# mapping, and the per-file scan status each group collects. Files without
# usable metadata have no status and go to "no_metadata".
# Which files a scan reads (#249 Activity-2): "new" skips the files the Library
# already links with a resolved link (get_known_resolved_paths) before reading
# them, so new, unlinked, unresolved and ambiguous files are read; "all" reads
# every file again. The workspace job defaults to "new"; legacy /scan is "all".
SCAN_SCOPES = ("new", "all")
DEFAULT_SCAN_SCOPE = "new"
SCAN_GROUP_NAMES = ("added", "known", "review", "no_metadata", "errors")
SCAN_GROUPS = {
    "added": "added",
    "linked": "added",
    "enriched": "added",
    "unchanged": "known",
    "skipped": "known",  # "new" scope: known and resolved, not read again
    "ambiguous": "review",
    "conflict": "review",
    "unresolved": "review",
    "failed": "errors",
}

# Bulk image tagging (#249 D1): Add tags and Auto-tag over selected images.
# Record-changing: it writes image tag sets and new vocabulary names.
IMAGE_TAG_JOB_KIND = "image-tags"
IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors")

ADMISSIONS = ("read", "guarded", "records")
# What each admission class is refused beside (besides its own kind).
_REFUSED_BESIDE = {
    "read": frozenset(),
    "guarded": frozenset({"records"}),
    "records": frozenset({"records", "guarded"}),
}

# Kinds of the legacy-page operations that hold the registry (D80; titles in
# their handlers end in "(admin page)", "(gallery page)" or "(legacy page)").
RESTORE_JOB_KIND = "restore"
MAINTENANCE_JOB_KIND = "maintenance"
DUPLICATE_CLEANUP_JOB_KIND = "duplicate_cleanup"
THUMBNAILS_JOB_KIND = "thumbnails"
LORA_IMPORT_JOB_KIND = "lora_import"
AUTOTAG_BATCH_JOB_KIND = "autotag"


def resolve_admission(admission=None, record_changing=None) -> str:
    """The admission class for a start call; ``record_changing`` is an alias.

    ``record_changing=True`` means ``records`` and ``False`` means ``read``;
    with neither given the class is ``read``. Raises ValueError for an
    unknown class or one that contradicts ``record_changing``.
    """
    if admission is None:
        return "records" if record_changing else "read"
    if admission not in ADMISSIONS:
        raise ValueError(f"Unknown job admission {admission!r}")
    if record_changing is not None and record_changing != (admission == "records"):
        raise ValueError(
            f"record_changing={record_changing!r} contradicts admission {admission!r}"
        )
    return admission


def _relative_path(path) -> str:
    """``path`` as a root-relative POSIX path; an absolute one keeps its name."""
    text = str(path or "")
    if (
        PurePosixPath(text).is_absolute()
        or PureWindowsPath(text).is_absolute()
        or PureWindowsPath(text).drive
    ):
        return re.split(r"[\\/]", text)[-1]
    return text.replace("\\", "/")


class JobBusy(Exception):
    """A job could not start because ``running`` holds its slot."""

    def __init__(self, running):
        super().__init__(f"{running.title} is already running")
        self.running = running


class Job:
    """One tracked operation: state, counters, outcomes, log and subscribers."""

    def __init__(
        self,
        kind,
        title,
        *,
        admission=None,
        record_changing=None,
        cancellable=True,
        scope=None,
        groups=(),
        clock=None,
    ):
        self._clock = clock or time.time
        self.id = uuid.uuid4().hex
        self.kind = kind
        self.title = title
        self.admission = resolve_admission(admission, record_changing)
        self.cancellable = bool(cancellable)
        self.scope = dict(scope or {})
        self.state = "running"
        self.started_at = self._clock()
        self._started_seq = next(_sequence)
        self._finished_seq = None
        self.finished_at = None
        self.progress = 0
        self.status_line = ""
        self.processed = 0
        self.total = 0
        self.counters = dict.fromkeys(groups, 0)
        self.outcomes = {group: [] for group in groups}
        self.log = []
        self.cancelled = False
        self.last_event = None
        self.task = None
        self._subscribers = set()

    @property
    def record_changing(self):
        return self.admission == "records"

    @property
    def running(self):
        return self.state == "running"

    @property
    def finished(self):
        return not self.running

    @property
    def subscriber_count(self):
        return len(self._subscribers)

    # ── progress ──────────────────────────────────────────────────────

    def update(self, *, progress=None, status_line=None, processed=None, total=None):
        """Set the summary's progress fields; ``None`` leaves one unchanged."""
        if progress is not None:
            self.progress = max(0, min(100, int(progress)))
        if status_line is not None:
            self.status_line = status_line
        if processed is not None:
            self.processed = processed
        if total is not None:
            self.total = total

    def publish(self, payload):
        """Record ``payload`` as the latest event and hand it to every subscriber."""
        if isinstance(payload.get("progress"), (int, float)):
            self.update(progress=payload["progress"])
        self.last_event = payload
        for queue in list(self._subscribers):
            queue.put_nowait(payload)

    def log_line(self, text):
        """Append a timestamped log line, up to MAX_LOG_LINES and one cap line."""
        if len(self.log) > MAX_LOG_LINES:
            return
        if len(self.log) == MAX_LOG_LINES:
            self.log.append(LOG_CAP_LINE)
            return
        stamp = time.strftime("%H:%M:%S", time.localtime(self._clock()))
        self.log.append(f"{stamp} {text}")

    def record(self, group, row, log_line):
        """Count one file outcome in ``group``, list its row and log it.

        ``row`` is ``{path, root, root_index, detail}``: the path relative to
        its output root, the root's public label and its listing index (the
        identity Images opens it by; None when unknown). An absolute path is
        cut to its file name, so a row never reveals the server's layout.
        """
        if group not in self.counters:
            raise ValueError(f"Unknown outcome group {group!r} for {self.kind}")
        self.counters[group] += 1
        raw = str(row.get("path") or "")
        path = _relative_path(raw)
        if raw and raw != path:
            log_line = log_line.replace(raw, path)
        rows = self.outcomes[group]
        index = row.get("root_index")
        if isinstance(index, bool) or not isinstance(index, int) or index < 0:
            index = None
        if len(rows) < MAX_OUTCOME_ROWS:
            rows.append(
                {
                    "path": path,
                    "root": str(row.get("root") or ""),
                    "root_index": index,
                    "detail": str(row.get("detail") or ""),
                }
            )
        self.log_line(log_line)

    def finish(self, state="finished"):
        """End the job in ``state`` and release every subscriber.

        Each subscriber still gets the events queued before the release. A
        second finish is a no-op, so the first final state stands.
        """
        if state not in FINAL_STATES:
            raise ValueError(f"{state!r} is not a final job state")
        if self.finished:
            return
        self.state = state
        self.finished_at = self._clock()
        self._finished_seq = next(_sequence)
        for queue in list(self._subscribers):
            queue.put_nowait(None)

    async def events(self):
        """Yield the latest event, then each later one, until the job finishes.

        Consume it under ``contextlib.aclosing`` so a consumer that stops early
        (a closed socket) drops its queue at once instead of at garbage collection.
        """
        queue = asyncio.Queue()
        self._subscribers.add(queue)
        try:
            if self.last_event is not None:
                yield self.last_event
            if self.finished:
                return
            while True:
                payload = await queue.get()
                if payload is None:
                    return
                yield payload
        finally:
            self._subscribers.discard(queue)

    # ── views ─────────────────────────────────────────────────────────

    def summary(self):
        """The job as the Activity list and pill show it."""
        return {
            "id": self.id,
            "kind": self.kind,
            "title": self.title,
            "state": self.state,
            "started_at": self.started_at,
            "finished_at": self.finished_at,
            "progress": self.progress,
            "status_line": self.status_line,
            "processed": self.processed,
            "total": self.total,
            "counters": dict(self.counters),
            "scope": dict(self.scope),
            # What detail() lists, so a view re-reads it only when it changed.
            "listed": {group: len(rows) for group, rows in self.outcomes.items()},
            "log_length": len(self.log),
            # A view offers Cancel only when the server can stop the job.
            "cancellable": self.cancellable,
            "admission": self.admission,
        }

    def detail(self):
        """The summary plus the listed outcome rows and the log."""
        return {
            **self.summary(),
            "outcomes": {
                group: [dict(row) for row in rows]
                for group, rows in self.outcomes.items()
            },
            "log": list(self.log),
        }

    def report(self):
        """The downloadable report: the summary, every outcome group and the log.

        Each group carries its exact count, its listed rows (at most
        MAX_OUTCOME_ROWS) and whether the list is cut short. Rows and log lines
        already hold root-relative paths only.
        """
        summary = self.summary()
        summary.pop("listed", None)
        summary.pop("log_length", None)
        return {
            "report_version": 1,
            "generated_at": self._clock(),
            "job": summary,
            "outcomes": {
                group: {
                    "count": self.counters.get(group, len(rows)),
                    "rows": [dict(row) for row in rows],
                    "truncated": self.counters.get(group, len(rows)) > len(rows),
                }
                for group, rows in self.outcomes.items()
            },
            "log": list(self.log),
        }

    def report_filename(self):
        """The download name: prompt-manager-KIND-YYYYMMDD-HHMM.json (local time)."""
        stamp = time.strftime("%Y%m%d-%H%M", time.localtime(self._clock()))
        return f"prompt-manager-{self.kind}-{stamp}.json"


class JobRegistry:
    """Every job of this server session, running and recently finished."""

    def __init__(self, logger_=None):
        self._logger = logger_ or logger
        self._jobs = {}

    def start(
        self,
        kind,
        title,
        *,
        admission=None,
        record_changing=None,
        cancellable=True,
        scope=None,
        runner,
        groups=(),
    ) -> Job:
        """Create a job and run ``runner(job)`` as a task.

        ``admission`` is one of ``ADMISSIONS``; ``record_changing=True`` is
        the older spelling of ``records`` and ``False`` of ``read``. One of
        the two is required (TypeError otherwise), so no new job runs as
        ``read`` beside a scan by omission. Raises JobBusy, naming the job in
        the way, when the admission class refuses it (see the module
        docstring). Must be called from the event loop.
        """
        if admission is None and record_changing is None:
            raise TypeError(
                f"JobRegistry.start({kind!r}) needs an admission class "
                "(admission= or record_changing=)"
            )
        job = self._admit(
            kind,
            title,
            admission=admission,
            record_changing=record_changing,
            cancellable=cancellable,
            scope=scope,
            groups=groups,
        )
        job.task = asyncio.ensure_future(self._run(job, runner))
        return job

    def hold(self, kind, title, *, admission, scope=None):
        """A running, non-cancellable job for a legacy-page operation.

        Returns an async context manager. Entering it registers the job and
        yields it, or raises JobBusy at once, without waiting, when the
        admission class refuses it; nothing is registered until then, so a
        hold that is never entered leaves no job behind. The job ends when
        the block exits: ``finished``, or ``failed`` when it raised (unless
        the body already ended it). A caller that answers a refusal in its
        own way checks ``blocker(kind, admission=...)`` first, with no await
        before the ``async with``.
        """
        resolve_admission(admission)  # an unknown class fails at the call
        return _Hold(
            lambda: self._admit(
                kind, title, admission=admission, cancellable=False, scope=scope
            )
        )

    def _admit(self, kind, title, *, admission, record_changing=None, **job_args):
        resolved = resolve_admission(admission, record_changing)
        blocker = self.blocker(kind, admission=resolved)
        if blocker is not None:
            raise JobBusy(blocker)
        self._prune()
        job = Job(kind, title, admission=resolved, **job_args)
        self._jobs[job.id] = job
        return job

    def blocker(self, kind, *, admission=None, record_changing=None):
        """The running job a new ``kind`` job would be refused for, or None.

        The same kind refuses first; then the oldest running job whose class
        the new job's admission class may not run beside.
        """
        same_kind = self.active(kind)
        if same_kind is not None:
            return same_kind
        refused = _REFUSED_BESIDE[resolve_admission(admission, record_changing)]
        running = [job for job in self._jobs.values() if job.running]
        running.sort(key=lambda job: job._started_seq)
        for job in running:
            if job.admission in refused:
                return job
        return None

    async def _run(self, job, runner):
        try:
            await runner(job)
        except asyncio.CancelledError:
            job.finish("cancelled")
            raise
        except Exception:
            self._logger.exception(f"Job {job.kind} failed")
            job.finish("failed")
        finally:
            if job.running:
                job.finish("cancelled" if job.cancelled else "finished")

    def _prune(self):
        """Drop the oldest finished jobs so a new one keeps the cap."""
        finished = [job for job in self._jobs.values() if job.finished]
        finished.sort(key=lambda job: job._finished_seq)
        for job in finished[: max(0, len(finished) - MAX_FINISHED_JOBS + 1)]:
            del self._jobs[job.id]

    def get(self, job_id):
        return self._jobs.get(job_id)

    def active(self, kind):
        """The running job of ``kind``, or None."""
        for job in self._jobs.values():
            if job.kind == kind and job.running:
                return job
        return None

    def list(self):
        """Running jobs first, then the last finished ones, newest first."""
        jobs = list(self._jobs.values())
        running = [job for job in jobs if job.running]
        finished = [job for job in jobs if job.finished]
        running.sort(key=lambda job: job._started_seq, reverse=True)
        finished.sort(key=lambda job: job._finished_seq, reverse=True)
        return running + finished[:MAX_FINISHED_JOBS]

    def cancel(self, job_id) -> bool:
        """Ask a running job to stop at its next safe boundary.

        Returns whether a cancellation was requested: False for a finished,
        unknown or non-cancellable job, so a cancel after the end is a no-op.
        """
        job = self._jobs.get(job_id)
        if job is None or job.finished or not job.cancellable:
            return False
        job.cancelled = True
        return True


class _Hold:
    """The async context manager ``JobRegistry.hold`` returns."""

    def __init__(self, admit):
        self._admit = admit
        self._job = None

    async def __aenter__(self):
        self._job = self._admit()
        return self._job

    async def __aexit__(self, exc_type, exc, tb):
        if self._job.running:
            self._job.finish("finished" if exc_type is None else "failed")
        return False