fix(threads): stop background work at exit, on cleanup and on the tests' unload (#301)

- __init__.py: _stop_background_work() detaches the monitor's callbacks,
  stops the monitor, closes the preview worker and shuts the prompt
  cleanup thread down against one 10 s deadline, once. It runs from
  atexit (ComfyUI never runs aiohttp's cleanup hooks), from on_cleanup
  where a host runs it, and from the tests' package unload; it removes
  its own atexit callback. Registered after the history hook, so at exit
  it runs before the history dispose.
- AutoThumbnailer.close/close_auto_thumbnails take a timeout and return
  whether the worker ended.
- Tests: teardown result, once-only, on_cleanup, atexit order, a failing
  step or a second Ctrl+C never skipping the tracker, one shared
  deadline, three repeated loads leaving no thread, and a child process
  that exits normally with the image worker alive.

tests/test_init_package.py

@@ -4,6 +4,8 @@ The package is loaded the way ComfyUI loads it (``__init__`` with relative
 imports) with ComfyUI's ``server`` and ``folder_paths`` modules stubbed.
 """
 
+import asyncio
+import atexit
 import contextlib
 import importlib.util
 import io
@@ -40,13 +42,20 @@ class FakePromptQueue:
         return self.results.pop(0) if self.results else None
 
 
-def fake_server_module():
-    """A ``server`` module whose PromptServer.instance accepts routes and hooks."""
+def fake_server_module(app=None, queue_task_done=False):
+    """A ``server`` module whose PromptServer.instance accepts routes and hooks.
+
+    ``app`` replaces the default app (one without aiohttp hooks);
+    ``queue_task_done`` gives the queue a ``task_done``, so history registers.
+    """
+    queue = FakePromptQueue()
+    if queue_task_done:
+        queue.task_done = lambda *a, **k: None
     instance = types.SimpleNamespace(
         routes=web.RouteTableDef(),
-        prompt_queue=FakePromptQueue(),
+        prompt_queue=queue,
         handlers=[],
-        app=types.SimpleNamespace(middlewares=[]),
+        app=app if app is not None else types.SimpleNamespace(middlewares=[]),
     )
     instance.add_on_prompt_handler = instance.handlers.append
     module = types.ModuleType("server")
@@ -92,7 +101,7 @@ def _import_package(name, db_path, output_dir, server=None):
             try:
                 spec.loader.exec_module(package)
             except BaseException:
-                cleanup_package(package)
+                cleanup_package(package, require_stopped=False)
                 raise
     finally:
         for key, module in saved.items():
@@ -103,20 +112,34 @@ def _import_package(name, db_path, output_dir, server=None):
     return package, server.PromptServer.instance
 
 
-def cleanup_package(package):
-    """Stop fixture-owned workers before releasing its other resources."""
+def cleanup_package(package, require_stopped=True):
+    """Stop fixture-owned workers before releasing its other resources.
+
+    The package's own teardown runs first (the reload case of #301); with
+    ``require_stopped`` every part it reports must have stopped. The direct
+    stops after it also cover an import that failed before the teardown
+    existed, and a test that made a teardown step fail on purpose.
+    """
+    teardown = getattr(package, "_stop_background_work", None)
+    if callable(teardown):
+        result = teardown()
+        if require_stopped:
+            assert all(result.values()), f"package teardown left work running: {result}"
+    monitor = getattr(package, "_global_image_monitor", None)
+    if monitor is not None:
+        assert monitor.stop_monitoring(), "package image monitor still running"
+    tracker = getattr(package, "_global_prompt_tracker", None)
+    if tracker is not None:
+        assert tracker.shutdown(), "package prompt cleanup thread still running"
     registration = getattr(package, "_history_registration", None)
     if registration is not None:
         assert registration.dispose(3), "package history worker leaked"
-    monitor = getattr(package, "_global_image_monitor", None)
-    if monitor is not None:
-        monitor.stop_monitoring()
     close_every_database()
     sys.modules.pop(package.__name__, None)
 
 
 @contextlib.contextmanager
-def load_package(name, db_path, output_dir, server=None):
+def load_package(name, db_path, output_dir, server=None, require_stopped=True):
     """Dispose package workers before closing its isolated logging handlers."""
     root_logger = logging.getLogger("prompt_manager")
     loggers = [root_logger] + [
@@ -144,7 +167,7 @@ def load_package(name, db_path, output_dir, server=None):
         try:
             yield package, instance
         finally:
-            cleanup_package(package)
+            cleanup_package(package, require_stopped=require_stopped)
     finally:
         current_loggers = [root_logger] + [
             logger
@@ -556,3 +579,208 @@ class TestPackageLoggerIsolation(unittest.TestCase):
         self.assert_logging_restored(
             captured["handlers"], captured["package"]._history_registration
         )
+
+
+PACKAGE_THREAD_NAMES = ("PromptManagerPromptCleanup", "PromptManagerImageWorker")
+
+
+def package_threads():
+    """Live threads this design owns: the tracker, the image worker, and the
+    watchdog observer and emitters (which have no name of their own)."""
+    from watchdog.utils import BaseThread
+
+    return {
+        thread
+        for thread in threading.enumerate()
+        if thread.name in PACKAGE_THREAD_NAMES or isinstance(thread, BaseThread)
+    }
+
+
+def hooked_app():
+    return types.SimpleNamespace(
+        middlewares=[], on_cleanup=[], on_startup=[], on_shutdown=[]
+    )
+
+
+class PackageTeardownCase(unittest.TestCase):
+    """Each test loads its own package (#301): the teardown runs once."""
+
+    def setUp(self):
+        self.tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(self.tmp.cleanup)
+        self.addCleanup(lambda: self.assertEqual(open_handles_under(self.tmp.name), []))
+        self.output_dir = os.path.join(self.tmp.name, "output")
+        os.makedirs(self.output_dir)
+        self.db_path = os.path.join(self.tmp.name, "prompts.db")
+
+    def load(self, server=None, require_stopped=True, suffix=""):
+        name = f"pm_teardown_{self._testMethodName}{suffix}"
+        stack = contextlib.ExitStack()
+        self.addCleanup(stack.close)
+        stack.enter_context(contextlib.redirect_stdout(io.StringIO()))
+        package, _ = stack.enter_context(
+            load_package(
+                name,
+                self.db_path,
+                self.output_dir,
+                server,
+                require_stopped=require_stopped,
+            )
+        )
+        return package
+
+    def busy_worker(self, package, name="busy.png"):
+        """The image worker alive in its processing delay."""
+        handler = package._global_image_monitor.handler
+        self.assertIsNotNone(handler, "the package's monitor did not start")
+        handler.processing_delay = 3600
+        # Outside the watched folder, so the observer does not schedule it first.
+        path = os.path.join(self.tmp.name, name)
+        with open(path, "wb") as image:
+            image.write(b"x" * 16)
+        self.assertTrue(handler.schedule(path))
+        self.assertTrue(handler._worker.is_alive())
+        return handler._worker
+
+
+class TestPackageTeardown(PackageTeardownCase):
+    def test_teardown_stops_every_part_once(self):
+        package = self.load()
+        monitor = package._global_image_monitor
+        observer = monitor.observer
+        tracker_thread = package._global_prompt_tracker.cleanup_thread
+        worker = self.busy_worker(package)
+
+        with self.assertLogs("prompt_manager.init", "INFO") as logs:
+            result = package._stop_background_work()
+
+        self.assertEqual(result, {"monitor": True, "thumbnails": True, "tracker": True})
+        self.assertEqual(list(result), ["monitor", "thumbnails", "tracker"])
+        self.assertIn("Background work stopped in", "\n".join(logs.output))
+        self.assertIn("monitor=ok thumbnails=ok tracker=ok", "\n".join(logs.output))
+        self.assertFalse(worker.is_alive())
+        self.assertFalse(observer.is_alive())
+        self.assertFalse(tracker_thread.is_alive())
+        result["monitor"] = False  # a copy: the stored result is unchanged
+        with mock.patch.object(monitor, "stop_monitoring") as again:
+            self.assertEqual(
+                package._stop_background_work(),
+                {"monitor": True, "thumbnails": True, "tracker": True},
+            )
+        again.assert_not_called()
+
+    def test_teardown_detaches_the_monitor_callbacks(self):
+        package = self.load()
+        monitor = package._global_image_monitor
+        self.assertIsNotNone(monitor.review_recorder)
+        package._stop_background_work()
+        self.assertIsNone(monitor.review_recorder)
+        self.assertIsNone(monitor.thumbnail_queue)
+        self.assertIsNone(monitor.output_listener)
+
+    def test_teardown_removes_its_atexit_callback(self):
+        package = self.load()
+        with mock.patch("atexit.unregister", wraps=atexit.unregister) as unregister:
+            package._stop_background_work()
+        unregister.assert_called_once_with(package._stop_background_work)
+
+    def test_on_cleanup_runs_the_teardown(self):
+        app = hooked_app()
+        package = self.load(fake_server_module(app=app))
+        self.assertIs(app.on_cleanup[0], package._stop_background_work_on_cleanup)
+        self.assertIsNone(package._background_work_result)
+
+        async def run_cleanup():
+            for callback in list(app.on_cleanup):
+                await callback(app)
+
+        asyncio.run(run_cleanup())
+        self.assertEqual(
+            package._background_work_result,
+            {"monitor": True, "thumbnails": True, "tracker": True},
+        )
+
+    def test_teardown_runs_before_history_dispose_at_exit(self):
+        with mock.patch("atexit.register", wraps=atexit.register) as register:
+            package = self.load(fake_server_module(queue_task_done=True))
+        registered = [call.args[0] for call in register.call_args_list]
+        registration = package._history_registration
+        self.assertIsNotNone(registration, "history did not register")
+        self.assertIn(registration.dispose, registered)
+        self.assertIn(package._stop_background_work, registered)
+        # atexit runs newest first: the teardown before the history dispose.
+        self.assertGreater(
+            registered.index(package._stop_background_work),
+            registered.index(registration.dispose),
+        )
+
+
+class TestPackageTeardownFailures(PackageTeardownCase):
+    """A failing step never skips the later ones (design D14, R2)."""
+
+    def assert_tracker_stops_after(self, error):
+        package = self.load(require_stopped=False)
+        monitor = package._global_image_monitor
+        tracker_thread = package._global_prompt_tracker.cleanup_thread
+        with (
+            mock.patch.object(monitor, "stop_monitoring", side_effect=error),
+            self.assertLogs("prompt_manager.init", "WARNING") as logs,
+        ):
+            result = package._stop_background_work()
+        self.assertFalse(result["monitor"])
+        self.assertTrue(result["tracker"])
+        self.assertFalse(tracker_thread.is_alive())
+        text = "\n".join(logs.output)
+        self.assertIn("Background work not stopped in time", text)
+        self.assertRegex(text, r"not stopped in time \([0-9.]+ s\): monitor$")
+        self.assertIn("Traceback", text)
+
+    def test_teardown_step_failure_does_not_skip_later_steps(self):
+        self.assert_tracker_stops_after(RuntimeError("stop failed"))
+
+    def test_second_interrupt_does_not_skip_the_tracker(self):
+        self.assert_tracker_stops_after(KeyboardInterrupt())
+
+    def test_teardown_shares_one_deadline(self):
+        package = self.load(require_stopped=False)
+        now = [100.0]
+        seen = []
+
+        def step(timeout, **_):
+            seen.append(timeout)
+            now[0] += 6.0
+            return True
+
+        with (
+            mock.patch.object(package, "_clock", lambda: now[0]),
+            mock.patch.object(
+                package._global_image_monitor, "stop_monitoring", side_effect=step
+            ),
+            mock.patch.object(
+                package.api,
+                "close_auto_thumbnails",
+                side_effect=lambda final=False, timeout=None: step(timeout),
+            ),
+            mock.patch.object(
+                package._global_prompt_tracker, "shutdown", side_effect=step
+            ),
+        ):
+            package._stop_background_work()
+        self.assertEqual(seen, [5.0, 4.0, 0.0])
+
+
+class TestRepeatedLoads(PackageTeardownCase):
+    def test_repeated_loads_leave_no_threads(self):
+        for round_ in range(3):
+            before = package_threads()
+            with contextlib.redirect_stdout(io.StringIO()):
+                with load_package(
+                    f"pm_teardown_repeat_{round_}",
+                    self.db_path,
+                    self.output_dir,
+                ) as (package, _):
+                    self.busy_worker(package, f"busy_{round_}.png")
+                    started = package_threads() - before
+                    self.assertTrue(started, "the package started no thread")
+            left = {t for t in package_threads() - before if t.is_alive()}
+            self.assertEqual(left, set(), f"round {round_}")

tests/test_package_exit.py

@@ -0,0 +1,125 @@
+"""The package teardown runs at a normal interpreter exit (#301, #302).
+
+ComfyUI catches Ctrl+C in ``main.py`` and returns, so its exit is a normal
+return: ``atexit`` callbacks then run, newest first. A child process imports
+the package with ComfyUI stubbed, leaves the image worker alive in its
+processing delay and returns. Its own ``atexit`` reporter, registered before
+the package, runs after the teardown and prints the result as one JSON line.
+"""
+
+import json
+import os
+import subprocess
+import sys
+import tempfile
+import textwrap
+import unittest
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+# The parent never waits longer than this for the child (the hang watchdog
+# allows 120 s per test).
+CHILD_CEILING_SECONDS = 60
+REPORT_PREFIX = "PM_EXIT_REPORT "
+
+CHILD = textwrap.dedent("""
+    import atexit
+    import contextlib
+    import io
+    import json
+    import os
+    import sys
+    import threading
+
+    root, tmp = sys.argv[1], sys.argv[2]
+    sys.path.insert(0, root)
+    state = {}
+
+    def report():
+        package = state.get("package")
+        result = getattr(package, "_background_work_result", None)
+        alive = sorted(
+            t.name
+            for t in threading.enumerate()
+            if t.name in ("PromptManagerImageWorker", "PromptManagerPromptCleanup")
+        )
+        print("PM_EXIT_REPORT " + json.dumps({"result": result, "alive": alive}))
+        sys.stdout.flush()
+
+    # First registered, so it runs last: after the package's teardown and
+    # after the history dispose.
+    atexit.register(report)
+
+    from tests.test_init_package import _import_package, fake_server_module
+
+    output = os.path.join(tmp, "output")
+    with contextlib.redirect_stdout(io.StringIO()):
+        package, _ = _import_package(
+            "pm_exit_child",
+            os.path.join(tmp, "prompts.db"),
+            output,
+            fake_server_module(queue_task_done=True),
+        )
+    state["package"] = package
+    monitor = package._global_image_monitor
+    handler = monitor.handler if monitor is not None else None
+    if handler is None:
+        sys.exit(3)
+    handler.processing_delay = 3600
+    path = os.path.join(tmp, "busy.png")  # outside the watched folder
+    with open(path, "wb") as image:
+        image.write(b"x" * 16)
+    if not handler.schedule(path) or not handler._worker.is_alive():
+        sys.exit(4)
+    # Return normally with the worker alive, as ComfyUI does after Ctrl+C.
+    """)
+
+
+class TestPackageExit(unittest.TestCase):
+    def test_teardown_runs_at_a_normal_exit_with_the_worker_alive(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            os.makedirs(os.path.join(tmp, "output"))
+            script = os.path.join(tmp, "child.py")
+            with open(script, "w", encoding="utf-8") as handle:
+                handle.write(CHILD)
+            env = dict(os.environ)
+            env.pop("PROMPT_MANAGER_CONFIG_PATH", None)
+            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",
+            )
+            # 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.
+        detail = f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}"
+        self.assertEqual(completed.returncode, 0, detail)
+        reports = [
+            line[len(REPORT_PREFIX) :]
+            for line in completed.stdout.splitlines()
+            if line.startswith(REPORT_PREFIX)
+        ]
+        self.assertEqual(len(reports), 1, detail)
+        report = json.loads(reports[0])
+        self.assertEqual(
+            report["result"],
+            {"monitor": True, "thumbnails": True, "tracker": True},
+            detail,
+        )
+        self.assertEqual(report["alive"], [], detail)
+        # The teardown's own warnings. Import-time warnings (the API's
+        # output-folder search finds no ComfyUI here) are not about exit.
+        for text in (
+            "Traceback",
+            "Background work not stopped in time",
+            "still running after",
+        ):
+            self.assertNotIn(text, completed.stderr, detail)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_thumbnail_auto.py

@@ -738,11 +738,27 @@ class TestShutdown(AutoThumbnailCase):
             patch.object(auto, "CLOSE_WAIT_SECONDS", 0),
             self.assertLogs("prompt_manager.api", "WARNING") as logs,
         ):
-            self.api.close_auto_thumbnails()
+            self.assertFalse(self.api.close_auto_thumbnails())
         process.kill.assert_called_once()
         self.assertIn("still running", "\n".join(logs.output))
         release.set()
 
+    async def test_close_reports_an_ended_worker(self):
+        path = make_png(self.output_dir / "a.png", size=(64, 64))
+        self.api.queue_capture_thumbnail(str(path))
+        worker = self.api._auto_thumbnails()._worker
+        await self._idle()
+        loop = asyncio.get_running_loop()
+        closed = await loop.run_in_executor(
+            None, lambda: self.api.close_auto_thumbnails(timeout=WAIT)
+        )
+        self.assertTrue(closed)
+        self.assertFalse(worker.is_alive())
+
+    def test_close_without_a_thumbnailer_returns_true(self):
+        self.assertIsNone(self.api._auto_thumbnailer)
+        self.assertTrue(self.api.close_auto_thumbnails())
+
 
 class TestHistorySharesTheAdmission(AutoThumbnailCase):
     async def _shared_flight(self, name, is_video, first):