From dc192d58368a557422e37811241ca2807bb92775 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Thu, 8 Oct 2026 16:37:15 -0700 Subject: [PATCH] fix: output scan runs as a background job that survives the browser going away Closing the dashboard tab or clicking outside the scan modal raised ClientConnectionResetError out of scan_images and abandoned the scan mid-batch, leaving a half-scanned database. Two overlapping scans could also run at once. - py/api/scan_job.py: ScanJob fans progress out to any number of SSE subscribers (latest snapshot first, then live events) - POST /prompt_manager/scan starts the scan as an asyncio task, or attaches to the one already running; a write to a dead socket ends that stream quietly while the scan continues - GET /prompt_manager/scan/status reports the running flag and last event - dashboard: dismissing the modal while a scan runs minimizes it to a progress pill, the page re-attaches to a running scan after a reload, a lost stream reconnects, and leaving the page mid-scan asks first --- py/api/__init__.py | 2 + py/api/admin.py | 106 +++++++++++----- py/api/scan_job.py | 61 +++++++++ tests/js/scan-progress.test.js | 78 ++++++++++++ tests/test_api_admin.py | 121 ++++++++++++++++++ tests/test_scan_job.py | 92 ++++++++++++++ web/admin.html | 7 ++ web/js/admin.js | 224 +++++++++++++++++++++------------ web/js/scan-progress.js | 109 ++++++++++++++++ 9 files changed, 685 insertions(+), 115 deletions(-) create mode 100644 py/api/scan_job.py create mode 100644 tests/js/scan-progress.test.js create mode 100644 tests/test_scan_job.py create mode 100644 web/js/scan-progress.js diff --git a/py/api/__init__.py b/py/api/__init__.py index bc1e715..4023f75 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -274,6 +274,8 @@ 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. + self._scan_job = None # Run cleanup on initialization to remove any existing duplicates try: diff --git a/py/api/admin.py b/py/api/admin.py index bf9581d..ea8362a 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -1,6 +1,7 @@ """Admin and maintenance API routes for PromptManager.""" import asyncio +import contextlib import datetime import hashlib import json @@ -11,6 +12,8 @@ from pathlib import Path from aiohttp import web +from .scan_job import ScanJob + try: from ...utils.validators import validate_result_timeout except ImportError: @@ -336,6 +339,10 @@ class AdminRoutesMixin: async def scan_images_route(request): return await self.scan_images(request) + @routes.get("/prompt_manager/scan/status") + async def scan_status_route(request): + return await self.scan_status(request) + LONG_JOB_DUPLICATES = "duplicates" async def scan_duplicates_endpoint(self, request): @@ -1105,9 +1112,43 @@ class AdminRoutesMixin: async def scan_images(self, request): """Scan ComfyUI output images for prompt metadata and add them to the database. - Streams progress as server-sent events; all filesystem and database - work runs through the executor. + 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. """ + job = self._scan_job + if job is None or job.finished: + job = self._start_scan_job() + return await self._stream_scan_job(request, job) + + async def scan_status(self, request): + """Whether an output scan is running, and its latest progress event.""" + job = self._scan_job + return web.json_response( + { + "success": True, + "running": bool(job and job.running), + "last_event": job.last_event if job else None, + } + ) + + def _start_scan_job(self): + """Create the job record synchronously, then run the scan as a task.""" + job = ScanJob() + 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() + + async def _stream_scan_job(self, request, job): + """Stream ``job`` as SSE; stop quietly if the client goes away.""" response = web.StreamResponse( status=200, reason="OK", @@ -1118,27 +1159,30 @@ class AdminRoutesMixin: }, ) await response.prepare(request) - - async for chunk in self._scan_images_events(): - await response.write(chunk.encode("utf-8")) - - await response.write_eof() + try: + async with contextlib.aclosing(job.events()) as events: + async for payload in events: + await response.write(_sse(payload).encode("utf-8")) + await response.write_eof() + except Exception as e: # the socket closed under us; the scan goes on + self.logger.info( + f"Scan progress client disconnected ({e}); " + "the scan continues in the background" + ) return response async def _scan_images_events(self): - """Yield SSE strings while scanning every configured output directory.""" + """Yield progress payloads while scanning every configured output directory.""" 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 _sse( - { - "type": "error", - "message": "No output directories found. " - "Configure scan directories in Settings.", - } - ) + yield { + "type": "error", + "message": "No output directories found. " + "Configure scan directories in Settings.", + } return dir_names = [self._public_path(d) for d in output_dirs] @@ -1162,18 +1206,14 @@ class AdminRoutesMixin: f"found={counts['found']}, new_prompts_added={counts['added']}, " f"images_linked_to_existing={counts['linked']}" ) - yield _sse({"type": "complete", **counts, "directories": dir_names}) + yield {"type": "complete", **counts, "directories": dir_names} except Exception: self.logger.exception("Scan error") - yield _sse( - { - "type": "error", - "message": ( - "An internal error occurred. Check server logs for details." - ), - } - ) + yield { + "type": "error", + "message": "An internal error occurred. Check server logs for details.", + } async def _scan_batches(self, media_files, counts): """Process ``media_files`` in executor batches, updating ``counts`` in place. @@ -1200,17 +1240,15 @@ class AdminRoutesMixin: @staticmethod def _scan_progress(progress, status, counts=None): - """SSE progress event for the output scan.""" + """Progress event payload for the output scan.""" counts = counts or {} - return _sse( - { - "type": "progress", - "progress": progress, - "status": status, - "processed": counts.get("processed", 0), - "found": counts.get("found", 0), - } - ) + return { + "type": "progress", + "progress": progress, + "status": status, + "processed": counts.get("processed", 0), + "found": counts.get("found", 0), + } def _collect_output_media_sync(self, output_dirs): """Collect all media files from the output directories (blocking I/O).""" diff --git a/py/api/scan_job.py b/py/api/scan_job.py new file mode 100644 index 0000000..d3eb9ce --- /dev/null +++ b/py/api/scan_job.py @@ -0,0 +1,61 @@ +"""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.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) diff --git a/tests/js/scan-progress.test.js b/tests/js/scan-progress.test.js new file mode 100644 index 0000000..6f0cfb9 --- /dev/null +++ b/tests/js/scan-progress.test.js @@ -0,0 +1,78 @@ +// Run with: node --test tests/js/ +const test = require("node:test"); +const assert = require("node:assert/strict"); +const { + initialState, + reduceScanEvent, + fromStatus, + dismissAction, + progressLabel, +} = require("../../web/js/scan-progress.js"); + +test("initialState is idle with zero counts", () => { + assert.deepEqual(initialState(), { + status: "idle", progress: 0, statusText: "", processed: 0, found: 0, added: 0, linked: 0, message: "", + }); +}); + +test("a progress event moves the scan to running and copies the counts", () => { + const before = initialState(); + const after = reduceScanEvent(before, { + type: "progress", progress: 42, status: "Processing file 84/200...", processed: 84, found: 12, + }); + assert.equal(after.status, "running"); + assert.equal(after.progress, 42); + assert.equal(after.statusText, "Processing file 84/200..."); + assert.equal(after.processed, 84); + assert.equal(after.found, 12); + assert.equal(before.status, "idle", "reducer must not mutate its input"); +}); + +test("a complete event finishes the scan with the final totals", () => { + const running = reduceScanEvent(initialState(), { type: "progress", progress: 90, status: "x", processed: 9, found: 1 }); + const done = reduceScanEvent(running, { type: "complete", processed: 10, found: 3, added: 2, linked: 1 }); + assert.equal(done.status, "done"); + assert.equal(done.progress, 100); + assert.deepEqual([done.processed, done.found, done.added, done.linked], [10, 3, 2, 1]); +}); + +test("an error event fails the scan and keeps the server message", () => { + const failed = reduceScanEvent(initialState(), { type: "error", message: "No output directories found." }); + assert.equal(failed.status, "failed"); + assert.equal(failed.message, "No output directories found."); +}); + +test("unknown events leave the state untouched", () => { + const state = reduceScanEvent(initialState(), { type: "progress", progress: 5, status: "s", processed: 1, found: 0 }); + assert.equal(reduceScanEvent(state, { type: "mystery" }), state); + assert.equal(reduceScanEvent(state, null), state); +}); + +test("fromStatus rebuilds a running scan from the status endpoint", () => { + const state = fromStatus({ + running: true, + last_event: { type: "progress", progress: 30, status: "Processing file 3/10...", processed: 3, found: 1 }, + }); + assert.equal(state.status, "running"); + assert.equal(state.progress, 30); + assert.equal(state.processed, 3); +}); + +test("fromStatus is idle when nothing is running, even if a last event exists", () => { + assert.equal(fromStatus({ running: false, last_event: { type: "complete", processed: 10 } }).status, "idle"); + assert.equal(fromStatus({ running: false, last_event: null }).status, "idle"); + assert.equal(fromStatus(null).status, "idle"); +}); + +test("dismissing the modal minimizes a running scan and closes otherwise", () => { + const running = reduceScanEvent(initialState(), { type: "progress", progress: 1, status: "s", processed: 0, found: 0 }); + assert.equal(dismissAction(running), "minimize"); + assert.equal(dismissAction(initialState()), "close"); + assert.equal(dismissAction(reduceScanEvent(running, { type: "complete" })), "close"); +}); + +test("progressLabel summarises a running scan for the minimized pill", () => { + const running = reduceScanEvent(initialState(), { type: "progress", progress: 42, status: "s", processed: 84, found: 12 }); + assert.equal(progressLabel(running), "Scanning 42% · 84 files · 12 prompts"); + assert.equal(progressLabel(initialState()), ""); +}); diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index 2ac3d1a..1f5b2e2 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -4,6 +4,7 @@ Uses aiohttp's test client against a PromptManagerAPI wired to a temporary database and a stub ``folder_paths`` module, so the tests run without ComfyUI. """ +import asyncio import json import os import shutil @@ -1508,5 +1509,125 @@ class TestOutputScanInternals(AdminAPITestCase): self.assertNotIn("boom", body) +class TestOutputScanIsABackgroundJob(AdminAPITestCase): + """The output scan keeps running when the browser goes away, and a second + request attaches to the running scan instead of starting another one.""" + + def _gate_collect(self): + """Block the media collection step until the returned event is set.""" + 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 test_scan_survives_client_disconnect_and_finishes(self): + from unittest.mock import AsyncMock, patch + + _write_png(self.output_dir / "gen.png", prompt_text="a red square") + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] + + with self.assertNoLogs("aiohttp.server", level="ERROR"): + with patch.object( + web.StreamResponse, + "write", + AsyncMock(side_effect=ConnectionResetError("closing transport")), + ): + resp = await self.client.request("POST", "/prompt_manager/scan") + self.assertEqual(resp.status, 200) + await resp.release() + await asyncio.wait_for(self.api._scan_job.task, 10) + + prompts = self.api.db.get_recent_prompts(limit=10)["prompts"] + self.assertEqual([p["text"] for p in prompts], ["a red square"]) + self.assertEqual(self.api._scan_job.last_event["type"], "complete") + self.assertFalse(self.api._scan_job.running) + + async def test_second_request_attaches_to_the_running_scan(self): + _write_png(self.output_dir / "gen.png", prompt_text="a red square") + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] + gate, calls = self._gate_collect() + + first = asyncio.ensure_future( + self.client.request("POST", "/prompt_manager/scan") + ) + await asyncio.sleep(0.05) + second = await asyncio.wait_for( + self.client.request("POST", "/prompt_manager/scan"), 5 + ) + self.assertEqual(second.status, 200) + self.assertIn("text/event-stream", second.headers["Content-Type"]) + + gate.set() + first_body = await (await first).text() + second_body = await second.text() + + self.assertEqual(len(calls), 1) + self.assertIn('"type": "complete"', first_body) + self.assertIn('"type": "complete"', second_body) + self.assertIn('"added": 1', second_body) + + async def test_status_reports_running_scan_and_its_last_event(self): + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] + gate, _ = self._gate_collect() + + first = asyncio.ensure_future( + self.client.request("POST", "/prompt_manager/scan") + ) + await asyncio.sleep(0.05) + resp = await self.client.request("GET", "/prompt_manager/scan/status") + running = await resp.json() + + gate.set() + await (await first).text() + resp = await self.client.request("GET", "/prompt_manager/scan/status") + finished = await resp.json() + + self.assertEqual(resp.status, 200) + self.assertTrue(running["success"]) + self.assertTrue(running["running"]) + self.assertEqual(running["last_event"]["type"], "progress") + self.assertFalse(finished["running"]) + self.assertEqual(finished["last_event"]["type"], "complete") + + async def test_status_before_any_scan_is_idle(self): + resp = await self.client.request("GET", "/prompt_manager/scan/status") + data = await resp.json() + self.assertEqual(resp.status, 200) + self.assertEqual(data, {"success": True, "running": False, "last_event": None}) + + async def test_scan_after_a_finished_one_starts_fresh(self): + _write_png(self.output_dir / "gen.png", prompt_text="a red square") + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] + gate, calls = self._gate_collect() + gate.set() + + await (await self.client.request("POST", "/prompt_manager/scan")).text() + body = await (await self.client.request("POST", "/prompt_manager/scan")).text() + + self.assertEqual(len(calls), 2) + self.assertIn('"linked": 1', body) + + async def test_failed_scan_is_not_left_running(self): + GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] + self.api._collect_output_media_sync = _raise + + await (await self.client.request("POST", "/prompt_manager/scan")).text() + data = await ( + await self.client.request("GET", "/prompt_manager/scan/status") + ).json() + + self.assertFalse(data["running"]) + self.assertEqual(data["last_event"]["type"], "error") + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_scan_job.py b/tests/test_scan_job.py new file mode 100644 index 0000000..dab064b --- /dev/null +++ b/tests/test_scan_job.py @@ -0,0 +1,92 @@ +"""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() diff --git a/web/admin.html b/web/admin.html index f988b4a..dc8b073 100644 --- a/web/admin.html +++ b/web/admin.html @@ -11,6 +11,7 @@ + @@ -823,6 +824,12 @@ + + +