From 2cf9f42e73fdd394869f8a119a8e085f3d328cfb Mon Sep 17 00:00:00 2001 From: Vito Date: Fri, 3 Apr 2026 06:22:55 -0700 Subject: [PATCH] fix: diagnostics and monitor respect config, singleton lifecycle (#132) (#133) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: diagnostics, backup, and restore now respect config.json db path (#132) * test: add regression tests for config-aware diagnostics and backup (#132) * fix(test): use fresh event loop to avoid aiohttp test pollution * fix: don't kill singleton image monitor when node instance is garbage collected The __del__ method was calling stop_monitoring() on the shared singleton ImageMonitor, killing it for ALL nodes whenever any single node instance was garbage collected between workflow runs. This caused the monitor to restart on every execution and miss image save events due to race conditions. * test: add regression tests for image monitor singleton lifecycle Ensures cleanup/del never stops the shared monitor, singleton is thread-safe, and observer survives across node instance lifecycles. * style: address PR review — log suppressed exceptions, remove unused imports --- prompt_manager_base.py | 16 +-- py/api/admin.py | 42 ++++-- tests/test_diagnostics_config.py | 189 ++++++++++++++++++++++++++ tests/test_image_monitor_singleton.py | 164 ++++++++++++++++++++++ utils/diagnostics.py | 74 +++++++--- 5 files changed, 445 insertions(+), 40 deletions(-) create mode 100644 tests/test_diagnostics_config.py create mode 100644 tests/test_image_monitor_singleton.py diff --git a/prompt_manager_base.py b/prompt_manager_base.py index 75106f8..9a393fc 100644 --- a/prompt_manager_base.py +++ b/prompt_manager_base.py @@ -219,14 +219,10 @@ class PromptManagerBase: } def cleanup_gallery_system(self): - """Clean up gallery system resources.""" - try: - if hasattr(self, "image_monitor"): - self.image_monitor.stop_monitoring() - self.logger.debug("Gallery system cleaned up") - except Exception as e: - self.logger.error(f"Error cleaning up gallery system: {e}") + """Clean up gallery system resources. - def __del__(self): - """Cleanup when object is destroyed.""" - self.cleanup_gallery_system() + Note: The image monitor is a singleton shared across all node instances, + so we intentionally do NOT stop it here. It should run for the entire + ComfyUI session. Stopping it would kill monitoring for all nodes. + """ + self.logger.debug("Gallery system cleaned up (monitor continues running)") diff --git a/py/api/admin.py b/py/api/admin.py index b4060c6..072deb7 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -497,7 +497,7 @@ class AdminRoutesMixin: # Check database try: - db_path = "prompts.db" + db_path = self.db.model.db_path if os.path.exists(db_path): with sqlite3.connect(db_path) as conn: conn.row_factory = sqlite3.Row @@ -559,12 +559,38 @@ class AdminRoutesMixin: # Check output directories output_dirs = [] - potential_dirs = ["output", "../output", "../../output"] - for dir_path in potential_dirs: - abs_path = os.path.abspath(dir_path) - if os.path.exists(abs_path): - output_dirs.append(abs_path) + # Check user-configured gallery directories from config.json + try: + from ..config import GalleryConfig + + for cfg_dir in GalleryConfig.MONITORING_DIRECTORIES: + abs_path = os.path.abspath(cfg_dir) + if os.path.exists(abs_path) and abs_path not in output_dirs: + output_dirs.append(abs_path) + except Exception as e: + self.logger.debug(f"Could not load gallery config: {e}") + + # Check ComfyUI's own output directory (respects --output-directory) + try: + import folder_paths + + comfyui_output = folder_paths.get_output_directory() + if comfyui_output and os.path.exists(comfyui_output): + abs_path = os.path.abspath(comfyui_output) + if abs_path not in output_dirs: + output_dirs.append(abs_path) + except ImportError: + self.logger.debug( + "folder_paths not available, skipping ComfyUI output dir" + ) + + # Fallback: check common relative paths + if not output_dirs: + for dir_path in ["output", "../output", "../../output"]: + abs_path = os.path.abspath(dir_path) + if os.path.exists(abs_path) and abs_path not in output_dirs: + output_dirs.append(abs_path) results["comfyui_output"] = { "status": "ok" if output_dirs else "warning", @@ -801,7 +827,7 @@ class AdminRoutesMixin: async def backup_database(self, request): """Backup the entire prompts.db database file.""" try: - db_path = "prompts.db" + db_path = self.db.model.db_path if not os.path.exists(db_path): return web.json_response( @@ -894,7 +920,7 @@ class AdminRoutesMixin: cursor = conn.execute("SELECT COUNT(*) as count FROM prompts") prompt_count = cursor.fetchone()["count"] - db_path = "prompts.db" + db_path = self.db.model.db_path backup_path = f"{db_path}.backup_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" if os.path.exists(db_path): diff --git a/tests/test_diagnostics_config.py b/tests/test_diagnostics_config.py new file mode 100644 index 0000000..cc69437 --- /dev/null +++ b/tests/test_diagnostics_config.py @@ -0,0 +1,189 @@ +""" +Tests for issue #132: Diagnostics, backup, and restore must respect config.json. + +Verifies that: +- GalleryDiagnostics resolves db_path from config (not hardcoded "prompts.db") +- The admin API diagnostics/backup/restore endpoints use the configured db_path +- Output directory checks consult GalleryConfig before falling back to heuristics +""" + +import asyncio +import json +import os +import sys +import tempfile +import types +import unittest +from unittest.mock import patch, MagicMock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +EXTENSION_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +# --------------------------------------------------------------------------- +# GalleryDiagnostics path resolution +# --------------------------------------------------------------------------- +class TestGalleryDiagnosticsDbPath(unittest.TestCase): + """GalleryDiagnostics should resolve db_path from config, not hardcode it.""" + + def test_default_resolves_from_config(self): + """GalleryDiagnostics() with no args should read config like PromptDatabase.""" + from utils.diagnostics import GalleryDiagnostics + + with tempfile.TemporaryDirectory() as tmpdir: + custom_path = os.path.join(tmpdir, "custom.db") + fake_config_mod = types.ModuleType("py.config") + fake_pm_config = type( + "PromptManagerConfig", (), {"DEFAULT_DB_PATH": custom_path} + ) + fake_config_mod.PromptManagerConfig = fake_pm_config + + with patch.dict(sys.modules, {"py.config": fake_config_mod}): + diag = GalleryDiagnostics() + + self.assertEqual(diag.db_path, custom_path) + + def test_default_without_config_falls_back_to_extension_root(self): + """Without config, should default to prompts.db in extension root.""" + from utils.diagnostics import GalleryDiagnostics + + diag = GalleryDiagnostics() + expected = os.path.join(EXTENSION_ROOT, "prompts.db") + self.assertEqual(diag.db_path, expected) + + def test_explicit_path_overrides_config(self): + """An explicit db_path argument should be used as-is.""" + from utils.diagnostics import GalleryDiagnostics + + diag = GalleryDiagnostics("/explicit/path.db") + self.assertEqual(diag.db_path, "/explicit/path.db") + + def test_config_relative_path_resolved_to_extension_root(self): + """A relative path from config should resolve against extension root.""" + from utils.diagnostics import GalleryDiagnostics + + fake_config_mod = types.ModuleType("py.config") + fake_pm_config = type( + "PromptManagerConfig", (), {"DEFAULT_DB_PATH": "data/prompts.db"} + ) + fake_config_mod.PromptManagerConfig = fake_pm_config + + with patch.dict(sys.modules, {"py.config": fake_config_mod}): + diag = GalleryDiagnostics() + + expected = os.path.join(EXTENSION_ROOT, "data", "prompts.db") + self.assertEqual(diag.db_path, expected) + + +# --------------------------------------------------------------------------- +# Admin API endpoints use self.db.model.db_path +# --------------------------------------------------------------------------- +class TestAdminEndpointsUseConfigPath(unittest.TestCase): + """Admin diagnostics, backup, and restore must use the configured db_path.""" + + def _make_api_stub(self, db_path): + """Create a minimal object that mimics the admin mixin's self.""" + stub = MagicMock() + stub.db.model.db_path = db_path + stub.logger = MagicMock() + return stub + + def _run_async(self, coro): + """Run a coroutine on a fresh event loop to avoid aiohttp pollution.""" + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete(coro) + finally: + loop.close() + + def test_run_diagnostics_uses_model_db_path(self): + """run_diagnostics should check self.db.model.db_path, not 'prompts.db'.""" + if "server" not in sys.modules: + sys.modules["server"] = MagicMock() + + from py.api.admin import AdminRoutesMixin + + with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as f: + db_file = f.name + + try: + import sqlite3 + + with sqlite3.connect(db_file) as conn: + conn.execute( + "CREATE TABLE prompts (id INTEGER PRIMARY KEY, text TEXT, created_at TEXT)" + ) + conn.execute( + "INSERT INTO prompts (text, created_at) VALUES ('test', '2024-01-01')" + ) + + stub = self._make_api_stub(db_file) + result = self._run_async( + AdminRoutesMixin.run_diagnostics(stub, MagicMock()) + ) + + body = json.loads(result.body) + self.assertTrue(body["success"]) + self.assertEqual(body["diagnostics"]["database"]["status"], "ok") + self.assertEqual(body["diagnostics"]["database"]["prompt_count"], 1) + finally: + os.unlink(db_file) + + def test_run_diagnostics_reports_missing_custom_path(self): + """If the configured db_path doesn't exist, diagnostics should report error.""" + if "server" not in sys.modules: + sys.modules["server"] = MagicMock() + + from py.api.admin import AdminRoutesMixin + + stub = self._make_api_stub("/nonexistent/custom/prompts.db") + + result = self._run_async(AdminRoutesMixin.run_diagnostics(stub, MagicMock())) + body = json.loads(result.body) + self.assertTrue(body["success"]) + self.assertEqual(body["diagnostics"]["database"]["status"], "error") + self.assertIn( + "/nonexistent/custom/prompts.db", + body["diagnostics"]["database"]["message"], + ) + + def test_backup_uses_model_db_path(self): + """backup_database should read from self.db.model.db_path.""" + if "server" not in sys.modules: + sys.modules["server"] = MagicMock() + + from py.api.admin import AdminRoutesMixin + + with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as f: + db_file = f.name + f.write(b"SQLite format 3\x00test data for backup") + + try: + stub = self._make_api_stub(db_file) + result = self._run_async( + AdminRoutesMixin.backup_database(stub, MagicMock()) + ) + + self.assertEqual(result.content_type, "application/octet-stream") + self.assertGreater(len(result.body), 0) + finally: + os.unlink(db_file) + + def test_backup_reports_missing_custom_path(self): + """backup_database should return 404 when configured path doesn't exist.""" + if "server" not in sys.modules: + sys.modules["server"] = MagicMock() + + from py.api.admin import AdminRoutesMixin + + stub = self._make_api_stub("/nonexistent/custom/prompts.db") + + result = self._run_async(AdminRoutesMixin.backup_database(stub, MagicMock())) + body = json.loads(result.body) + self.assertFalse(body["success"]) + self.assertEqual(result.status, 404) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_image_monitor_singleton.py b/tests/test_image_monitor_singleton.py new file mode 100644 index 0000000..08d6fed --- /dev/null +++ b/tests/test_image_monitor_singleton.py @@ -0,0 +1,164 @@ +""" +Tests for image monitor singleton lifecycle. + +Verifies that: +- The singleton image monitor survives node garbage collection +- Node __del__ / cleanup_gallery_system does NOT stop the shared monitor +- Multiple node instances share the same monitor +- The monitor observer stays alive across node lifecycles +""" + +import os +import sys +import threading +import unittest +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +# Mock ComfyUI server before importing anything that touches config +_mock_server = MagicMock() +_mock_server.PromptServer.instance.routes = MagicMock() +sys.modules["server"] = _mock_server + +# Mock comfyui_integration to avoid import issues +sys.modules["utils.comfyui_integration"] = MagicMock() + +from utils.image_monitor import get_image_monitor +import utils.image_monitor as im_mod + + +class TestMonitorSingletonLifecycle(unittest.TestCase): + """The singleton monitor must survive node instance garbage collection.""" + + def setUp(self): + # Reset singleton for each test + im_mod._monitor_instance = None + + def tearDown(self): + im_mod._monitor_instance = None + + def test_get_image_monitor_returns_singleton(self): + """Multiple calls should return the exact same instance.""" + db = MagicMock() + tracker = MagicMock() + + m1 = get_image_monitor(db, tracker) + m2 = get_image_monitor(db, tracker) + + self.assertIs(m1, m2) + + def test_singleton_survives_across_different_callers(self): + """Different db/tracker args on subsequent calls still return the same instance.""" + db1, tracker1 = MagicMock(), MagicMock() + db2, tracker2 = MagicMock(), MagicMock() + + m1 = get_image_monitor(db1, tracker1) + m2 = get_image_monitor(db2, tracker2) + + self.assertIs(m1, m2) + + def test_cleanup_does_not_stop_monitor(self): + """PromptManagerBase.cleanup_gallery_system must NOT stop the monitor.""" + from prompt_manager_base import PromptManagerBase + + with patch.object(PromptManagerBase, "__init__", lambda self, **kw: None): + node = PromptManagerBase() + node.logger = MagicMock() + + mock_monitor = MagicMock() + node.image_monitor = mock_monitor + + node.cleanup_gallery_system() + + mock_monitor.stop_monitoring.assert_not_called() + + def test_del_does_not_stop_monitor(self): + """Node __del__ must NOT stop the singleton monitor.""" + from prompt_manager_base import PromptManagerBase + + with patch.object(PromptManagerBase, "__init__", lambda self, **kw: None): + node = PromptManagerBase() + node.logger = MagicMock() + + mock_monitor = MagicMock() + node.image_monitor = mock_monitor + + # Simulate garbage collection + del node + + mock_monitor.stop_monitoring.assert_not_called() + + def test_observer_stays_alive_after_node_cleanup(self): + """A running observer must remain alive after node cleanup.""" + db = MagicMock() + tracker = MagicMock() + monitor = get_image_monitor(db, tracker) + + # Simulate a running observer + mock_observer = MagicMock() + mock_observer.is_alive.return_value = True + monitor.observer = mock_observer + monitor.handler = MagicMock() + monitor.monitored_directories = ["/fake/output"] + + from prompt_manager_base import PromptManagerBase + + with patch.object(PromptManagerBase, "__init__", lambda self, **kw: None): + node = PromptManagerBase() + node.logger = MagicMock() + node.image_monitor = monitor + + node.cleanup_gallery_system() + + # Observer must still be alive + self.assertIsNotNone(monitor.observer) + self.assertTrue(monitor.observer.is_alive()) + self.assertEqual(monitor.monitored_directories, ["/fake/output"]) + + def test_monitor_start_not_called_when_already_running(self): + """start_monitoring should be a no-op if observer is already active.""" + db = MagicMock() + tracker = MagicMock() + monitor = get_image_monitor(db, tracker) + + # Set up as if already running + mock_observer = MagicMock() + monitor.observer = mock_observer + + monitor.start_monitoring() + + # Should not create a new observer + self.assertIs(monitor.observer, mock_observer) + + +class TestMonitorThreadSafety(unittest.TestCase): + """Singleton creation must be thread-safe.""" + + def setUp(self): + im_mod._monitor_instance = None + + def tearDown(self): + im_mod._monitor_instance = None + + def test_concurrent_get_image_monitor_returns_same_instance(self): + """Multiple threads calling get_image_monitor must get the same instance.""" + results = [] + barrier = threading.Barrier(5) + + def get_monitor(): + barrier.wait() + m = get_image_monitor(MagicMock(), MagicMock()) + results.append(id(m)) + + threads = [threading.Thread(target=get_monitor) for _ in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + + self.assertEqual(len(set(results)), 1, "All threads must get the same instance") + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/diagnostics.py b/utils/diagnostics.py index 60bb7c6..ffed61a 100644 --- a/utils/diagnostics.py +++ b/utils/diagnostics.py @@ -58,15 +58,34 @@ class GalleryDiagnostics: - Additional data specific to the diagnostic """ - def __init__(self, db_path: str = "prompts.db"): + def __init__(self, db_path: str = None): """Initialize the diagnostics system. Args: - db_path: Path to the SQLite database file to diagnose + db_path: Path to the SQLite database file to diagnose. + If None, resolves from config.json or defaults to + prompts.db in the extension root. """ + if db_path is None: + db_path = self._resolve_db_path() self.db_path = db_path self.logger = get_logger("prompt_manager.diagnostics") + @staticmethod + def _resolve_db_path() -> str: + """Resolve database path from config, matching PromptDatabase logic.""" + extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + try: + from py.config import PromptManagerConfig + + db_path = PromptManagerConfig.DEFAULT_DB_PATH + except Exception: + db_path = "prompts.db" + + if not os.path.isabs(db_path): + db_path = os.path.join(extension_root, db_path) + return db_path + def run_full_diagnostic(self) -> Dict[str, Any]: """Run a complete diagnostic check. @@ -268,23 +287,10 @@ class GalleryDiagnostics: output_dirs = [] - # Try to detect ComfyUI output directories - potential_dirs = [ - "output", - "../output", - "../../output", - "ComfyUI/output", - "../ComfyUI/output", - "../../ComfyUI/output", - ] - - for dir_path in potential_dirs: - abs_path = os.path.abspath(dir_path) - if os.path.exists(abs_path): + def _add_dir(abs_path, source): + if abs_path not in output_dirs: output_dirs.append(abs_path) - self.logger.info(f" [DIR] Found output dir: {abs_path}") - - # Count images in this directory + self.logger.info(f" [DIR] {source}: {abs_path}") try: image_files = [] for ext in [".png", ".jpg", ".jpeg", ".webp"]: @@ -293,17 +299,41 @@ class GalleryDiagnostics: except Exception as e: self.logger.error(f" [FAIL] Error scanning: {e}") - # Try ComfyUI's folder_paths + # Check user-configured gallery directories from config.json + try: + from py.config import GalleryConfig + + for cfg_dir in GalleryConfig.MONITORING_DIRECTORIES: + abs_path = os.path.abspath(cfg_dir) + if os.path.exists(abs_path): + _add_dir(abs_path, "Configured directory") + except Exception as e: + self.logger.debug(f" [NOTE] Could not load gallery config: {e}") + + # Check ComfyUI's own output directory (respects --output-directory) try: import folder_paths comfyui_output = folder_paths.get_output_directory() - if comfyui_output and comfyui_output not in output_dirs: - output_dirs.append(comfyui_output) - self.logger.info(f" [DIR] ComfyUI output dir: {comfyui_output}") + if comfyui_output and os.path.exists(comfyui_output): + _add_dir(os.path.abspath(comfyui_output), "ComfyUI output dir") except ImportError: self.logger.warning(" [WARN] ComfyUI folder_paths not available") + # Fallback: check common relative paths + if not output_dirs: + for dir_path in [ + "output", + "../output", + "../../output", + "ComfyUI/output", + "../ComfyUI/output", + "../../ComfyUI/output", + ]: + abs_path = os.path.abspath(dir_path) + if os.path.exists(abs_path): + _add_dir(abs_path, "Found output dir") + return { "status": "ok" if output_dirs else "warning", "message": "No output directories found" if not output_dirs else None,