feat(activity): shared session job registry with the scan as its first kind

The output scan now runs as the first kind of an in-memory JobRegistry
(py/api/jobs.py, replacing py/api/scan_job.py). Each job keeps exact
counters, at most 200 outcome rows per group and a log capped at 2,000
lines, with root-relative paths and public root labels only. New routes
list, inspect, follow (SSE) and cancel jobs and start the scan
(POST /prompt_manager/jobs/scan, 409 busy while another record-changing
job runs). The legacy /prompt_manager/scan routes attach to the same job
and stream unchanged payloads.
This commit is contained in:
Vito Sansevero
2026-10-21 07:13:44 -07:00
parent 762669758b
commit 89f3fa8aae
10 changed files with 1280 additions and 205 deletions
+10 -2
View File
@@ -6,6 +6,7 @@ Split into domain-specific mixins for maintainability:
- AdminRoutesMixin: duplicates, stats, settings, diagnostics, maintenance, backup, scan
- LoggingRoutesMixin: log management endpoints
- AutotagRoutesMixin: auto-tagging model management and tagging
- JobRoutesMixin: the session job registry behind Activity (py/api/jobs.py)
"""
import asyncio
@@ -28,6 +29,8 @@ from .vocabulary import VocabularyRoutesMixin
from .history import HistoryRoutesMixin
from .images import ImageRoutesMixin
from .admin import AdminRoutesMixin
from .job_routes import JobRoutesMixin
from .jobs import JobRegistry
from .logging_routes import LoggingRoutesMixin
from .autotag_routes import AutotagRoutesMixin
from .lora_integration import LoraIntegrationMixin
@@ -252,6 +255,7 @@ class PromptManagerAPI(
PromptRoutesMixin,
ImageRoutesMixin,
AdminRoutesMixin,
JobRoutesMixin,
LoggingRoutesMixin,
AutotagRoutesMixin,
LoraIntegrationMixin,
@@ -297,9 +301,12 @@ class PromptManagerAPI(
# _acquire_long_job(). Created lazily so no lock binds to an event
# loop before the server's loop exists.
self._long_jobs: dict = {}
# The current output scan (py/api/scan_job.ScanJob), if one was started.
# This session's jobs (py/api/jobs.py): the metadata scan, and later
# every other long operation the Activity page tracks.
self._jobs = JobRegistry(self.logger)
# The latest metadata scan job (also in self._jobs), if one was started.
self._scan_job = None
# The current duplicate scan (also a ScanJob), if one was started.
# The current duplicate scan (a py/api/jobs.Job outside the registry).
self._duplicate_job = None
# Run cleanup on initialization to remove any existing duplicates
@@ -436,6 +443,7 @@ class PromptManagerAPI(
self._register_prompt_routes(routes)
self._register_image_routes(routes)
self._register_admin_routes(routes)
self._register_job_routes(routes)
self._register_logging_routes(routes)
self._register_autotag_routes(routes)
self._register_lora_routes(routes)
+211 -45
View File
@@ -6,16 +6,18 @@ from contextlib import closing
import datetime
import hashlib
import json
import math
import os
import sqlite3
import tempfile
import threading
import time
from pathlib import Path
from aiohttp import web
from .prompts import thumbnail_url_for
from .scan_job import ScanJob
from .jobs import SCAN_GROUP_NAMES, SCAN_GROUPS, Job, JobBusy
try:
from ...utils.parallel import map_parallel
@@ -47,8 +49,10 @@ except ImportError:
try:
from ...database.library import display_title
from ...utils.hashing import generate_prompt_hash
except ImportError:
from database.library import display_title
from utils.hashing import generate_prompt_hash
IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff")
@@ -138,6 +142,93 @@ def _sse(payload):
return f"data: {json.dumps(payload)}\n\n"
SCAN_JOB_KIND = "scan"
SCAN_JOB_TITLE = "Scan for metadata"
# The legacy final event type of a scan and the job state it ends in.
SCAN_FINAL_STATES = {
"complete": "finished",
"cancelled": "cancelled",
"error": "failed",
}
MIN_BATCHES_FOR_ESTIMATE = 2 # "about X min left" waits for this many batches
# Log line prefix per outcome group (spec: + added, = known, ? review, ! error,
# - no usable metadata).
SCAN_LOG_PREFIXES = {
"added": "+",
"known": "=",
"review": "?",
"errors": "!",
"no_metadata": "-",
}
# What an outcome row says about a file: where it went (with the prompt's
# title) or why it was skipped. "error" is a file whose ingestion raised;
# None is a file without prompt metadata.
SCAN_TITLED_DETAILS = {
"added": "Added to the Library as “{title}”",
"linked": "Linked to “{title}”",
"enriched": "Metadata added to “{title}”",
"unchanged": "Already linked to “{title}”",
}
SCAN_DETAILS = {
"added": "Added to the Library",
"linked": "Linked to an existing prompt",
"enriched": "Metadata added to its prompt",
"unchanged": "Already linked to its prompt",
"ambiguous": "Needs review: more than one prompt could match",
"unresolved": "Needs review: the prompt could not be traced in the workflow",
"conflict": "Needs review: already linked to a different prompt",
"failed": "Error: the prompt could not be stored",
"error": "Error: the file could not be processed (see the server log)",
"found": "Metadata found, but no prompt text to store",
None: "No prompt metadata in the file",
}
def _plural(count, word):
return f"{count} {word}" if count == 1 else f"{count} {word}s"
def _scan_status_line(done, total, batches, elapsed):
""" "Processing file N of M", with an estimate once enough batches ran."""
line = f"Processing file {done} of {total}"
if batches < MIN_BATCHES_FOR_ESTIMATE or done <= 0 or done >= total:
return line
minutes = max(1, math.ceil(elapsed / done * (total - done) / 60))
return f"{line} · about {minutes} min left"
def _scan_roots(output_dirs):
"""(absolute root, public label) pairs, the most specific root first."""
public_path, _ = _public_helpers()
roots = [(os.path.abspath(str(d)), public_path(d)) for d in output_dirs]
return sorted(roots, key=lambda root: len(root[0]), reverse=True)
def _scan_location(media_file, roots):
"""``media_file`` relative to its output root, and that root's label."""
full = os.path.abspath(str(media_file))
folded = os.path.normcase(full)
for root, label in roots:
base = os.path.normcase(root).rstrip("\\/") + os.sep
if folded.startswith(base):
return Path(os.path.relpath(full, root)).as_posix(), label
return Path(full).name, ""
def _scan_outcome(media_file, roots, status, title=None):
"""The (group, row, log line) one scanned file records in its job."""
if status in SCAN_GROUPS:
group = SCAN_GROUPS[status]
else:
group = "errors" if status == "error" else "no_metadata"
template = SCAN_TITLED_DETAILS.get(status) if title else None
detail = template.format(title=title) if template else SCAN_DETAILS[status]
path, label = _scan_location(media_file, roots)
row = {"path": path, "root": label, "detail": detail}
where = f"{label}/{path}" if label else path
return group, row, f"{SCAN_LOG_PREFIXES[group]} {where} · {detail}"
def _inspect_restore_upload(path):
"""Checks on a restore upload beyond PromptModel.verify_database_file.
@@ -437,7 +528,7 @@ class AdminRoutesMixin:
guard = self._acquire_long_job(self.LONG_JOB_DUPLICATES)
if guard is None:
return self._long_job_busy_response("Duplicate scan or delete")
job = ScanJob()
job = Job("duplicates", "Duplicate scan")
self._duplicate_job = job
job.task = asyncio.ensure_future(self._run_duplicate_job(job, guard))
return await self._stream_scan_job(request, job)
@@ -1319,14 +1410,18 @@ class AdminRoutesMixin:
async def scan_images(self, request):
"""Scan ComfyUI output images for prompt metadata and add them to the database.
The scan runs as a background task that outlives this request, so a
closed tab never leaves a half-finished scan. Progress streams as
server-sent events; while a scan is running, a new request attaches
to it instead of starting another one.
The scan runs as the registry's ``scan`` job, which outlives this
request, so a closed tab never leaves a half-finished scan. Progress
streams as the legacy server-sent events; while a scan is running
(started here or from the workspace), a new request attaches to it
instead of starting another one.
"""
job = self._scan_job
if job is None or job.finished:
job = self._start_scan_job()
job = self._jobs.active(SCAN_JOB_KIND)
if job is None:
try:
job = self._start_scan_job()
except JobBusy as busy:
return self._long_job_busy_response(busy.running.title)
return await self._stream_scan_job(request, job)
async def cancel_scan(self, request):
@@ -1349,18 +1444,33 @@ class AdminRoutesMixin:
)
def _start_scan_job(self):
"""Create the job record synchronously, then run the scan as a task."""
job = ScanJob()
"""Register the scan job synchronously; its runner starts on the loop.
Raises JobBusy while another record-changing job runs.
"""
job = self._jobs.start(
SCAN_JOB_KIND,
SCAN_JOB_TITLE,
record_changing=True,
scope={"mode": "all"},
runner=self._run_scan_job,
groups=SCAN_GROUP_NAMES,
)
self._scan_job = job
job.task = asyncio.ensure_future(self._run_scan_job(job))
return job
async def _run_scan_job(self, job):
try:
async for payload in self._scan_images_events():
job.publish(payload)
finally:
job.finish()
final = None
async for payload in self._scan_images_events(job):
job.publish(payload)
final = payload.get("type")
state = SCAN_FINAL_STATES.get(final, "finished")
if state == "failed":
job.log_line(f"Scan failed · {job.status_line}")
else:
verb = "cancelled" if state == "cancelled" else "finished"
job.log_line(f"Scan {verb} · {_plural(job.processed, 'file')} processed")
job.finish(state)
async def _stream_scan_job(self, request, job):
"""Stream ``job`` as SSE; stop quietly if the client goes away."""
@@ -1386,21 +1496,30 @@ class AdminRoutesMixin:
)
return response
async def _scan_images_events(self):
"""Yield progress payloads while scanning every configured output directory."""
async def _scan_images_events(self, job):
"""Yield progress payloads while scanning every configured output directory.
The payloads are the legacy SSE events, unchanged; ``job`` also gets
its status line, totals and one outcome per processed file.
"""
try:
self.logger.info("Starting image scan operation")
output_dirs = self._get_all_output_dirs()
if not output_dirs:
self.logger.error("No output directories found")
yield {
"type": "error",
"message": "No output directories found. "
"Configure scan directories in Settings.",
}
message = (
"No output directories found. "
"Configure scan directories in Settings."
)
job.update(status_line=message)
yield {"type": "error", "message": message}
return
dir_names = [self._public_path(d) for d in output_dirs]
folders = _plural(len(output_dirs), "output folder")
job.scope["output_folders"] = len(output_dirs)
job.log_line(f"Scan started · {folders}")
job.update(status_line=f"Looking for media files in {folders}")
yield self._scan_progress(
0, f"Scanning {len(output_dirs)} directory(ies) for media files..."
)
@@ -1409,11 +1528,16 @@ class AdminRoutesMixin:
self._collect_output_media_sync, output_dirs
)
counts = dict.fromkeys(SCAN_COUNT_KEYS, 0)
job.update(
total=len(media_files),
status_line=_scan_status_line(0, len(media_files), 0, 0),
)
if media_files:
yield self._scan_progress(
5, f"Found {len(media_files)} media files to process...", counts
)
async for event in self._scan_batches(media_files, counts):
roots = _scan_roots(output_dirs)
async for event in self._scan_batches(media_files, counts, job, roots):
yield event
self.logger.info(
@@ -1421,37 +1545,55 @@ class AdminRoutesMixin:
f"found={counts['found']}, new_prompts_added={counts['added']}, "
f"images_linked_to_existing={counts['linked']}"
)
job = self._scan_job
final_type = "cancelled" if job and job.cancelled else "complete"
final_type = "cancelled" if job.cancelled else "complete"
if final_type == "complete":
job.update(progress=100)
yield {"type": final_type, **counts, "directories": dir_names}
except Exception:
self.logger.exception("Scan error")
job.update(status_line="The scan stopped on an internal error.")
yield {
"type": "error",
"message": "An internal error occurred. Check server logs for details.",
}
async def _scan_batches(self, media_files, counts):
async def _scan_batches(self, media_files, counts, job, roots=()):
"""Process ``media_files`` in executor batches, updating ``counts`` in place.
Yields one progress event per batch.
Records each file's outcome in ``job`` and yields one progress event
per batch.
"""
total = len(media_files)
started = time.monotonic()
batches = 0
for batch_start in range(0, total, SCAN_BATCH_SIZE):
if self._scan_job and self._scan_job.cancelled:
if job.cancelled:
return
batch = media_files[batch_start : batch_start + SCAN_BATCH_SIZE]
batch_results = await self._run_in_executor(
self._extract_batch_metadata_sync, batch
)
outcomes = []
batch_counts = await self._run_in_executor(
self._ingest_scan_batch_sync, batch_results
self._ingest_scan_batch_sync,
batch_results,
roots=roots,
outcomes=outcomes,
)
for key in counts:
counts[key] += batch_counts[key]
for group, row, log_line in outcomes:
job.record(group, row, log_line)
batches += 1
done = min(batch_start + len(batch), total)
job.update(
processed=counts["processed"],
status_line=_scan_status_line(
done, total, batches, time.monotonic() - started
),
)
yield self._scan_progress(
int(done / total * 100), f"Processing file {done}/{total}...", counts
)
@@ -1490,9 +1632,19 @@ class AdminRoutesMixin:
return map_parallel(extract, file_batch, PromptManagerConfig.WORKER_THREADS)
def _ingest_scan_batch_sync(self, batch_results):
"""Store the prompts found in a metadata batch (blocking DB work)."""
def _ingest_scan_batch_sync(self, batch_results, roots=(), outcomes=None):
"""Store the prompts found in a metadata batch (blocking DB work).
Returns the legacy counts. When ``outcomes`` is a list, every processed
file appends its ``(group, row, log_line)`` for the scan job, with the
path relative to its root in ``roots`` (see ``_scan_roots``).
"""
counts = dict.fromkeys(SCAN_COUNT_KEYS, 0)
def note(media_file, status, title=None):
if outcomes is not None:
outcomes.append(_scan_outcome(media_file, roots, status, title))
for media_file, metadata in batch_results:
if self._scan_job and self._scan_job.cancelled:
break
@@ -1502,25 +1654,21 @@ class AdminRoutesMixin:
except Exception as e:
self.logger.error(f"Error processing {media_file.name}: {e}")
counts["failed"] += 1
note(media_file, "error")
continue
if outcome is None:
note(media_file, None)
continue
counts["found"] += 1
status = outcome["status"] if isinstance(outcome, dict) else outcome
if status in (
"added",
"linked",
"enriched",
"unchanged",
"ambiguous",
"unresolved",
"conflict",
"failed",
):
if status in SCAN_GROUPS:
counts[status] += 1
if isinstance(outcome, dict):
counts["pairs_added"] += outcome.get("pairs_added", 0)
counts["pair_links_added"] += outcome.get("links_added", 0)
note(media_file, status, outcome.get("title"))
else:
note(media_file, status)
return counts
@staticmethod
@@ -1539,7 +1687,13 @@ class AdminRoutesMixin:
return positive.strip() or None
def _ingest_scanned_file(self, media_file, metadata, detailed=False):
"""Enrich exact pairs; retain legacy positive-only import for absent pairs."""
"""Enrich exact pairs; retain legacy positive-only import for absent pairs.
Returns the file's scan status, or None without prompt metadata. With
``detailed``, a stored or linked file returns a dict instead: its
``status``, the prompt's display ``title`` and, for a resolved pair,
the pair counts.
"""
if not metadata:
return None
canonical = dict(metadata)
@@ -1557,7 +1711,12 @@ class AdminRoutesMixin:
result = ingest_resolved_pair(
self.db, str(media_file), canonical, resolution
)
return result if detailed else result["status"]
if not detailed:
return result["status"]
return {
**result,
"title": display_title((), resolution["positive"]["text"]),
}
if canonical.get("prompt") and resolution["state"] in (
"ambiguous",
"unresolved",
@@ -1585,6 +1744,13 @@ class AdminRoutesMixin:
return "found"
prompt_text = prompt_text.strip()
status = self._store_scanned_prompt(media_file, canonical, prompt_text)
if not detailed or status not in SCAN_TITLED_DETAILS:
return status
return {"status": status, "title": display_title((), prompt_text)}
def _store_scanned_prompt(self, media_file, canonical, prompt_text):
"""Link ``media_file`` to ``prompt_text``'s prompt, saving it when new."""
prompt_hash = generate_prompt_hash(prompt_text)
existing = self.db.get_prompt_by_hash(prompt_hash)
if existing:
+109
View File
@@ -0,0 +1,109 @@
"""Session job routes (Activity): list, inspect, follow and cancel jobs.
The registry itself is ``py/api/jobs.py``; the scan job's runner is in
``py/api/admin.py``. Every response leaves out absolute server paths: job
rows already carry root-relative paths and public root labels.
"""
import contextlib
import json
from aiohttp import web
from .jobs import JobBusy
def _not_found():
return web.json_response({"success": False, "error": "Job not found"}, status=404)
def _summary_event(job):
return f"data: {json.dumps(job.summary())}\n\n".encode("utf-8")
class JobRoutesMixin:
"""Mixin providing the /prompt_manager/jobs endpoints."""
def _register_job_routes(self, routes):
@routes.get("/prompt_manager/jobs")
async def list_jobs_route(request):
return await self.list_jobs(request)
@routes.post("/prompt_manager/jobs/scan")
async def start_scan_job_route(request):
return await self.start_scan_job(request)
@routes.get("/prompt_manager/jobs/{job_id}")
async def job_detail_route(request):
return await self.job_detail(request)
@routes.get("/prompt_manager/jobs/{job_id}/events")
async def job_events_route(request):
return await self.job_events(request)
@routes.post("/prompt_manager/jobs/{job_id}/cancel")
async def cancel_job_route(request):
return await self.cancel_job(request)
async def list_jobs(self, request):
"""This session's jobs: running first, then finished, newest first."""
jobs = [job.summary() for job in self._jobs.list()]
return web.json_response({"success": True, "jobs": jobs})
async def job_detail(self, request):
job = self._jobs.get(request.match_info["job_id"])
if job is None:
return _not_found()
return web.json_response({"success": True, "job": job.detail()})
async def cancel_job(self, request):
"""Stop a job at its next file boundary; a finished job is a no-op."""
job_id = request.match_info["job_id"]
if self._jobs.get(job_id) is None:
return _not_found()
requested = self._jobs.cancel(job_id)
return web.json_response({"success": True, "cancellation_requested": requested})
async def start_scan_job(self, request):
"""Start the metadata scan as a job (202), or 409 naming the busy job."""
try:
job = self._start_scan_job()
except JobBusy as busy:
running = busy.running
return web.json_response(
{
"success": False,
"code": "busy",
"error": str(busy),
"running": {"id": running.id, "title": running.title},
},
status=409,
)
return web.json_response({"success": True, "job": job.summary()}, status=202)
async def job_events(self, request):
"""Stream the job's summary as SSE: now, after each event, and at the end."""
job = self._jobs.get(request.match_info["job_id"])
if job is None:
return _not_found()
response = web.StreamResponse(
status=200,
reason="OK",
headers={
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
"Connection": "keep-alive",
},
)
await response.prepare(request)
try:
await response.write(_summary_event(job))
async with contextlib.aclosing(job.events()) as events:
async for _payload in events:
await response.write(_summary_event(job))
# The last event is published before the job takes its final state.
await response.write(_summary_event(job))
await response.write_eof()
except Exception as e: # the socket closed under us; the job goes on
self.logger.info(f"Job progress client disconnected ({e})")
return response
+314
View File
@@ -0,0 +1,314 @@
"""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. Record-changing jobs (the metadata scan, and later maintenance, cleanup
and restore) never run together, and one job of a kind runs at a time.
"""
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".
SCAN_GROUP_NAMES = ("added", "known", "review", "no_metadata", "errors")
SCAN_GROUPS = {
"added": "added",
"linked": "added",
"enriched": "added",
"unchanged": "known",
"ambiguous": "review",
"conflict": "review",
"unresolved": "review",
"failed": "errors",
}
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, *, record_changing=False, scope=None, groups=(), clock=None
):
self._clock = clock or time.time
self.id = uuid.uuid4().hex
self.kind = kind
self.title = title
self.record_changing = record_changing
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 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, detail}``: the path relative to its output
root and the root's public label. 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]
if len(rows) < MAX_OUTCOME_ROWS:
rows.append(
{
"path": path,
"root": str(row.get("root") or ""),
"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),
}
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),
}
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, *, record_changing, scope, runner, groups=()) -> Job:
"""Create a job and run ``runner(job)`` as a task.
Raises JobBusy, naming the job in the way, when a job of ``kind`` is
running or when ``record_changing`` and another record-changing job
runs. Must be called from the event loop.
"""
same_kind = self.active(kind)
if same_kind is not None:
raise JobBusy(same_kind)
if record_changing:
for job in self._jobs.values():
if job.running and job.record_changing:
raise JobBusy(job)
self._prune()
job = Job(
kind, title, record_changing=record_changing, scope=scope, groups=groups
)
self._jobs[job.id] = job
job.task = asyncio.ensure_future(self._run(job, runner))
return job
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 or
unknown job, so a cancel after the end is a no-op.
"""
job = self._jobs.get(job_id)
if job is None or job.finished:
return False
job.cancelled = True
return True
-62
View File
@@ -1,62 +0,0 @@
"""Progress fan-out for the background output scan.
The scan runs as an asyncio task that outlives the HTTP request that started
it, so an accidental tab close does not abandon a half-finished scan. Every
SSE subscriber (the request that started the scan, a reconnecting page, a
second tab) receives the latest event as a snapshot and then every event
published after it, until the job finishes.
"""
import asyncio
class ScanJob:
"""One scan run: its latest event, completion flag and live subscribers."""
def __init__(self):
self.last_event = None
self.finished = False
self.cancelled = False
self.task = None
self._subscribers = set()
@property
def running(self):
return not self.finished
@property
def subscriber_count(self):
return len(self._subscribers)
def publish(self, payload):
"""Record ``payload`` as the latest event and hand it to every subscriber."""
self.last_event = payload
for queue in list(self._subscribers):
queue.put_nowait(payload)
def finish(self):
"""Mark the job done and release every subscriber after its queued events."""
self.finished = True
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)
+226
View File
@@ -1965,6 +1965,232 @@ class TestOutputScanIsABackgroundJob(AdminAPITestCase):
self.assertEqual(data["last_event"]["type"], "error")
class TestActivityJobs(AdminAPITestCase):
"""The scan runs as a registry job; legacy routes keep their payloads."""
def _gate_collect(self):
import threading
gate = threading.Event()
calls = []
original = self.api._collect_output_media_sync
def gated(output_dirs):
calls.append(output_dirs)
gate.wait(5)
return original(output_dirs)
self.api._collect_output_media_sync = gated
return gate, calls
async def _json(self, method, url):
resp = await self.client.request(method, url)
return resp.status, await resp.json()
async def test_legacy_scan_route_reuses_registry_job_and_payload_shape(self):
from py.api.admin import SCAN_COUNT_KEYS
_write_png(self.output_dir / "gen.png", prompt_text="a red square")
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
gate, calls = self._gate_collect()
# The workspace starts the job; the legacy page then attaches to it.
status, started = await self._json("POST", "/prompt_manager/jobs/scan")
self.assertEqual(status, 202)
self.assertTrue(started["success"])
legacy = asyncio.ensure_future(
self.client.request("POST", "/prompt_manager/scan", data="{}")
)
await asyncio.sleep(0.05)
gate.set()
body = await (await legacy).text()
await asyncio.wait_for(self.api._scan_job.task, 10)
self.assertEqual(len(calls), 1)
payloads = _sse_payloads(body)
progress = [p for p in payloads if p["type"] == "progress"]
self.assertTrue(progress)
legacy_keys = set(SCAN_COUNT_KEYS) | {"type", "progress", "status", "processed"}
for payload in progress:
self.assertEqual(set(payload), legacy_keys)
final = payloads[-1]
self.assertEqual(final["type"], "complete")
self.assertEqual(set(final), set(SCAN_COUNT_KEYS) | {"type", "directories"})
self.assertEqual(final["added"], 1)
self.assertEqual(final["directories"], ["output"])
status, listed = await self._json("GET", "/prompt_manager/jobs")
self.assertEqual(status, 200)
self.assertEqual(len(listed["jobs"]), 1)
job = listed["jobs"][0]
self.assertEqual(job["id"], started["job"]["id"])
self.assertEqual(job["id"], self.api._scan_job.id)
self.assertEqual((job["kind"], job["state"]), ("scan", "finished"))
self.assertEqual(job["counters"]["added"], 1)
self.assertEqual(job["progress"], 100)
self.assertEqual(job["scope"]["output_folders"], 1)
async def test_scan_records_outcomes_by_group(self):
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
(self.output_dir / "sub").mkdir()
names = {
"sub/added.png": {"status": "added", "title": "A red fox"},
"linked.png": {"status": "linked", "title": "A grey cat"},
"unchanged.png": {"status": "unchanged", "title": "A brown dog"},
"ambiguous.png": {"status": "ambiguous"},
"failed.png": RuntimeError(f"{self.tmpdir}/secret boom"),
"none.png": None,
}
for name in names:
(self.output_dir / name).write_bytes(b"not an image")
def fake_ingest(media_file, metadata, detailed=False):
key = Path(media_file).relative_to(self.output_dir).as_posix()
outcome = names[key]
if isinstance(outcome, Exception):
raise outcome
return outcome
self.api._ingest_scanned_file = fake_ingest
with self.assertLogs("prompt_manager.api", level="ERROR"):
status, started = await self._json("POST", "/prompt_manager/jobs/scan")
self.assertEqual(status, 202)
await asyncio.wait_for(self.api._scan_job.task, 10)
resp = await self.client.request(
"GET", f"/prompt_manager/jobs/{started['job']['id']}"
)
text = await resp.text()
self.assertEqual(resp.status, 200)
self.assertNotIn(self.tmpdir, text)
job = json.loads(text)["job"]
self.assertEqual(
job["counters"],
{"added": 2, "known": 1, "review": 1, "no_metadata": 1, "errors": 1},
)
outcomes = job["outcomes"]
self.assertEqual(
{k: len(v) for k, v in outcomes.items()},
{"added": 2, "known": 1, "review": 1, "no_metadata": 1, "errors": 1},
)
added = {row["path"]: row for row in outcomes["added"]}
self.assertEqual(set(added), {"sub/added.png", "linked.png"})
self.assertIn("A red fox", added["sub/added.png"]["detail"])
self.assertIn("A grey cat", added["linked.png"]["detail"])
self.assertEqual(added["linked.png"]["root"], "output")
self.assertIn("A brown dog", outcomes["known"][0]["detail"])
self.assertIn("review", outcomes["review"][0]["detail"].lower())
self.assertEqual(outcomes["errors"][0]["path"], "failed.png")
self.assertTrue(outcomes["errors"][0]["detail"])
self.assertEqual(outcomes["no_metadata"][0]["path"], "none.png")
self.assertTrue(outcomes["no_metadata"][0]["detail"])
prefixes = {line.split(" ", 2)[1] for line in job["log"][1:-1]}
self.assertEqual(prefixes, {"+", "=", "?", "!", "-"})
self.assertEqual(job["state"], "finished")
self.assertTrue(job["status_line"].startswith("Processing file 6 of 6"))
legacy = self.api._scan_job.last_event
self.assertEqual(
(legacy["processed"], legacy["found"], legacy["failed"]), (6, 4, 1)
)
async def test_jobs_routes_404_and_409(self):
_write_png(self.output_dir / "gen.png", prompt_text="a red square")
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
for method, url in (
("GET", "/prompt_manager/jobs/nope"),
("GET", "/prompt_manager/jobs/nope/events"),
("POST", "/prompt_manager/jobs/nope/cancel"),
):
status, data = await self._json(method, url)
self.assertEqual(status, 404, url)
self.assertFalse(data["success"])
gate, calls = self._gate_collect()
status, started = await self._json("POST", "/prompt_manager/jobs/scan")
self.assertEqual(status, 202)
job_id = started["job"]["id"]
self.assertEqual(started["job"]["state"], "running")
status, busy = await self._json("POST", "/prompt_manager/jobs/scan")
self.assertEqual(status, 409)
self.assertEqual(
busy["running"], {"id": job_id, "title": started["job"]["title"]}
)
self.assertEqual((busy["success"], busy["code"]), (False, "busy"))
watcher = asyncio.ensure_future(
self.client.request("GET", f"/prompt_manager/jobs/{job_id}/events")
)
await asyncio.sleep(0.05)
status, cancel = await self._json(
"POST", f"/prompt_manager/jobs/{job_id}/cancel"
)
self.assertEqual((status, cancel["cancellation_requested"]), (200, True))
gate.set()
events = _sse_payloads(await (await watcher).text())
await asyncio.wait_for(self.api._scan_job.task, 10)
self.assertEqual(len(calls), 1)
self.assertEqual(events[0]["id"], job_id)
self.assertEqual(events[-1]["state"], "cancelled")
status, cancel = await self._json(
"POST", f"/prompt_manager/jobs/{job_id}/cancel"
)
self.assertEqual((status, cancel["cancellation_requested"]), (200, False))
status, legacy = await self._json("GET", "/prompt_manager/scan/status")
self.assertEqual(legacy["last_event"]["type"], "cancelled")
async def test_scan_without_directories_fails_the_job(self):
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.outside_dir)]
sys.modules.pop("folder_paths", None)
self.api._find_comfyui_output_dir = lambda: None
status, started = await self._json("POST", "/prompt_manager/jobs/scan")
await asyncio.wait_for(self.api._scan_job.task, 10)
status, data = await self._json(
"GET", f"/prompt_manager/jobs/{started['job']['id']}"
)
self.assertEqual(data["job"]["state"], "failed")
self.assertIn("No output directories", data["job"]["status_line"])
async def test_legacy_scan_refused_while_another_record_job_runs(self):
hold = asyncio.Event()
async def runner(job):
await hold.wait()
other = self.api._jobs.start(
"maintenance", "Maintenance", record_changing=True, scope={}, runner=runner
)
resp = await self.client.request("POST", "/prompt_manager/scan", data="{}")
self.assertEqual(resp.status, 409)
self.assertIn("Maintenance", (await resp.json())["error"])
hold.set()
await asyncio.wait_for(other.task, 5)
class TestScanStatusLine(unittest.TestCase):
def test_estimate_appears_after_two_batches(self):
from py.api.admin import _scan_status_line
self.assertEqual(
_scan_status_line(50, 200, 1, 10.0), "Processing file 50 of 200"
)
self.assertEqual(
_scan_status_line(100, 1000, 2, 30.0),
"Processing file 100 of 1000 · about 5 min left",
)
self.assertEqual(
_scan_status_line(100, 120, 2, 1.0),
"Processing file 100 of 120 · about 1 min left",
)
self.assertEqual(
_scan_status_line(120, 120, 3, 9.0), "Processing file 120 of 120"
)
class TestWorkerThreadsSetting(AdminAPITestCase):
"""The worker-thread count is exposed, validated, persisted and used by the scan."""
+406
View File
@@ -0,0 +1,406 @@
"""Unit tests for py/api/jobs.py: the in-memory session job registry.
Fake runners only: no files, no database. Ordering is driven by events and
awaited tasks, never by wall-clock bounds.
"""
import asyncio
import contextlib
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py.api.jobs import ( # noqa: E402
LOG_CAP_LINE,
MAX_FINISHED_JOBS,
MAX_LOG_LINES,
MAX_OUTCOME_ROWS,
SCAN_GROUP_NAMES,
SCAN_GROUPS,
Job,
JobBusy,
JobRegistry,
)
async def _collect(job, stop_after=None):
seen = []
async with contextlib.aclosing(job.events()) as events:
async for payload in events:
seen.append(payload)
if stop_after is not None and len(seen) >= stop_after:
break
return seen
def _gated_runner(gate, publish=None):
"""A runner that publishes ``publish`` (if any) and waits for ``gate``."""
async def runner(job):
if publish is not None:
job.publish(publish)
await gate.wait()
return runner
async def _noop(job):
return None
class TestJobFanOut(unittest.IsolatedAsyncioTestCase):
"""The SSE fan-out the scan job always had, now on every Job."""
async def test_new_job_is_running_with_no_event(self):
job = Job("scan", "Scan")
self.assertTrue(job.running)
self.assertIsNone(job.last_event)
self.assertEqual(job.state, "running")
self.assertEqual(len(job.id), 32)
async def test_subscriber_receives_events_until_finish(self):
job = Job("scan", "Scan")
consumer = asyncio.ensure_future(_collect(job))
await asyncio.sleep(0)
job.publish({"type": "progress", "progress": 10})
job.publish({"type": "complete"})
job.finish("finished")
self.assertEqual(
await asyncio.wait_for(consumer, 2),
[{"type": "progress", "progress": 10}, {"type": "complete"}],
)
self.assertFalse(job.running)
async def test_late_subscriber_gets_latest_snapshot_then_live_events(self):
job = Job("scan", "Scan")
job.publish({"type": "progress", "progress": 10})
job.publish({"type": "progress", "progress": 40})
consumer = asyncio.ensure_future(_collect(job))
await asyncio.sleep(0)
job.publish({"type": "complete"})
job.finish()
self.assertEqual(
await asyncio.wait_for(consumer, 2),
[{"type": "progress", "progress": 40}, {"type": "complete"}],
)
async def test_subscriber_after_finish_gets_only_the_final_event(self):
job = Job("scan", "Scan")
job.publish({"type": "progress", "progress": 99})
job.publish({"type": "complete"})
job.finish()
self.assertEqual(await _collect(job), [{"type": "complete"}])
async def test_departed_subscriber_is_forgotten(self):
job = Job("scan", "Scan")
job.publish({"type": "progress", "progress": 1})
await _collect(job, stop_after=1)
self.assertEqual(job.subscriber_count, 0)
job.publish({"type": "complete"}) # must not raise or queue anywhere
job.finish()
async def test_second_subscriber_attaches_to_running_job(self):
registry = JobRegistry()
gate = asyncio.Event()
job = registry.start(
"scan",
"Scan",
record_changing=True,
scope={},
runner=_gated_runner(gate, {"type": "progress", "progress": 30}),
)
await asyncio.sleep(0) # the runner publishes its first event
first = asyncio.ensure_future(_collect(job))
second = asyncio.ensure_future(_collect(job))
await asyncio.sleep(0)
self.assertEqual(job.subscriber_count, 2)
self.assertIs(registry.active("scan"), job)
job.publish({"type": "progress", "progress": 60})
gate.set()
await asyncio.wait_for(job.task, 2)
expected = [
{"type": "progress", "progress": 30},
{"type": "progress", "progress": 60},
]
self.assertEqual(await asyncio.wait_for(first, 2), expected)
self.assertEqual(await asyncio.wait_for(second, 2), expected)
self.assertEqual(job.summary()["progress"], 60)
async def test_finish_rejects_unknown_state_and_is_final(self):
job = Job("scan", "Scan")
with self.assertRaises(ValueError):
job.finish("running")
job.finish("cancelled")
job.finish("finished") # a second finish keeps the first state
self.assertEqual(job.state, "cancelled")
self.assertIsNotNone(job.finished_at)
class TestJobRegistry(unittest.IsolatedAsyncioTestCase):
async def test_start_lists_and_finishes_a_job(self):
registry = JobRegistry()
gate = asyncio.Event()
job = registry.start(
"scan",
"Scan for metadata",
record_changing=True,
scope={"mode": "all"},
runner=_gated_runner(gate),
)
self.assertEqual(job.state, "running")
self.assertIs(registry.get(job.id), job)
self.assertIsNone(registry.get("missing"))
summary = job.summary()
self.assertEqual(
set(summary),
{
"id",
"kind",
"title",
"state",
"started_at",
"finished_at",
"progress",
"status_line",
"processed",
"total",
"counters",
"scope",
},
)
self.assertEqual(summary["scope"], {"mode": "all"})
self.assertIsNone(summary["finished_at"])
gate.set()
await asyncio.wait_for(job.task, 2)
self.assertEqual(job.state, "finished")
self.assertIsNone(registry.active("scan"))
# Running jobs list first, then finished ones newest first.
older = registry.start(
"thumbs", "Thumbnails", record_changing=False, scope={}, runner=_noop
)
await asyncio.wait_for(older.task, 2)
hold = asyncio.Event()
running = registry.start(
"other",
"Other",
record_changing=False,
scope={},
runner=_gated_runner(hold),
)
listed = registry.list()
self.assertEqual([j.id for j in listed], [running.id, older.id, job.id])
# Only the last MAX_FINISHED_JOBS finished jobs are kept.
for index in range(MAX_FINISHED_JOBS + 5):
done = registry.start(
f"k{index}", "Fake", record_changing=False, scope={}, runner=_noop
)
await asyncio.wait_for(done.task, 2)
listed = registry.list()
finished = [j for j in listed if not j.running]
self.assertEqual(len(finished), MAX_FINISHED_JOBS)
self.assertEqual(finished[0].id, done.id)
self.assertIsNone(registry.get(job.id))
self.assertIs(listed[0], running)
hold.set()
await asyncio.wait_for(running.task, 2)
async def test_record_changing_jobs_are_exclusive(self):
registry = JobRegistry()
gate = asyncio.Event()
first = registry.start(
"scan",
"Scan for metadata",
record_changing=True,
scope={},
runner=_gated_runner(gate),
)
with self.assertRaises(JobBusy) as raised:
registry.start(
"maintenance",
"Maintenance",
record_changing=True,
scope={},
runner=_noop,
)
self.assertIs(raised.exception.running, first)
self.assertIn("Scan for metadata", str(raised.exception))
reader = registry.start(
"report", "Report", record_changing=False, scope={}, runner=_noop
)
await asyncio.wait_for(reader.task, 2)
self.assertEqual(reader.state, "finished")
# One job per kind, even a read-only one.
hold = asyncio.Event()
registry.start(
"report",
"Report",
record_changing=False,
scope={},
runner=_gated_runner(hold),
)
with self.assertRaises(JobBusy):
registry.start(
"report", "Report", record_changing=False, scope={}, runner=_noop
)
gate.set()
hold.set()
await asyncio.wait_for(first.task, 2)
async def test_cancel_is_idempotent_and_keeps_partial_counts(self):
registry = JobRegistry()
reached = asyncio.Event()
proceed = asyncio.Event()
async def runner(job):
for index in range(3):
if job.cancelled:
return
job.record("added", {"path": f"a{index}.png", "root": "output"}, "+")
if index == 0:
reached.set()
await proceed.wait()
job = registry.start(
"scan",
"Scan",
record_changing=True,
scope={},
runner=runner,
groups=SCAN_GROUP_NAMES,
)
await asyncio.wait_for(reached.wait(), 2)
self.assertTrue(registry.cancel(job.id))
self.assertTrue(registry.cancel(job.id)) # still running: still requested
proceed.set()
await asyncio.wait_for(job.task, 2)
self.assertEqual(job.state, "cancelled")
self.assertEqual(job.counters["added"], 1) # the work done before the stop
self.assertFalse(registry.cancel(job.id)) # finished: a no-op
self.assertFalse(registry.cancel("missing"))
self.assertEqual(job.state, "cancelled")
async def test_failing_runner_ends_failed_and_frees_the_slot(self):
registry = JobRegistry()
async def runner(job):
raise RuntimeError("boom")
job = registry.start(
"scan", "Scan", record_changing=True, scope={}, runner=runner
)
with self.assertLogs("py.api.jobs", level="ERROR"):
await asyncio.wait_for(job.task, 2)
self.assertEqual(job.state, "failed")
self.assertIsNone(registry.active("scan"))
async def test_runner_that_finishes_its_own_job_keeps_its_state(self):
registry = JobRegistry()
async def runner(job):
job.finish("failed")
job = registry.start(
"scan", "Scan", record_changing=True, scope={}, runner=runner
)
await asyncio.wait_for(job.task, 2)
self.assertEqual(job.state, "failed")
class TestJobOutcomesAndLog(unittest.TestCase):
def test_outcome_and_log_caps(self):
job = Job("scan", "Scan", groups=SCAN_GROUP_NAMES)
for index in range(MAX_OUTCOME_ROWS + 50):
job.record("errors", {"path": f"f{index}.png", "root": "output"}, "!")
self.assertEqual(len(job.outcomes["errors"]), MAX_OUTCOME_ROWS)
self.assertEqual(job.counters["errors"], MAX_OUTCOME_ROWS + 50)
self.assertEqual(job.outcomes["errors"][0]["path"], "f0.png")
job = Job("scan", "Scan")
for index in range(MAX_LOG_LINES + 100):
job.log_line(f"line {index}")
self.assertEqual(len(job.log), MAX_LOG_LINES + 1)
self.assertEqual(job.log[-1], LOG_CAP_LINE)
self.assertTrue(job.log[0].endswith("line 0"))
self.assertEqual(LOG_CAP_LINE, "[Log capped at 2,000 lines]")
def test_paths_are_root_relative(self):
job = Job("scan", "Scan", groups=SCAN_GROUP_NAMES)
absolute = os.path.abspath(os.path.join("outputs", "sub", "a.png"))
job.record(
"added",
{"path": absolute, "root": "output", "detail": "x"},
f"+ {absolute}",
)
job.record("known", {"path": "sub/b.png", "root": "output"}, "= sub/b.png")
job.record("review", {"path": "C:\\out\\c.png", "root": "output"}, "?")
rows = [row for group in job.outcomes.values() for row in group]
for row in rows:
self.assertFalse(os.path.isabs(row["path"]), row)
self.assertNotIn(":", row["path"])
self.assertEqual(job.outcomes["added"][0]["path"], "a.png")
self.assertEqual(job.outcomes["known"][0]["path"], "sub/b.png")
self.assertEqual(job.outcomes["review"][0]["path"], "c.png")
self.assertEqual(job.outcomes["known"][0]["detail"], "")
self.assertNotIn(absolute, "\n".join(job.log))
def test_record_rejects_unknown_group(self):
job = Job("scan", "Scan", groups=SCAN_GROUP_NAMES)
with self.assertRaises(ValueError):
job.record("bogus", {"path": "a.png", "root": "output"}, "x")
def test_detail_and_progress(self):
job = Job("scan", "Scan", scope={"mode": "all"}, groups=SCAN_GROUP_NAMES)
job.update(progress=150, status_line="Processing file 3 of 4", processed=3)
job.update(total=4)
job.record("known", {"path": "a.png", "root": "output"}, "= a.png")
detail = job.detail()
self.assertEqual(detail["progress"], 100)
self.assertEqual(detail["status_line"], "Processing file 3 of 4")
self.assertEqual((detail["processed"], detail["total"]), (3, 4))
self.assertEqual(detail["counters"]["known"], 1)
self.assertEqual(detail["outcomes"]["known"][0]["path"], "a.png")
self.assertEqual(len(detail["log"]), 1)
detail["outcomes"]["known"].clear() # a copy, not the live list
self.assertEqual(len(job.outcomes["known"]), 1)
job.publish({"type": "progress", "progress": -5})
self.assertEqual(job.summary()["progress"], 0)
def test_scan_groups_follow_the_counter_mapping(self):
self.assertEqual(
SCAN_GROUPS,
{
"added": "added",
"linked": "added",
"enriched": "added",
"unchanged": "known",
"ambiguous": "review",
"conflict": "review",
"unresolved": "review",
"failed": "errors",
},
)
self.assertEqual(
SCAN_GROUP_NAMES, ("added", "known", "review", "no_metadata", "errors")
)
if __name__ == "__main__":
unittest.main()
+3 -3
View File
@@ -115,9 +115,9 @@ class TestPairScan(AdminAPITestCase):
self.assertEqual(self.api._ingest_scanned_file(path, meta), "unchanged")
async def test_cancel_finishes_at_safe_boundary_and_restart_is_idempotent(self):
from py.api.scan_job import ScanJob
from py.api.jobs import Job
job = ScanJob()
job = Job("scan", "Scan")
self.api._scan_job = job
response = await self.client.post("/prompt_manager/scan/cancel")
self.assertEqual(response.status, 200)
@@ -128,7 +128,7 @@ class TestPairScan(AdminAPITestCase):
self.assertEqual(counts["processed"], 0)
self.assertEqual(self.api.db.list_prompt_pairs()["total"], 0)
job.finish()
self.api._scan_job = ScanJob()
self.api._scan_job = Job("scan", "Scan")
first = self.api._ingest_scan_batch_sync(
[(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})]
)
-92
View File
@@ -1,92 +0,0 @@
"""Unit tests for py/api/scan_job.py: fan-out of scan progress to SSE subscribers."""
import asyncio
import contextlib
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py.api.scan_job import ScanJob # noqa: E402
async def _collect(job, stop_after=None):
seen = []
async with contextlib.aclosing(job.events()) as events:
async for payload in events:
seen.append(payload)
if stop_after is not None and len(seen) >= stop_after:
break
return seen
class TestScanJob(unittest.IsolatedAsyncioTestCase):
async def test_new_job_is_running_with_no_event(self):
job = ScanJob()
self.assertTrue(job.running)
self.assertIsNone(job.last_event)
async def test_subscriber_receives_events_until_finish(self):
job = ScanJob()
consumer = asyncio.ensure_future(_collect(job))
await asyncio.sleep(0)
job.publish({"type": "progress", "progress": 10})
job.publish({"type": "complete"})
job.finish()
self.assertEqual(
await asyncio.wait_for(consumer, 2),
[{"type": "progress", "progress": 10}, {"type": "complete"}],
)
self.assertFalse(job.running)
async def test_late_subscriber_gets_latest_snapshot_then_live_events(self):
job = ScanJob()
job.publish({"type": "progress", "progress": 10})
job.publish({"type": "progress", "progress": 40})
consumer = asyncio.ensure_future(_collect(job))
await asyncio.sleep(0)
job.publish({"type": "complete"})
job.finish()
self.assertEqual(
await asyncio.wait_for(consumer, 2),
[{"type": "progress", "progress": 40}, {"type": "complete"}],
)
async def test_subscriber_after_finish_gets_only_the_final_event(self):
job = ScanJob()
job.publish({"type": "progress", "progress": 99})
job.publish({"type": "complete"})
job.finish()
self.assertEqual(await _collect(job), [{"type": "complete"}])
async def test_two_subscribers_both_see_every_event(self):
job = ScanJob()
first = asyncio.ensure_future(_collect(job))
second = asyncio.ensure_future(_collect(job))
await asyncio.sleep(0)
job.publish({"type": "progress", "progress": 1})
job.finish()
self.assertEqual(await first, [{"type": "progress", "progress": 1}])
self.assertEqual(await second, [{"type": "progress", "progress": 1}])
async def test_departed_subscriber_is_forgotten(self):
job = ScanJob()
job.publish({"type": "progress", "progress": 1})
await _collect(job, stop_after=1)
self.assertEqual(job.subscriber_count, 0)
job.publish({"type": "complete"}) # must not raise or queue anywhere
job.finish()
if __name__ == "__main__":
unittest.main()
+1 -1
View File
@@ -1,7 +1,7 @@
/**
* State of the output-image scan as seen by the dashboard.
*
* The scan runs on the server as a background job (see py/api/scan_job.py) and
* The scan runs on the server as a background job (see py/api/jobs.py) and
* reports progress over SSE. This module turns those events into a small
* immutable state object and answers the UI questions that depend on it:
* what dismissing the modal should do, what the minimized pill should say, and