fix(threads): refuse new previews before the monitor stops; isolate the exit child (#301)

Review findings:
- The teardown makes the API refuse new previews before it detaches the
  monitor's callbacks: an image in hand that already read the preview
  queue no longer starts a preview worker during exit (Python 3.12+
  refuses that thread and logs a traceback). New
  refuse_auto_thumbnails().
- The teardown imports its bounds inside each step and falls back to the
  standard logger, so a failing import never skips the other steps.
- The exit child keeps its library and log in its own folder
  (PROMPT_MANAGER_DATA_DIR), reports watchdog's threads too, fails with
  the child's output on a hang, and rejects "Error while stopping".
- TESTING.md: "Background threads at exit".

TESTING.md

@@ -275,6 +275,72 @@ Start fixing as soon as a failed log shows the cause, but push the fix only afte
 the run for that commit has finished: pull-request runs use `cancel-in-progress`,
 so a push cancels the running jobs and can hide a second failure.
 
+## Background threads at exit
+
+PromptManager's own background threads stop through one package teardown,
+`_stop_background_work()` in `__init__.py` (#301, #302). It runs once and
+works against one 10 s deadline. Each step gets the smaller of its own bound
+and the time left. The steps, in order:
+
+1. Detach the monitor's callbacks (Needs review recorder, preview queue,
+   Images listing listener).
+2. `ImageMonitor.stop_monitoring` closes the image handler, which joins the
+   image worker (5 s). It then calls watchdog's `observer.stop()`, which is
+   not bounded here, and joins the observer (5 s).
+3. `close_auto_thumbnails(final=True)` stops the preview worker (30 s) and
+   kills ffmpeg.
+4. `PromptTracker.shutdown` stops the prompt cleanup thread (2 s).
+
+A failing step, or a second Ctrl+C, is logged with its traceback and does not
+skip the later steps. A clean run logs one INFO line to PromptManager's log:
+`Background work stopped in 0.4 s: monitor=ok thumbnails=ok tracker=ok`. If
+any part is still running, one WARNING goes to the console:
+`Background work not stopped in time (10.0 s): monitor`.
+
+ComfyUI never runs aiohttp's `on_cleanup` hooks, because it never calls
+`runner.cleanup()`. So `atexit` is the live host's path. It runs on Ctrl+C
+(`main.py` catches `KeyboardInterrupt` and returns) and on a normal exit.
+The teardown is registered after the execution history hook, so at exit
+(newest first) it runs before the history writer's 2 s dispose.
+It also runs from `on_cleanup` on hosts that do run it (our test servers) and
+from the tests' package unload (`cleanup_package`), and it removes its own
+`atexit` callback. Nothing runs when the process is killed: SIGTERM, a
+ComfyUI-Manager restart through `os.execv`, a closed Windows console, or
+ComfyUI Desktop stopping its server. SQLite's WAL journal covers those.
+
+Stop contract: `PromptTracker.shutdown(timeout)`,
+`ImageGenerationHandler.shutdown(timeout)` (restartable),
+`ImageGenerationHandler.close(timeout)` (final: `schedule()` refuses and a
+sealed run's re-check is dropped) and `ImageMonitor.stop_monitoring(timeout)`
+each signal their thread and join it within the bound. They return True only
+when it ended, and log a warning naming the part when it did not. Called on
+their own thread, they only signal.
+
+Test rules:
+
+- Every test that constructs a real `PromptTracker`, `ImageGenerationHandler`
+  or `ImageMonitor` ends with its final stop and an assertion
+  (`tests/background_stops.py`: `close_handler`, `stop_monitor`,
+  `stop_tracker`). The stop comes before the folder and databases are removed.
+  `tearDown` runs before `addCleanup` callbacks, so a class whose `tearDown`
+  removes files stops the thread there.
+- A test waits for a state through an `Event` set by the code under test, or
+  by a wrapper around it, with a 30 s ceiling. It never sleeps or polls for
+  that state.
+- A test that asserts "still running" first waits for an `entered` event
+  inside the held call, then stops with `timeout=0.05`.
+- A test that asserts a history writer's `dispose()`/`close()` is True right
+  after queuing work drains first (`drain_writer`), or holds the write itself
+  (#370).
+- Patch the module attribute of the code under test (for example
+  `utils.prompt_tracker.time`, or the package's `_clock`), or patch with
+  `wraps=`. Never replace a `time` or `atexit` function for the whole
+  process.
+- `tests/test_package_exit.py` runs a child process that imports the
+  package, leaves the image worker alive and returns. Its `atexit` reporter
+  checks the teardown's result. This is the interpreter path ComfyUI takes
+  after Ctrl+C, on all three OS.
+
 ## Performance budgets
 
 A spec performance budget gets two tests:

__init__.py

@@ -250,6 +250,18 @@ def _teardown_step(name, stop, logger):
         return False
 
 
+def _teardown_logger():
+    """PromptManager's init logger; the standard one if that cannot load."""
+    try:
+        from .utils.logging_config import get_logger
+
+        return get_logger("prompt_manager.init")
+    except Exception:
+        import logging
+
+        return logging.getLogger("prompt_manager.init")
+
+
 def _stop_background_work():
     """Stop the image monitor, the preview worker and the prompt cleanup
     thread against one TEARDOWN_SECONDS deadline, once.
@@ -259,14 +271,10 @@ def _stop_background_work():
     copy of the first result.
     """
     global _background_work_result
-    from .utils.image_monitor import OBSERVER_JOIN_SECONDS
-    from .utils.logging_config import get_logger
-    from .utils.prompt_tracker import TRACKER_JOIN_SECONDS
-
     with _background_work_lock:
         if _background_work_result is not None:
             return dict(_background_work_result)
-        logger = get_logger("prompt_manager.init")
+        logger = _teardown_logger()
         began = _clock()
         deadline = began + TEARDOWN_SECONDS
 
@@ -282,8 +290,14 @@ def _stop_background_work():
         if monitor is not None:
 
             def stop_monitor():
+                from .utils.image_monitor import OBSERVER_JOIN_SECONDS
+
                 # Callbacks first: an image in hand queues no preview into a
-                # closing preview worker.
+                # closing preview worker. The worker may have read the queue
+                # callback already, so the API refuses new previews too (at
+                # exit, Python 3.12+ refuses to start their worker thread).
+                if thumbnails is not None:
+                    thumbnails.refuse_auto_thumbnails()
                 monitor.set_review_recorder(None)
                 monitor.set_thumbnail_queue(None)
                 monitor.set_output_listener(None)
@@ -293,20 +307,26 @@ def _stop_background_work():
                 "the image monitor", stop_monitor, logger
             )
         if thumbnails is not None:
-            from .py.api.thumbnail_auto import CLOSE_WAIT_SECONDS
+
+            def close_previews():
+                from .py.api.thumbnail_auto import CLOSE_WAIT_SECONDS
+
+                return thumbnails.close_auto_thumbnails(
+                    final=True, timeout=left(CLOSE_WAIT_SECONDS)
+                )
 
             result["thumbnails"] = _teardown_step(
-                "the preview worker",
-                lambda: thumbnails.close_auto_thumbnails(
-                    final=True, timeout=left(CLOSE_WAIT_SECONDS)
-                ),
-                logger,
+                "the preview worker", close_previews, logger
             )
         if tracker is not None:
+
+            def stop_tracker():
+                from .utils.prompt_tracker import TRACKER_JOIN_SECONDS
+
+                return tracker.shutdown(left(TRACKER_JOIN_SECONDS))
+
             result["tracker"] = _teardown_step(
-                "the prompt cleanup thread",
-                lambda: tracker.shutdown(left(TRACKER_JOIN_SECONDS)),
-                logger,
+                "the prompt cleanup thread", stop_tracker, logger
             )
         elapsed = _clock() - began
         failed = [name for name, ok in result.items() if not ok]

py/api/thumbnail_auto.py

@@ -619,6 +619,12 @@ class AutoThumbnailMixin:
                 )
             return self._auto_thumbnailer
 
+    def refuse_auto_thumbnails(self):
+        """Refuse every later capture without stopping the worker yet (the
+        package teardown, before ``close_auto_thumbnails(final=True)``)."""
+        with AutoThumbnailMixin._auto_thumbnailer_lock:
+            self._auto_thumbnails_shut = True
+
     def close_auto_thumbnails(self, final=False, timeout=None):
         """Stop the capture worker and the view pool, and forget them.
 

tests/test_init_package.py

@@ -18,6 +18,7 @@ import threading
 import types
 import unittest
 import unittest.mock as mock
+from pathlib import Path
 
 from aiohttp import web
 
@@ -678,6 +679,39 @@ class TestPackageTeardown(PackageTeardownCase):
         self.assertIsNone(monitor.thumbnail_queue)
         self.assertIsNone(monitor.output_listener)
 
+    def test_an_image_in_hand_queues_no_preview_during_teardown(self):
+        package = self.load()
+        monitor = package._global_image_monitor
+        real_stop = monitor.stop_monitoring
+        image = os.path.join(self.output_dir, "in_hand.png")
+        from PIL import Image
+
+        Image.new("RGB", (8, 8), "red").save(image)
+        seen = {}
+
+        def stop_with_an_image_in_hand(timeout):
+            # The worker read ``thumbnail_queue`` before the detach and now
+            # hands its finished image over (at exit, Python 3.12+ refuses a
+            # new thread here and logs a traceback).
+            seen["queued"] = package.api.queue_capture_thumbnail(image)
+            seen["worker"] = package.api._auto_thumbnailer
+            return real_stop(timeout)
+
+        with (
+            mock.patch.object(
+                package.api,
+                "_get_all_output_dirs",
+                return_value=[Path(os.path.realpath(self.output_dir))],
+            ),
+            mock.patch.object(
+                monitor, "stop_monitoring", side_effect=stop_with_an_image_in_hand
+            ),
+        ):
+            result = package._stop_background_work()
+        self.assertEqual(seen["queued"], "skipped")
+        self.assertIsNone(seen["worker"])
+        self.assertTrue(all(result.values()), result)
+
     def test_teardown_removes_its_atexit_callback(self):
         package = self.load()
         with mock.patch("atexit.unregister", wraps=atexit.unregister) as unregister:

tests/test_package_exit.py

@@ -35,12 +35,16 @@ CHILD = textwrap.dedent("""
     state = {}
 
     def report():
+        from watchdog.utils import BaseThread
+
         package = state.get("package")
         result = getattr(package, "_background_work_result", None)
+        # The image worker, the tracker, and watchdog's observer and emitters.
         alive = sorted(
             t.name
             for t in threading.enumerate()
             if t.name in ("PromptManagerImageWorker", "PromptManagerPromptCleanup")
+            or isinstance(t, BaseThread)
         )
         print("PM_EXIT_REPORT " + json.dumps({"result": result, "alive": alive}))
         sys.stdout.flush()
@@ -83,16 +87,26 @@ class TestPackageExit(unittest.TestCase):
                 handle.write(CHILD)
             env = dict(os.environ)
             env.pop("PROMPT_MANAGER_CONFIG_PATH", None)
+            # The child's library and log file live in its own folder, never
+            # in the checkout's logs/ (which this runner may hold open) or a
+            # data folder the caller set.
+            env["PROMPT_MANAGER_DATA_DIR"] = tmp
             env["PYTHONIOENCODING"] = "utf-8"
-            completed = subprocess.run(
-                [sys.executable, script, ROOT, tmp],
-                cwd=ROOT,
-                env=env,
-                capture_output=True,
-                timeout=CHILD_CEILING_SECONDS,
-                encoding="utf-8",
-                errors="replace",
-            )
+            try:
+                completed = subprocess.run(
+                    [sys.executable, script, ROOT, tmp],
+                    cwd=ROOT,
+                    env=env,
+                    capture_output=True,
+                    timeout=CHILD_CEILING_SECONDS,
+                    encoding="utf-8",
+                    errors="replace",
+                )
+            except subprocess.TimeoutExpired as hung:
+                self.fail(
+                    f"the child did not exit within {CHILD_CEILING_SECONDS} s\n"
+                    f"stdout:\n{hung.stdout}\nstderr:\n{hung.stderr}"
+                )
             # Leaving the block removes the folder: on Windows that fails
             # if the child left a handle open, which it cannot once it exited,
             # so the check here is that the child exited at all.
@@ -105,6 +119,8 @@ class TestPackageExit(unittest.TestCase):
         ]
         self.assertEqual(len(reports), 1, detail)
         report = json.loads(reports[0])
+        # "thumbnails" is True here because no preview worker started; its
+        # close result is tested in test_thumbnail_auto and test_init_package.
         self.assertEqual(
             report["result"],
             {"monitor": True, "thumbnails": True, "tracker": True},
@@ -117,6 +133,7 @@ class TestPackageExit(unittest.TestCase):
             "Traceback",
             "Background work not stopped in time",
             "still running after",
+            "Error while stopping",
         ):
             self.assertNotIn(text, completed.stderr, detail)