From 4c6dbc92fe02e3d2a5bf779a986ab17cf82b27de Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Thu, 8 Oct 2026 18:39:46 -0700 Subject: [PATCH] fix: serve every gallery video format, validate settings before applying, dedupe output dirs Review findings on #153 (Codex): - the image serve route allowed only three video types while the gallery lists seven, so .avi/.mkv/.m4v/.wmv cards opened to a 403; the serve allow-list now covers every gallery video format - a settings payload with one invalid field still applied the valid ones before returning 400; the whole payload is validated first and nothing changes on failure. webui_display_mode is validated too Third Windows CI run: - the same output directory spelled twice (different case or a link) was listed twice in diagnostics and served roots; both dedupe on the resolved, case-folded path - the root-less subfolder test expects the "/" form Code-quality bot: index definitions are (name, target) pairs instead of implicitly concatenated strings; redundant local asyncio imports and two trivial lambdas in test stubs removed. --- database/models.py | 33 +++++++-------- py/api/__init__.py | 12 ++++-- py/api/admin.py | 86 ++++++++++++++++++++++----------------- py/api/images.py | 9 ++-- tests/test_api_admin.py | 55 ++++++++++++++++++++++--- tests/test_api_autotag.py | 2 +- tests/test_api_images.py | 12 ++++++ tests/test_api_lora.py | 2 +- tests/test_database.py | 5 ++- 9 files changed, 146 insertions(+), 70 deletions(-) diff --git a/database/models.py b/database/models.py index e6761c2..84e7392 100644 --- a/database/models.py +++ b/database/models.py @@ -250,26 +250,23 @@ class PromptModel: - Hash values for duplicate detection """ indexes = [ - "CREATE INDEX IF NOT EXISTS idx_prompts_text ON prompts(text)", - "CREATE INDEX IF NOT EXISTS idx_prompts_category ON prompts(category)", - "CREATE INDEX IF NOT EXISTS idx_prompts_created_at ON prompts(created_at)", - "CREATE INDEX IF NOT EXISTS idx_prompts_hash ON prompts(hash)", - "CREATE INDEX IF NOT EXISTS idx_prompts_rating ON prompts(rating)", - "CREATE INDEX IF NOT EXISTS idx_prompts_last_used ON prompts(last_used_at)", - "CREATE INDEX IF NOT EXISTS idx_prompts_run_count ON prompts(run_count)", - "CREATE INDEX IF NOT EXISTS idx_prompt_images ON " - "generated_images(prompt_id)", - "CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)", - "CREATE INDEX IF NOT EXISTS idx_image_file_path ON " - "generated_images(file_path)", - "CREATE INDEX IF NOT EXISTS idx_generation_time ON " - "generated_images(generation_time)", - "CREATE INDEX IF NOT EXISTS idx_prompt_tags_tag ON prompt_tags(tag_id)", - "CREATE INDEX IF NOT EXISTS idx_tags_name ON tags(name)", + ("idx_prompts_text", "prompts(text)"), + ("idx_prompts_category", "prompts(category)"), + ("idx_prompts_created_at", "prompts(created_at)"), + ("idx_prompts_hash", "prompts(hash)"), + ("idx_prompts_rating", "prompts(rating)"), + ("idx_prompts_last_used", "prompts(last_used_at)"), + ("idx_prompts_run_count", "prompts(run_count)"), + ("idx_prompt_images", "generated_images(prompt_id)"), + ("idx_image_path", "generated_images(image_path)"), + ("idx_image_file_path", "generated_images(file_path)"), + ("idx_generation_time", "generated_images(generation_time)"), + ("idx_prompt_tags_tag", "prompt_tags(tag_id)"), + ("idx_tags_name", "tags(name)"), ] - for index_sql in indexes: - conn.execute(index_sql) + for name, target in indexes: + conn.execute(f"CREATE INDEX IF NOT EXISTS {name} ON {target}") def get_connection(self) -> sqlite3.Connection: """ diff --git a/py/api/__init__.py b/py/api/__init__.py index 1a7d5b6..747fa30 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -639,12 +639,18 @@ class PromptManagerAPI( configured = list(GalleryConfig.MONITORING_DIRECTORIES) output_dirs = [] + seen = set() for d in configured: ok, reason = GalleryConfig.validate_gallery_root(d) - if ok: - output_dirs.append(Path(os.path.realpath(d))) - else: + if not ok: self.logger.warning(f"Ignoring configured gallery root: {reason}") + continue + real = os.path.realpath(d) + key = os.path.normcase(real) # one entry per directory, however spelled + if key in seen: + continue + seen.add(key) + output_dirs.append(Path(real)) if output_dirs: return output_dirs diff --git a/py/api/admin.py b/py/api/admin.py index 6a92dc4..d13c034 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -242,16 +242,23 @@ def _diagnose_output_dirs(configured_dirs): pass output_dirs = [] + seen = set() + + def add(path): + # One entry per directory: resolve links and fold case (Windows) + real = os.path.realpath(path) + key = os.path.normcase(real) + if os.path.exists(real) and key not in seen: + seen.add(key) + output_dirs.append(real) + for candidate in candidates: - abs_path = os.path.abspath(candidate) if candidate else None - if abs_path and os.path.exists(abs_path) and abs_path not in output_dirs: - output_dirs.append(abs_path) + if candidate: + add(candidate) if not output_dirs: for rel in ("output", "../output", "../../output"): - abs_path = os.path.abspath(rel) - if os.path.exists(abs_path) and abs_path not in output_dirs: - output_dirs.append(abs_path) + add(rel) return { "status": "ok" if output_dirs else "warning", @@ -742,7 +749,7 @@ class AdminRoutesMixin: return list(GalleryConfig.MONITORING_DIRECTORIES) async def save_settings(self, request): - """Save settings.""" + """Save settings: the whole payload is validated before anything changes.""" try: from ..config import PromptManagerConfig @@ -750,41 +757,19 @@ class AdminRoutesMixin: if error_response is not None: return error_response - if "result_timeout" in data: - try: - PromptManagerConfig.RESULT_TIMEOUT = validate_result_timeout( - data["result_timeout"] - ) - except ValueError as ve: - return web.json_response( - {"success": False, "error": str(ve)}, status=400 - ) - if "webui_display_mode" in data: - PromptManagerConfig.WEBUI_DISPLAY_MODE = data["webui_display_mode"] - if "infinite_scroll" in data: - if not isinstance(data["infinite_scroll"], bool): - return web.json_response( - { - "success": False, - "error": "infinite_scroll must be true or false", - }, - status=400, - ) - PromptManagerConfig.INFINITE_SCROLL = data["infinite_scroll"] - if "worker_threads" in data: - try: - PromptManagerConfig.WORKER_THREADS = validate_worker_threads( - data["worker_threads"], PromptManagerConfig.max_worker_threads() - ) - except ValueError as ve: - return web.json_response( - {"success": False, "error": str(ve)}, status=400 - ) + try: + updates = self._validate_settings(data) + except ValueError as ve: + return web.json_response( + {"success": False, "error": str(ve)}, status=400 + ) new_roots, error = self._parse_gallery_roots(data) if error: return web.json_response({"success": False, "error": error}, status=400) + for name, value in updates.items(): + setattr(PromptManagerConfig, name, value) restart_required = False if new_roots is not None: restart_required = self._apply_gallery_roots(new_roots) @@ -805,6 +790,33 @@ class AdminRoutesMixin: status=500, ) + @staticmethod + def _validate_settings(data): + """Scalar settings of a payload as ``{CONFIG_ATTR: value}``. + + Raises ValueError with a user-facing message on the first bad value; + nothing is applied until every value passed. + """ + from ..config import PromptManagerConfig + + updates = {} + if "result_timeout" in data: + updates["RESULT_TIMEOUT"] = validate_result_timeout(data["result_timeout"]) + if "webui_display_mode" in data: + mode = data["webui_display_mode"] + if mode not in ("popup", "newtab"): + raise ValueError("webui_display_mode must be 'popup' or 'newtab'") + updates["WEBUI_DISPLAY_MODE"] = mode + if "infinite_scroll" in data: + if not isinstance(data["infinite_scroll"], bool): + raise ValueError("infinite_scroll must be true or false") + updates["INFINITE_SCROLL"] = data["infinite_scroll"] + if "worker_threads" in data: + updates["WORKER_THREADS"] = validate_worker_threads( + data["worker_threads"], PromptManagerConfig.max_worker_threads() + ) + return updates + def _parse_gallery_roots(self, data): """Extract and validate gallery roots from a settings payload. diff --git a/py/api/images.py b/py/api/images.py index 27aa0c1..d03e27d 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -79,7 +79,7 @@ def _media_access_error(path, allowed_dirs): Fails closed: an empty *allowed_dirs* denies everything. """ - if Path(path).suffix.lower() not in IMAGE_EXTENSIONS: + if Path(path).suffix.lower() not in SERVABLE_EXTENSIONS: return _forbidden("Only media files can be served") if not allowed_dirs: return _forbidden("No allowed image directories configured") @@ -92,6 +92,9 @@ def _media_access_error(path, allowed_dirs): GALLERY_IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".webp", ".gif"}) VIDEO_EXTENSIONS = frozenset({".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"}) MEDIA_EXTENSIONS = GALLERY_IMAGE_EXTENSIONS | VIDEO_EXTENSIONS +# What the serve routes hand out: every image type above plus every video type +# the gallery lists, so a card never points at a 403. +SERVABLE_EXTENSIONS = IMAGE_EXTENSIONS | VIDEO_EXTENSIONS # Caps on request-triggered filesystem work. MAX_SCAN_DEPTH = 12 @@ -707,7 +710,7 @@ class ImageRoutesMixin: filepath = request.match_info["filepath"] if _has_traversal(filepath): return _forbidden() - if Path(filepath).suffix.lower() not in IMAGE_EXTENSIONS: + if Path(filepath).suffix.lower() not in SERVABLE_EXTENSIONS: return _forbidden("Only media files can be served") output_dirs = self._select_root( @@ -1347,7 +1350,7 @@ class ImageRoutesMixin: allowed_dirs = list(self._get_all_output_dirs()) resolved = self._resolve_client_media_path(image_path, allowed_dirs) if resolved is None: - if Path(image_path).suffix.lower() not in IMAGE_EXTENSIONS: + if Path(image_path).suffix.lower() not in SERVABLE_EXTENSIONS: return _forbidden("Only media files can be served") return web.json_response( {"success": False, "error": "Image file not found"}, status=404 diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index a6a5e04..be48a55 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -303,6 +303,23 @@ class TestNoAbsolutePathsInResponses(AdminAPITestCase): [os.path.normcase(os.path.realpath(renders))], ) + @unittest.skipIf(sys.platform == "win32", "symlinks need privileges on Windows") + async def test_same_output_dir_spelled_twice_is_listed_once(self): + alias = self.comfy_dir / "alias" + os.symlink(self.output_dir, alias, target_is_directory=True) + # Two spellings of one directory (Windows adds case differences too) + GalleryConfig.MONITORING_DIRECTORIES = [str(alias), self._root(self.output_dir)] + self._install_monitor([self._root(self.output_dir)]) + + resp, body = await self._body("GET", "/prompt_manager/diagnostics") + + self.assertEqual(resp.status, 200) + data = json.loads(body) + self.assertEqual( + data["diagnostics"]["comfyui_output"]["output_dirs"], ["output"] + ) + self.assertEqual(len(self.api._get_all_output_dirs()), 1) + async def test_diagnostics_paths_are_relative(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] self._install_monitor([self._root(self.output_dir)]) @@ -882,8 +899,6 @@ class TestDuplicateJobsAreSingleFlight(AdminAPITestCase): """A duplicate scan or delete refuses to overlap with a running one.""" def _block_scan(self): - import asyncio - gate = asyncio.Event() async def slow_scan(): @@ -894,8 +909,6 @@ class TestDuplicateJobsAreSingleFlight(AdminAPITestCase): return gate async def test_second_scan_while_one_runs_is_409(self): - import asyncio - gate = self._block_scan() first = asyncio.ensure_future( self.client.request("GET", "/prompt_manager/scan_duplicates") @@ -914,8 +927,6 @@ class TestDuplicateJobsAreSingleFlight(AdminAPITestCase): self.assertEqual((await first).status, 200) async def test_delete_while_scan_runs_is_409_and_scan_after_is_fine(self): - import asyncio - gate = self._block_scan() first = asyncio.ensure_future( self.client.request("GET", "/prompt_manager/scan_duplicates") @@ -948,6 +959,38 @@ class TestDuplicateJobsAreSingleFlight(AdminAPITestCase): class TestSettingsMisc(AdminAPITestCase): + async def test_invalid_payload_changes_nothing(self): + before = ( + PromptManagerConfig.RESULT_TIMEOUT, + PromptManagerConfig.WEBUI_DISPLAY_MODE, + PromptManagerConfig.WORKER_THREADS, + PromptManagerConfig.INFINITE_SCROLL, + list(GalleryConfig.MONITORING_DIRECTORIES), + ) + + resp = await self.client.request( + "POST", + "/prompt_manager/settings", + json={ + "result_timeout": 42, + "webui_display_mode": "popup", + "worker_threads": 1, + "infinite_scroll": True, + "gallery_root_paths": [os.path.abspath(os.sep)], # rejected + }, + ) + + self.assertEqual(resp.status, 400) + after = ( + PromptManagerConfig.RESULT_TIMEOUT, + PromptManagerConfig.WEBUI_DISPLAY_MODE, + PromptManagerConfig.WORKER_THREADS, + PromptManagerConfig.INFINITE_SCROLL, + list(GalleryConfig.MONITORING_DIRECTORIES), + ) + self.assertEqual(after, before) + self.assertFalse(os.path.exists(self.config_path)) + async def test_infinite_scroll_default_is_reported_saved_and_validated(self): data = await ( await self.client.request("GET", "/prompt_manager/settings") diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index be765c2..1259e65 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -27,7 +27,7 @@ if "folder_paths" not in sys.modules: _stub = types.ModuleType("folder_paths") _stub.models_dir = tempfile.gettempdir() _stub.base_path = tempfile.gettempdir() - _stub.get_output_directory = lambda: tempfile.gettempdir() + _stub.get_output_directory = tempfile.gettempdir sys.modules["folder_paths"] = _stub from aiohttp import web diff --git a/tests/test_api_images.py b/tests/test_api_images.py index 7c77e31..9da9829 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -187,6 +187,18 @@ class TestServeOutputImage(ImageAPITestCase): self.assertEqual(resp.status, 403) + async def test_every_gallery_video_format_is_served(self): + # The gallery lists these as playable cards; opening them must not be a 403. + for ext in sorted(images_module.VIDEO_EXTENSIONS): + (self.output_dir / f"clip{ext}").write_bytes(b"\x00\x00\x00\x18ftyp") + + for ext in sorted(images_module.VIDEO_EXTENSIONS): + resp = await self.client.request( + "GET", f"/prompt_manager/images/serve/clip{ext}" + ) + await resp.read() + self.assertEqual(resp.status, 200, ext) + async def test_missing_file_is_404(self): resp = await self.client.request("GET", "/prompt_manager/images/serve/nope.png") diff --git a/tests/test_api_lora.py b/tests/test_api_lora.py index fe7e226..1c340b1 100644 --- a/tests/test_api_lora.py +++ b/tests/test_api_lora.py @@ -20,7 +20,7 @@ if "folder_paths" not in sys.modules: _stub = types.ModuleType("folder_paths") _stub.models_dir = tempfile.gettempdir() _stub.base_path = tempfile.gettempdir() - _stub.get_output_directory = lambda: tempfile.gettempdir() + _stub.get_output_directory = tempfile.gettempdir _stub.get_folder_paths = lambda name: [] sys.modules["folder_paths"] = _stub diff --git a/tests/test_database.py b/tests/test_database.py index e5f58b7..a73fcee 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -1069,7 +1069,10 @@ class TestMaintenanceOperations(DatabaseTestCase): self.assertIn("2026/08-Aug", nested) without_roots = self.db.get_prompt_subfolders() - self.assertIn(os.path.join(root, "2026", "08-Aug", "2026-08-06"), without_roots) + self.assertIn( + os.path.join(root, "2026", "08-Aug", "2026-08-06").replace(os.sep, "/"), + without_roots, + ) def test_prompts_by_tags_and_or_modes(self): both = self._save("both", tags=["x", "y"])