diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 4e0a1cf..660793e 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -2,11 +2,23 @@ import asyncio import json +import os from pathlib import Path +from typing import Optional from aiohttp import web +def _public_error(exc: Exception) -> str: + """Describe an exception without leaking absolute server paths.""" + if isinstance(exc, OSError) and exc.strerror: + filename = exc.filename + if filename: + return f"{exc.strerror}: {os.path.basename(str(filename))}" + return exc.strerror + return str(exc) + + class AutotagRoutesMixin: """Mixin providing auto-tagging API endpoints.""" @@ -15,11 +27,11 @@ class AutotagRoutesMixin: async def get_autotag_models_route(request): return await self.get_autotag_models(request) - @routes.get("/prompt_manager/autotag/download/{model_type}") + @routes.post("/prompt_manager/autotag/download/{model_type}") async def download_autotag_model_route(request): return await self.download_autotag_model(request) - @routes.get("/prompt_manager/autotag/start") + @routes.post("/prompt_manager/autotag/start") async def start_autotag_route(request): return await self.start_autotag(request) @@ -116,17 +128,39 @@ class AutotagRoutesMixin: await response.write_eof() return response + async def _read_start_params(self, request) -> Optional[dict]: + """Merge query-string and optional JSON-body parameters for autotag/start. + + Returns None when a body is present but is not a JSON object. + """ + params = dict(request.query) + if request.can_read_body: + try: + body = await request.json() + except (json.JSONDecodeError, UnicodeDecodeError): + return None + if not isinstance(body, dict): + return None + params.update(body) + return params + async def start_autotag(self, request): - """Start batch auto-tagging with streaming progress.""" - model_type = request.query.get("model_type", "gguf") - custom_prompt = request.query.get("prompt", "") - skip_tagged = request.query.get("skip_tagged", "true").lower() == "true" - keep_in_memory = request.query.get("keep_in_memory", "true").lower() == "true" + """Start batch auto-tagging with streaming progress (POST only).""" + params = await self._read_start_params(request) + if params is None: + return web.json_response( + {"success": False, "error": "Request body must be a JSON object"}, + status=400, + ) + model_type = params.get("model_type", "gguf") + custom_prompt = params.get("prompt", "") + skip_tagged = str(params.get("skip_tagged", "true")).lower() == "true" + keep_in_memory = str(params.get("keep_in_memory", "true")).lower() == "true" use_gpu = True # WD14 threshold params - general_threshold = request.query.get("general_threshold") - character_threshold = request.query.get("character_threshold") + general_threshold = params.get("general_threshold") + character_threshold = params.get("character_threshold") try: if general_threshold is not None: general_threshold = float(general_threshold) @@ -326,11 +360,90 @@ class AutotagRoutesMixin: await response.write_eof() return response + def _resolve_allowed_image_path(self, raw_path: str) -> Optional[Path]: + """Return the resolved path when it lies inside an output directory.""" + try: + resolved = Path(raw_path).resolve() + except (OSError, RuntimeError): + return None + for output_dir in self._get_all_output_dirs(): + try: + if resolved.is_relative_to(Path(output_dir).resolve()): + return resolved + except (OSError, RuntimeError): + continue + return None + + async def _select_autotag_target(self, data: dict): + """Pick the image to tag from an autotag/single body. + + Returns ``(image_path, prompt_id, None)`` on success or + ``(None, None, error_response)`` when the request must be rejected. + """ + image_id = data.get("image_id") + if image_id is not None: + if isinstance(image_id, bool) or not isinstance(image_id, (int, str)): + return None, None, self._bad_request("image_id must be an integer") + try: + image_id = int(image_id) + except (TypeError, ValueError): + return None, None, self._bad_request("image_id must be an integer") + record = await self._run_in_executor(self.db.get_image_by_id, image_id) + if not record or not record.get("image_path"): + return ( + None, + None, + web.json_response( + {"success": False, "error": f"Image {image_id} not found"}, + status=404, + ), + ) + return record["image_path"], record.get("prompt_id"), None + + raw_path = data.get("path") or data.get("image_path") + if not raw_path or not isinstance(raw_path, str): + return None, None, self._bad_request("image_id or path is required") + + resolved = self._resolve_allowed_image_path(raw_path) + if resolved is None: + return ( + None, + None, + web.json_response( + { + "success": False, + "error": ( + f"{os.path.basename(raw_path)} is outside the " + "configured output directories" + ), + }, + status=403, + ), + ) + + prompt_id = None + try: + prompt_id = await self._run_in_executor( + self.db.get_prompt_id_for_image, raw_path + ) + if prompt_id is None and str(resolved) != raw_path: + prompt_id = await self._run_in_executor( + self.db.get_prompt_id_for_image, str(resolved) + ) + except Exception as e: + self.logger.warning(f"Could not find linked prompt: {e}") + return str(resolved), prompt_id, None + + @staticmethod + def _bad_request(message: str): + return web.json_response({"success": False, "error": message}, status=400) + async def autotag_single(self, request): - """Generate tags for a single image.""" + """Generate tags for a single image selected by image_id or gallery path.""" try: data = await request.json() - image_path = data.get("image_path") + if not isinstance(data, dict): + return self._bad_request("Request body must be a JSON object") model_type = data.get("model_type", "gguf") custom_prompt = data.get("prompt") use_gpu = data.get("use_gpu", True) @@ -342,18 +455,13 @@ class AutotagRoutesMixin: if character_threshold is not None: character_threshold = float(character_threshold) except (ValueError, TypeError): - return web.json_response( - { - "success": False, - "error": "general_threshold and character_threshold must be numeric", - }, - status=400, + return self._bad_request( + "general_threshold and character_threshold must be numeric" ) - if not image_path: - return web.json_response( - {"success": False, "error": "image_path is required"}, status=400 - ) + image_path, prompt_id, error = await self._select_autotag_target(data) + if error is not None: + return error from ..autotag import get_autotag_service @@ -363,44 +471,34 @@ class AutotagRoutesMixin: not service.is_model_loaded() or service.get_loaded_model_type() != model_type ): - loop = asyncio.get_event_loop() - await loop.run_in_executor( - None, lambda: service.load_model(model_type, use_gpu) - ) + await self._run_in_executor(service.load_model, model_type, use_gpu) if custom_prompt: service.custom_prompt = custom_prompt - loop = asyncio.get_event_loop() - tags = await loop.run_in_executor( - None, - lambda: service.generate_tags( - image_path, - general_threshold=general_threshold, - character_threshold=character_threshold, - ), + tags = await self._run_in_executor( + service.generate_tags, + image_path, + general_threshold=general_threshold, + character_threshold=character_threshold, ) - prompt_id = None - try: - prompt_id = await self._run_in_executor( - self.db.get_prompt_id_for_image, image_path - ) - except Exception as e: - self.logger.warning(f"Could not find linked prompt: {e}") - return web.json_response( { "success": True, "tags": tags, "prompt_id": prompt_id, - "image_path": image_path, + "filename": os.path.basename(image_path), } ) + except json.JSONDecodeError: + return self._bad_request("Request body must be valid JSON") except Exception as e: self.logger.error(f"AutoTag single error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _public_error(e)}, status=500 + ) async def apply_autotag(self, request): """Apply selected tags to a prompt.""" diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py new file mode 100644 index 0000000..3c151fb --- /dev/null +++ b/tests/test_api_autotag.py @@ -0,0 +1,296 @@ +"""Route tests for the autotag endpoints. + +The ML engine is replaced by a fake service so no model is ever loaded and +no network request is made. +""" + +import json +import os +import shutil +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +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() + sys.modules["folder_paths"] = _stub + +from aiohttp import web +from aiohttp.test_utils import AioHTTPTestCase + +from database.operations import PromptDatabase +from py.api import PromptManagerAPI +from utils.hashing import generate_prompt_hash + + +class FakeAutotagService: + """Stand-in for AutoTagService that records calls and returns fixed tags.""" + + def __init__(self, tags=None, downloaded=True): + self.tags = tags if tags is not None else ["1girl", "smile"] + self.downloaded = downloaded + self.generate_calls = [] + self.load_calls = [] + self.download_calls = [] + self.unload_calls = 0 + self.custom_prompt = "" + self.default_prompt = "default prompt" + self.wd14_general_threshold = 0.35 + self.wd14_character_threshold = 0.85 + self.models_config = {"gguf": {}, "wd14-vit": {}} + self._loaded_type = None + self.generate_error = None + self.download_result = True + self.load_error = None + + def get_models_status(self): + return { + name: {"downloaded": self.downloaded, "model_path": None} + for name in self.models_config + } + + def is_model_loaded(self): + return self._loaded_type is not None + + def get_loaded_model_type(self): + return self._loaded_type + + def load_model(self, model_type, use_gpu=True): + if self.load_error is not None: + raise self.load_error + self.load_calls.append((model_type, use_gpu)) + self._loaded_type = model_type + return True + + def unload_model(self): + self.unload_calls += 1 + self._loaded_type = None + + def download_model(self, model_type, progress_callback=None): + self.download_calls.append(model_type) + if progress_callback: + progress_callback("Downloading...", 50) + return self.download_result + + def generate_tags( + self, image_path, prompt=None, general_threshold=None, character_threshold=None + ): + self.generate_calls.append( + (image_path, prompt, general_threshold, character_threshold) + ) + if self.generate_error is not None: + raise self.generate_error + return list(self.tags) + + +class AutotagAPITestCase(AioHTTPTestCase): + """Stands up the PromptManager routes with a temp DB, temp output dir and a fake engine.""" + + async def get_application(self): + self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") + self._temp_db.close() + self.tmp_root = Path(tempfile.mkdtemp()) + self.output_dir = self.tmp_root / "output" + self.output_dir.mkdir() + self.secret_dir = self.tmp_root / "secret" + self.secret_dir.mkdir() + + self.service = FakeAutotagService() + self._service_patch = patch( + "py.autotag.get_autotag_service", return_value=self.service + ) + self._service_patch.start() + + app = web.Application() + routes = web.RouteTableDef() + self.api = PromptManagerAPI() + self.api.db = PromptDatabase(self._temp_db.name) + self.api._get_all_output_dirs = lambda: [self.output_dir] + self.api.add_routes(routes) + app.router.add_routes(routes) + return app + + async def tearDownAsync(self): + self._service_patch.stop() + for path in ( + self._temp_db.name, + self._temp_db.name + "-wal", + self._temp_db.name + "-shm", + ): + if os.path.exists(path): + os.unlink(path) + shutil.rmtree(self.tmp_root, ignore_errors=True) + + # ── helpers ──────────────────────────────────────────────────────── + + def _save_prompt(self, text="Test prompt", tags=None): + return self.api.db.save_prompt( + text=text, + category=None, + tags=tags or [], + rating=None, + prompt_hash=generate_prompt_hash(text), + ) + + def _make_image(self, directory, name="img.png"): + path = directory / name + path.write_bytes(b"\x89PNG\r\n\x1a\n") + return path + + def _link_image(self, prompt_id, path): + return self.api.db.link_image_to_prompt(prompt_id, str(path)) + + async def _post_json(self, url, body): + return await self.client.request( + "POST", + url, + data=json.dumps(body), + headers={"Content-Type": "application/json"}, + ) + + @staticmethod + def _sse_events(text): + events = [] + for line in text.splitlines(): + if line.startswith("data: "): + events.append(json.loads(line[len("data: ") :])) + return events + + +class TestAutotagSingleContainment(AutotagAPITestCase): + + async def test_image_id_tags_the_linked_image(self): + prompt_id = self._save_prompt() + image = self._make_image(self.output_dir) + image_id = self._link_image(prompt_id, image) + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"image_id": image_id} + ) + + self.assertEqual(resp.status, 200) + data = await resp.json() + self.assertTrue(data["success"]) + self.assertEqual(data["tags"], ["1girl", "smile"]) + self.assertEqual(data["prompt_id"], prompt_id) + self.assertEqual(len(self.service.generate_calls), 1) + self.assertEqual( + os.path.realpath(self.service.generate_calls[0][0]), + os.path.realpath(str(image)), + ) + + async def test_path_inside_output_dir_is_tagged(self): + image = self._make_image(self.output_dir) + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"path": str(image)} + ) + + self.assertEqual(resp.status, 200) + data = await resp.json() + self.assertTrue(data["success"]) + self.assertEqual(data["tags"], ["1girl", "smile"]) + self.assertEqual(len(self.service.generate_calls), 1) + + async def test_path_outside_output_dir_is_forbidden(self): + secret = self._make_image(self.secret_dir, "password.png") + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"path": str(secret)} + ) + + self.assertEqual(resp.status, 403) + data = await resp.json() + self.assertFalse(data["success"]) + self.assertNotIn(str(self.secret_dir), data["error"]) + self.assertEqual(self.service.generate_calls, []) + + async def test_traversal_path_is_forbidden(self): + self._make_image(self.secret_dir, "password.png") + traversal = os.path.join(str(self.output_dir), "..", "secret", "password.png") + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"path": traversal} + ) + + self.assertEqual(resp.status, 403) + self.assertEqual(self.service.generate_calls, []) + + async def test_legacy_image_path_field_is_still_contained(self): + secret = self._make_image(self.secret_dir, "password.png") + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"image_path": str(secret)} + ) + + self.assertEqual(resp.status, 403) + self.assertEqual(self.service.generate_calls, []) + + async def test_missing_selector_is_400(self): + resp = await self._post_json("/prompt_manager/autotag/single", {}) + self.assertEqual(resp.status, 400) + self.assertFalse((await resp.json())["success"]) + + async def test_unknown_image_id_is_404(self): + resp = await self._post_json( + "/prompt_manager/autotag/single", {"image_id": 9999} + ) + self.assertEqual(resp.status, 404) + self.assertFalse((await resp.json())["success"]) + + async def test_non_integer_image_id_is_400(self): + resp = await self._post_json( + "/prompt_manager/autotag/single", {"image_id": "abc"} + ) + self.assertEqual(resp.status, 400) + + +class TestAutotagSideEffectRoutesArePostOnly(AutotagAPITestCase): + + async def test_get_download_is_405(self): + resp = await self.client.request("GET", "/prompt_manager/autotag/download/gguf") + self.assertEqual(resp.status, 405) + self.assertEqual(self.service.download_calls, []) + + async def test_get_start_is_405(self): + resp = await self.client.request("GET", "/prompt_manager/autotag/start") + self.assertEqual(resp.status, 405) + self.assertEqual(self.service.load_calls, []) + + async def test_post_download_streams_completion(self): + resp = await self.client.request( + "POST", "/prompt_manager/autotag/download/gguf" + ) + + self.assertEqual(resp.status, 200) + events = self._sse_events(await resp.text()) + self.assertEqual(events[-1]["type"], "complete") + self.assertEqual(self.service.download_calls, ["gguf"]) + + async def test_post_start_with_no_images_completes(self): + resp = await self.client.request( + "POST", "/prompt_manager/autotag/start?model_type=gguf" + ) + + self.assertEqual(resp.status, 200) + events = self._sse_events(await resp.text()) + self.assertEqual(events[-1]["type"], "complete") + self.assertEqual(events[-1]["processed"], 0) + self.assertEqual(self.service.load_calls, [("gguf", True)]) + + async def test_models_status_stays_get(self): + resp = await self.client.request("GET", "/prompt_manager/autotag/models") + self.assertEqual(resp.status, 200) + self.assertTrue((await resp.json())["success"]) + + +if __name__ == "__main__": + unittest.main()