feat(images): set, remove and list images' own categories (#332)

PUT /prompt_manager/images/category sets the own category of up to 500
listed files, DELETE removes it (listed or not), and
GET /prompt_manager/images/categories lists the values whose file is
listed. Targets are the Images listing identity of an output source,
checked by media_target (shared with Collections). The routes never
touch prompts, recipes or links; POST /prompt_manager/bulk/category is
unchanged.

TESTING.md

@@ -4470,6 +4470,64 @@ concept fixture stars the concept's favorites (each other prompt's first image).
   restore. Old backups without the table restore, and backup plus restore keeps
   WAL-resident favorites and closes every handle (`tests/test_database_backup.py`).
 
+Image categories (#332, store and routes): an image file may carry its own
+category; without one it shows the category its link supplies.
+
+- The table is `image_categories(root_key, relative_path, category, updated_at)`
+  with primary key `(root_key, relative_path)` and the index
+  `idx_image_categories_category`. `migrate_image_categories` creates both empty
+  inside a savepoint, with no backfill. It reads only `sqlite_master` and its own
+  table's PRAGMAs, and writes only its two objects. `tests/test_image_categories.py`
+  checks this with an SQLite authorizer on a 23,000-link library, and checks that
+  `prompts` and `prompt_pairs` `(id, category, updated_at)` are byte-for-byte equal
+  afterwards.
+- The **file key** is `image_category_key(sources, absolute path)` in
+  `py/api/images.py`: the deepest output source that contains the path (compared
+  lexically, never resolving the path, so a symlinked subfolder keeps its listed
+  spelling) and the path below it, keyed with the favorites helpers. Without
+  nested sources it equals a favorite's key. Every write and read uses it. A path
+  that no source contains has no key, so it has no own category.
+- `image_categories_schema_error` checks the structure at startup (a value saved
+  under older rules never stops the library opening). Restore also checks the
+  values: text columns, keys within the favorites bound, `category` 1 to 100
+  characters with no character below U+0020 and no leading or trailing space, and
+  `updated_at` of at most 64 bytes. Backups made before #332 restore without the
+  table; a migrated restore stage must have it (`require_library_state`).
+- `PUT /prompt_manager/images/category` `{items, category}` answers
+  `{success, category, updated, unchanged, not_found}`.
+  `DELETE /prompt_manager/images/category` `{items}` answers `{success, cleared}`.
+  `GET /prompt_manager/images/categories` answers `{success, categories}`: the
+  sorted distinct own values whose file is in some source's listing.
+- `items` are 1 to 500 Images targets `{root_index, root_scope, path}`, checked by
+  `py/api/media_target.py` (shared with Collections, #331). Refusals:
+  - 400: a malformed body, a field other than `items` and `category`, a
+    `root_index` outside the output sources (a LoRA folder), or a category that is
+    blank after stripping, over 100 characters, or holds a control character;
+  - 403: traversal, a drive, a backslash, a non-media name ("Only media files can
+    have a category") or a key over the stored bound;
+  - 409 `source-changed`: a stale `root_scope`;
+  - 413: a body over 1 MiB.
+  Any invalid item fails the whole request before a write. Targets that fold to one
+  key are written and counted once.
+- PUT writes only files the source's listing holds and counts the rest in
+  `not_found`; when none is listed it answers 404 `files-not-found` and writes
+  nothing. A file that already has the value counts as `unchanged` and keeps its
+  `updated_at`. DELETE accepts any well-formed target, so a missing file's
+  category can still be removed. An unexpected failure is a logged 500 ("The
+  category could not be changed." / "Image categories could not be read.") without
+  the exception text. The routes never read or write `prompts`, `prompt_pairs` or
+  `generated_images`; `POST /prompt_manager/bulk/category` is unchanged.
+- Duplicate cleanup (System › Maintenance): the kept copy's own category wins. A
+  kept copy without one receives the removed copy's (`COPY_SQL`, part of the
+  cleanup's query-plan check). The removed copy's row is never deleted, so a
+  give-back (the kept copy changed) and a failed delete lose nothing. The job log
+  says "Kept N image categories on the kept copies".
+  `tests/test_duplicate_image_categories.py` covers nested sources in both orders.
+- A moved, renamed or deleted file's row stays and is hidden; a file that later
+  appears at the same path shows it. The real-host check is the last test in
+  `tests/comfyui/image-attribution.spec.mjs`: PUT, list, DELETE on a generated
+  output.
+
 Image annotations schema (#249, `database/annotation_schema.py`): the group
 `media_annotation_metadata` (one row: schema version, `image_index_revision`),
 `media_annotation_files` (one private locator per physical file,

py/api/__init__.py

@@ -72,6 +72,7 @@ from .diagnostics_report import DiagnosticsReportMixin
 from .autotag_routes import AutotagRoutesMixin
 from .lora_integration import LoraIntegrationMixin
 from .lora_import_routes import LoraImportRoutesMixin
+from .image_categories import ImageCategoryRoutesMixin
 
 try:
     from ...database.operations import PromptDatabase
@@ -336,6 +337,7 @@ class PromptManagerAPI(
     AutotagRoutesMixin,
     LoraIntegrationMixin,
     LoraImportRoutesMixin,
+    ImageCategoryRoutesMixin,
 ):
     """REST API handler for PromptManager operations and web interface.
 
@@ -700,6 +702,7 @@ class PromptManagerAPI(
         self._register_autotag_routes(routes)
         self._register_lora_routes(routes)
         self._register_lora_import_routes(routes)
+        self._register_image_category_routes(routes)
 
         # Register gzip compression middleware (once)
         global _gzip_registered

py/api/annotations.py

@@ -27,7 +27,12 @@ from .annotation_suggest import (
     validate_options,
 )
 from .autotag_routes import model_busy_response
-from .images import _forbidden, _has_traversal, _unavailable_link, _verified_title
+
+# The target check is shared with the collection routes (#331); both names
+# stay importable from here.
+from .media_target import MAX_TARGET_PATH  # noqa: F401
+from .media_target import media_target as _target
+from .images import _unavailable_link, _verified_title
 from .prompts import bad_request
 from .. import model_gate
 
@@ -40,7 +45,7 @@ try:
     )
     from ...database.execution_history import subscribe_history_restore
     from ...database.library_commands import LibraryCommandError
-    from ...utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore
+    from ...utils.media_capture import CaptureError, CaptureStore
 except ImportError:
     from database.annotations import (
         MAX_IMAGE_TAGS,
@@ -50,15 +55,13 @@ except ImportError:
     )
     from database.execution_history import subscribe_history_restore
     from database.library_commands import LibraryCommandError
-    from utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore
+    from utils.media_capture import CaptureError, CaptureStore
 
 # Request bodies of the annotation routes are small (a target plus a tag set).
 MAX_BODY_BYTES = 64 * 1024
-MAX_TARGET_PATH = 1024
 CAPTURE_URL = "/prompt_manager/image-annotations/capture/"
 _PREVIEW_CHUNK = 64 * 1024
 
-_TARGET_FIELDS = frozenset({"root_index", "root_scope", "path"})
 _OBSERVE_FIELDS = frozenset({"target"})
 _SAVE_FIELDS = frozenset(
     {
@@ -188,37 +191,6 @@ def _capture_error(exc: CaptureError) -> web.Response:
     return _coded(exc.code, _MESSAGES.get(exc.code) or str(exc))
 
 
-def _target(data: object):
-    """``((index, scope, relative), None)`` for a target, else ``(None, error)``.
-
-    Only the listing identity is accepted: an int root index, the root's
-    64-hex scope and a bounded root-relative POSIX media path, never an
-    absolute, drive-qualified, backslash or ``..`` path. Nothing here touches
-    the filesystem or the database.
-    """
-    if not isinstance(data, dict) or set(data) != _TARGET_FIELDS:
-        return None, bad_request("target must name root_index, root_scope and path")
-    index, scope, relative = data["root_index"], data["root_scope"], data["path"]
-    if type(index) is not int or index < 0:
-        return None, bad_request("Invalid gallery root")
-    if not isinstance(scope, str) or _HEX64.fullmatch(scope) is None:
-        return None, bad_request("Invalid gallery root scope")
-    if not isinstance(relative, str) or not relative or len(relative) > MAX_TARGET_PATH:
-        return None, bad_request("path must be a relative media path")
-    parts = relative.split("/")
-    if (
-        "\0" in relative
-        or "\\" in relative
-        or relative.startswith("/")
-        or _has_traversal(relative)
-        or any(part in ("", ".", "..") or ":" in part for part in parts)
-    ):
-        return None, _forbidden()
-    if os.path.splitext(parts[-1])[1].lower() not in CAPTURE_EXTENSIONS:
-        return None, _forbidden("Only media files can be tagged")
-    return (index, scope, relative), None
-
-
 def _expires_at(seconds: float | None) -> str:
     moment = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
         seconds=max(0.0, seconds or 0.0)

py/api/image_categories.py

@@ -0,0 +1,285 @@
+"""Image category routes (#332): an image file's own category.
+
+``PUT /prompt_manager/images/category`` sets the own category of up to 500
+listed files; ``DELETE`` removes it (listed or not, so a missing file's
+category can still be removed); ``GET /prompt_manager/images/categories``
+lists the own-category values of files the sources list. Targets are the
+Images listing identity (``{root_index, root_scope, path}``, checked by
+``media_target``) of an output source; LoRA folders are not sources here.
+Each file is stored under its file key (``images.image_category_key``).
+These routes never read or write prompts, recipes or image links.
+"""
+
+import os
+from pathlib import Path
+
+from aiohttp import web
+
+from .image_tag_jobs import MAX_JOB_TARGETS, _parse_targets
+from .images import (
+    _favorite_relative_key,
+    _favorite_root_key,
+    _forbidden,
+    _json_body_or_error,
+    image_category_key,
+)
+from .prompts import bad_request
+
+try:
+    from ...database.image_categories import MAX_CATEGORY_LENGTH, MAX_KEY_LENGTH
+    from ...utils.validators import validate_category
+except ImportError:
+    from database.image_categories import MAX_CATEGORY_LENGTH, MAX_KEY_LENGTH
+    from utils.validators import validate_category
+
+CATEGORY_MEDIA_MESSAGE = "Only media files can have a category"
+_PUT_FIELDS = frozenset({"items", "category"})
+_DELETE_FIELDS = frozenset({"items"})
+SOURCE_CHANGED_MESSAGE = "The sources changed. Refresh the images and try again."
+WRITE_FAILED = "The category could not be changed."
+READ_FAILED = "Image categories could not be read."
+
+
+def _coded(status: int, code: str, message: str) -> web.Response:
+    return web.json_response(
+        {"success": False, "code": code, "error": message},
+        status=status,
+        headers={"Cache-Control": "no-store"},
+    )
+
+
+def _category_value(raw) -> tuple[str | None, web.Response | None]:
+    """The stripped category, or a 400 (an empty value is a 400: use DELETE)."""
+    if not isinstance(raw, str):
+        return None, bad_request("category must be a string")
+    value = raw.strip()
+    if not value:
+        return None, bad_request("category must not be empty")
+    if len(value) > MAX_CATEGORY_LENGTH:
+        return None, bad_request(
+            f"Category cannot exceed {MAX_CATEGORY_LENGTH} characters"
+        )
+    try:
+        validate_category(value)
+    except ValueError as exc:
+        return None, bad_request(str(exc))
+    return value, None
+
+
+def _nested_sources(root: Path, sources) -> list[tuple[str, str]]:
+    """``[(relative key prefix, source root key)]`` of sources inside *root*.
+
+    Deepest first, so the first prefix that matches is the file's source.
+    """
+    base = os.path.join(os.path.normcase(os.path.realpath(str(root))), "")
+    nested = []
+    for source in sources:
+        real = os.path.realpath(str(source))
+        if not os.path.normcase(real).startswith(base):
+            continue
+        relative = os.path.relpath(real, os.path.realpath(str(root)))
+        nested.append(
+            (
+                _favorite_relative_key(relative.replace(os.sep, "/")),
+                _favorite_root_key(source),
+            )
+        )
+    nested.sort(key=lambda item: item[0].count("/"), reverse=True)
+    return nested
+
+
+def _source_prefix(relative_key: str, nested) -> str | None:
+    """The prefix of the deepest nested source holding *relative_key*."""
+    for prefix, _source_key in nested:
+        if relative_key.startswith(prefix + "/"):
+            return prefix
+    return None
+
+
+class ImageCategoryRoutesMixin:
+    """Routes for images' own categories (#332)."""
+
+    def _register_image_category_routes(self, routes):
+        # Fixed paths beside DELETE /prompt_manager/images/{image_id}, as the
+        # favorite routes are: matched before the dynamic resource.
+        @routes.put("/prompt_manager/images/category")
+        async def set_image_category_route(request):
+            return await self.set_image_category(request)
+
+        @routes.delete("/prompt_manager/images/category")
+        async def clear_image_category_route(request):
+            return await self.clear_image_category(request)
+
+        @routes.get("/prompt_manager/images/categories")
+        async def list_image_categories_route(request):
+            return await self.list_image_categories(request)
+
+    # -- reading ---------------------------------------------------------------
+
+    def _own_category_map(self, sources, listings) -> dict:
+        """``{root / listed path: own category}`` for the listed files (blocking).
+
+        One primary-key range read per source; each stored row is matched to
+        the listing of every selected root that holds its source, so the work
+        follows the stored rows, not the files. A file under a nested source
+        is keyed by the deepest one (``image_category_key``). A row whose file
+        is gone matches nothing. A failed read raises.
+        """
+        sources = list(sources)
+        rows = self.db.own_categories_of(_favorite_root_key(s) for s in sources)
+        if not rows:
+            return {}
+        by_source: dict[str, dict[str, str]] = {}
+        for (root_key, relative), category in rows.items():
+            by_source.setdefault(root_key, {})[relative] = category
+        found = {}
+        for _index, root, files in listings:
+            listing = self._favorite_listing(root, files)
+            root_key = _favorite_root_key(root)
+            nested = [
+                item for item in _nested_sources(root, sources) if item[1] != root_key
+            ]
+            for relative, category in by_source.get(root_key, {}).items():
+                listed = listing.get(relative)
+                if listed is not None and _source_prefix(relative, nested) is None:
+                    found[root / listed] = category
+            for prefix, source_key in nested:
+                for relative, category in by_source.get(source_key, {}).items():
+                    full = f"{prefix}/{relative}"
+                    listed = listing.get(full)
+                    if listed is not None and _source_prefix(full, nested) == prefix:
+                        found[root / listed] = category
+        return found
+
+    async def _source_listings(self, roots) -> list[tuple]:
+        return [
+            (index, root, await self._get_gallery_files(root))
+            for index, root in enumerate(roots)
+        ]
+
+    async def list_image_categories(self, request):
+        """GET: the own-category values of listed files, sorted."""
+        try:
+            roots = list(self._get_all_output_dirs())
+            listings = await self._source_listings(roots)
+            found = await self._run_in_executor(self._own_category_map, roots, listings)
+        except Exception:
+            self.logger.exception("Image categories read failed")
+            return _coded(500, "categories-failed", READ_FAILED)
+        return web.json_response(
+            {"success": True, "categories": sorted(set(found.values()))},
+            headers={"Cache-Control": "no-store"},
+        )
+
+    # -- writing ---------------------------------------------------------------
+
+    async def _category_request(self, request, fields: frozenset):
+        """``(data, targets, roots, None)`` or ``(None, None, None, error)``."""
+        data, error = await _json_body_or_error(request)
+        if error is not None:
+            return None, None, None, error
+        if set(data) != fields:
+            return (
+                None,
+                None,
+                None,
+                bad_request(
+                    "Request must contain exactly " + ", ".join(sorted(fields))
+                ),
+            )
+        items = data["items"]
+        if not isinstance(items, list) or not 1 <= len(items) <= MAX_JOB_TARGETS:
+            return (
+                None,
+                None,
+                None,
+                bad_request(f"items must list 1 to {MAX_JOB_TARGETS} images"),
+            )
+        targets, error = _parse_targets(items, media_message=CATEGORY_MEDIA_MESSAGE)
+        if error is not None:
+            return None, None, None, error
+        roots = list(self._get_all_output_dirs())
+        for index, scope in dict.fromkeys((t[0], t[1]) for t in targets):
+            if index >= len(roots):
+                return None, None, None, bad_request("Invalid gallery root")
+            error = self._gallery_scope_error(
+                {"root_scope": scope, "root": str(index)}, roots
+            )
+            if error is not None:
+                if error.status == 409:
+                    error = _coded(409, "source-changed", SOURCE_CHANGED_MESSAGE)
+                return None, None, None, error
+        return data, targets, roots, None
+
+    async def set_image_category(self, request):
+        """PUT ``{items, category}``: set the own category of listed files."""
+        data, targets, roots, error = await self._category_request(request, _PUT_FIELDS)
+        if error is not None:
+            return error
+        category, error = _category_value(data["category"])
+        if error is not None:
+            return error
+        try:
+            listings = {}
+            listed_keys, missing_keys = {}, {}
+            for index, _scope, relative in targets:
+                root = roots[index]
+                if index not in listings:
+                    files = await self._get_gallery_files(root)
+                    listings[index] = await self._run_in_executor(
+                        self._favorite_listing, root, files
+                    )
+                listed = listings[index].get(_favorite_relative_key(relative))
+                key = image_category_key(roots, root / (listed or relative))
+                if key is None or max(len(key[0]), len(key[1])) > MAX_KEY_LENGTH:
+                    return _forbidden("Path is too long")
+                (missing_keys if listed is None else listed_keys)[key] = True
+            for key in listed_keys:
+                missing_keys.pop(key, None)
+            if not listed_keys:
+                return _coded(
+                    404,
+                    "files-not-found",
+                    "None of the selected images could be found.",
+                )
+            counts = await self._run_in_executor(
+                self.db.set_image_categories, list(listed_keys), category
+            )
+        except Exception:
+            self.logger.exception("Image category update failed")
+            return _coded(500, "category-failed", WRITE_FAILED)
+        return web.json_response(
+            {
+                "success": True,
+                "category": category,
+                "updated": counts["updated"],
+                "unchanged": counts["unchanged"],
+                "not_found": len(missing_keys),
+            },
+            headers={"Cache-Control": "no-store"},
+        )
+
+    async def clear_image_category(self, request):
+        """DELETE ``{items}``: remove the own category, listed or not."""
+        _data, targets, roots, error = await self._category_request(
+            request, _DELETE_FIELDS
+        )
+        if error is not None:
+            return error
+        keys = {}
+        for index, _scope, relative in targets:
+            key = image_category_key(roots, roots[index] / relative)
+            if key is None or max(len(key[0]), len(key[1])) > MAX_KEY_LENGTH:
+                return _forbidden("Path is too long")
+            keys[key] = True
+        try:
+            cleared = await self._run_in_executor(
+                self.db.clear_image_categories, list(keys)
+            )
+        except Exception:
+            self.logger.exception("Image category removal failed")
+            return _coded(500, "category-failed", WRITE_FAILED)
+        return web.json_response(
+            {"success": True, "cleared": cleared},
+            headers={"Cache-Control": "no-store"},
+        )

py/api/image_tag_jobs.py

@@ -193,13 +193,17 @@ class _Spec:
         return scope
 
 
-def _parse_targets(raw):
-    """Unique validated targets in request order, or ``(None, error)``."""
+def _parse_targets(raw, **target_options):
+    """Unique validated targets in request order, or ``(None, error)``.
+
+    ``target_options`` reach the target check (the image category routes
+    pass their own ``media_message``, #332).
+    """
     if not isinstance(raw, list) or not 1 <= len(raw) <= MAX_JOB_TARGETS:
         return None, bad_request(f"targets must list 1 to {MAX_JOB_TARGETS} images")
     seen, targets = set(), []
     for item in raw:
-        target, error = _target(item)
+        target, error = _target(item, **target_options)
         if error is not None:
             return None, error
         if target not in seen:

py/api/media_target.py

@@ -0,0 +1,54 @@
+"""The listing identity of one media file, as a request names it.
+
+Shared by the image-annotation routes and the collection routes (#331): a
+target is ``{root_index, root_scope, path}``, checked here without touching
+the filesystem or the database. The scope is compared with the current roots
+by the caller.
+"""
+
+import os
+import re
+
+from .images import _forbidden, _has_traversal
+from .prompts import bad_request
+
+try:
+    from ...utils.media_capture import CAPTURE_EXTENSIONS
+except ImportError:
+    from utils.media_capture import CAPTURE_EXTENSIONS
+
+MAX_TARGET_PATH = 1024
+TARGET_FIELDS = frozenset({"root_index", "root_scope", "path"})
+_HEX64 = re.compile(r"[0-9a-f]{64}")
+TAGGING_MEDIA_MESSAGE = "Only media files can be tagged"
+
+
+def media_target(data: object, *, media_message: str = TAGGING_MEDIA_MESSAGE):
+    """``((index, scope, relative), None)`` for a target, else ``(None, error)``.
+
+    Only the listing identity is accepted: an int root index, the root's
+    64-hex scope and a bounded root-relative POSIX media path, never an
+    absolute, drive-qualified, backslash or ``..`` path (a 400 for a malformed
+    shape, a 403 for an unsafe or non-media path, carrying *media_message*).
+    """
+    if not isinstance(data, dict) or set(data) != TARGET_FIELDS:
+        return None, bad_request("target must name root_index, root_scope and path")
+    index, scope, relative = data["root_index"], data["root_scope"], data["path"]
+    if type(index) is not int or index < 0:
+        return None, bad_request("Invalid gallery root")
+    if not isinstance(scope, str) or _HEX64.fullmatch(scope) is None:
+        return None, bad_request("Invalid gallery root scope")
+    if not isinstance(relative, str) or not relative or len(relative) > MAX_TARGET_PATH:
+        return None, bad_request("path must be a relative media path")
+    parts = relative.split("/")
+    if (
+        "\0" in relative
+        or "\\" in relative
+        or relative.startswith("/")
+        or _has_traversal(relative)
+        or any(part in ("", ".", "..") or ":" in part for part in parts)
+    ):
+        return None, _forbidden()
+    if os.path.splitext(parts[-1])[1].lower() not in CAPTURE_EXTENSIONS:
+        return None, _forbidden(media_message)
+    return (index, scope, relative), None

tests/comfyui/image-attribution.spec.mjs

@@ -140,3 +140,26 @@ test('an output saved without metadata links once its recorded run names a store
         expect((await request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true);
     }
 });
+
+// #332 on the real host: the migration ran at startup and the image category routes
+// set, list and remove an own category on a generated output.
+test('an output gets its own image category and loses it again', async ({ execution }, testInfo) => {
+    const request = execution.request;
+    const filename = await saved(request, graph({ text: `category host ${randomUUID()}`, prefix: 'pm332/category' }));
+    let image;
+    await expect.poll(async () => {
+        const listing = await (await request.get('/prompt_manager/images/output?mode=workspace&q=pm332&rescan=1')).json();
+        image = (listing.images || []).find(entry => entry.filename === filename || entry.path?.endsWith(`/${filename}`));
+        return Boolean(image);
+    }, { timeout: 20000, message: `${filename} is listed by Images` }).toBe(true);
+    const items = [{ root_index: image.root_index, root_scope: image.root_scope, path: image.path }];
+    const value = `PM332 host ${testInfo.project.name}`;
+    const set = await request.put('/prompt_manager/images/category', { data: { items, category: value } });
+    expect(set.status()).toBe(200);
+    expect(await set.json()).toEqual({ success: true, category: value, updated: 1, unchanged: 0, not_found: 0 });
+    expect((await (await request.get('/prompt_manager/images/categories')).json()).categories).toContain(value);
+    const cleared = await request.delete('/prompt_manager/images/category', { data: { items } });
+    expect(cleared.status()).toBe(200);
+    expect(await cleared.json()).toEqual({ success: true, cleared: 1 });
+    expect((await (await request.get('/prompt_manager/images/categories')).json()).categories).not.toContain(value);
+});

tests/test_api_image_categories.py

@@ -0,0 +1,303 @@
+"""Image category routes (#332): PUT/DELETE /images/category, GET /images/categories."""
+
+import json
+import ntpath
+import sqlite3
+from unittest.mock import patch
+
+try:
+    from test_api_images import ImageAPITestCase, make_png
+    from open_handles import assert_closed
+    from test_image_categories import category_checksum
+except ImportError:
+    from tests.test_api_images import ImageAPITestCase, make_png
+    from tests.open_handles import assert_closed
+    from tests.test_image_categories import category_checksum
+
+from database.image_categories import MAX_KEY_LENGTH, TABLE
+from py.api.images import _favorite_root_key, image_category_key
+
+CATEGORY = "/prompt_manager/images/category"
+CATEGORIES = "/prompt_manager/images/categories"
+
+
+class ImageCategoryTestCase(ImageAPITestCase):
+    async def tearDownAsync(self):
+        self.api.db.close_all()
+        assert_closed(self, self._temp_db.name, self._output_tmp.name)
+        await super().tearDownAsync()
+
+    def target(self, path, index=0):
+        roots = self.api._get_all_output_dirs()
+        return {
+            "root_index": index,
+            "root_scope": self.api._gallery_root_scope(roots[index], index),
+            "path": path,
+        }
+
+    async def send(self, method, body, status=200):
+        response = await self.client.request(method, CATEGORY, json=body)
+        text = await response.text()
+        self.assertEqual(response.status, status, text)
+        return json.loads(text)
+
+    async def put(self, paths, category, status=200):
+        items = [p if isinstance(p, dict) else self.target(p) for p in paths]
+        return await self.send("PUT", {"items": items, "category": category}, status)
+
+    async def delete(self, paths, status=200):
+        items = [p if isinstance(p, dict) else self.target(p) for p in paths]
+        return await self.send("DELETE", {"items": items}, status)
+
+    async def listed(self, status=200):
+        response = await self.client.get(CATEGORIES)
+        text = await response.text()
+        self.assertEqual(response.status, status, text)
+        return json.loads(text)
+
+    def rows(self):
+        return [
+            tuple(row)
+            for row in self.api.db.model.get_connection().execute(
+                f"SELECT root_key, relative_path, category FROM {TABLE} "
+                "ORDER BY root_key, relative_path"
+            )
+        ]
+
+    def key(self, relative, root=None):
+        roots = self.api._get_all_output_dirs()
+        return image_category_key(roots, (root or self.output_dir) / relative)
+
+
+class TestSetAndClear(ImageCategoryTestCase):
+    async def test_put_sets_listed_and_counts_missing(self):
+        make_png(self.output_dir / "a.png")
+        make_png(self.output_dir / "sub" / "b.png")
+        data = await self.put(["a.png", "sub/b.png", "missing.png"], "  Portraits ")
+        self.assertEqual(
+            data,
+            {
+                "success": True,
+                "category": "Portraits",
+                "updated": 2,
+                "unchanged": 0,
+                "not_found": 1,
+            },
+        )
+        root_key = _favorite_root_key(self.output_dir)
+        self.assertEqual(
+            self.rows(),
+            [(root_key, "a.png", "Portraits"), (root_key, "sub/b.png", "Portraits")],
+        )
+        data = await self.put(["a.png", "sub/b.png"], "Portraits")
+        self.assertEqual((data["updated"], data["unchanged"]), (0, 2))
+        data = await self.put(["a.png"], "Scans")
+        self.assertEqual((data["updated"], data["unchanged"]), (1, 0))
+        self.assertEqual(self.api.db.image_category(*self.key("a.png")), "Scans")
+
+    async def test_put_with_no_listed_file_is_404_and_writes_nothing(self):
+        data = await self.put(["missing.png", "gone/x.png"], "Scans", status=404)
+        self.assertEqual(data["code"], "files-not-found")
+        self.assertEqual(self.rows(), [])
+
+    async def test_duplicate_targets_are_written_and_counted_once(self):
+        make_png(self.output_dir / "a.png")
+        item = self.target("a.png")
+        data = await self.put([item, dict(item)], "Scans")
+        self.assertEqual((data["updated"], data["not_found"]), (1, 0))
+
+    async def test_delete_clears_listed_and_missing_files(self):
+        make_png(self.output_dir / "a.png")
+        await self.put(["a.png"], "Scans")
+        self.api.db.set_image_categories([self.key("gone.png")], "Scans")
+        data = await self.delete(["a.png", "gone.png", "never.png"])
+        self.assertEqual(data, {"success": True, "cleared": 2})
+        self.assertEqual(self.rows(), [])
+        self.assertEqual((await self.delete(["a.png"]))["cleared"], 0)
+
+    async def test_put_never_writes_prompts(self):
+        make_png(self.output_dir / "a.png")
+        make_png(self.output_dir / "b.png")
+        positive = self._save_prompt("castle at dusk", category="Landscapes")
+        self.api.db.link_image_to_prompt(positive, str(self.output_dir / "a.png"))
+        negative = self._save_prompt("blurry", category="Negatives")
+        self.api.db.save_prompt_pair("castle at dusk", "blurry", category="Recipes")
+        conn = self.api.db.model.get_connection()
+        before = category_checksum(conn)
+        links = conn.execute("SELECT * FROM generated_images").fetchall()
+        await self.put(["a.png", "b.png"], "Portraits")
+        await self.listed()
+        await self.delete(["a.png"])
+        self.assertEqual(category_checksum(conn), before)
+        self.assertEqual(
+            conn.execute("SELECT * FROM generated_images").fetchall(), links
+        )
+        self.assertIsNotNone(negative)
+
+    async def test_two_sources_in_one_request(self):
+        second = self.comfy_dir / "renders"
+        make_png(self.output_dir / "a.png")
+        make_png(second / "a.png")
+        self.api._get_all_output_dirs = lambda: [self.output_dir, second]
+        data = await self.put([self.target("a.png", 0), self.target("a.png", 1)], "X")
+        self.assertEqual(data["updated"], 2)
+        self.assertEqual(
+            {row[0] for row in self.rows()},
+            {_favorite_root_key(self.output_dir), _favorite_root_key(second)},
+        )
+        self.assertEqual((await self.delete([self.target("a.png", 1)]))["cleared"], 1)
+
+    async def test_nested_sources_key_by_the_deepest(self):
+        inner = self.output_dir / "nested"
+        make_png(inner / "a.png")
+        self.api._get_all_output_dirs = lambda: [self.output_dir, inner]
+        # Named through the outer source, stored under the inner one.
+        data = await self.put([self.target("nested/a.png", 0)], "Scans")
+        self.assertEqual(data["updated"], 1)
+        self.assertEqual(self.rows(), [(_favorite_root_key(inner), "a.png", "Scans")])
+        # The same file through the inner source is the same key.
+        data = await self.put([self.target("a.png", 1)], "Scans")
+        self.assertEqual(data["unchanged"], 1)
+        self.assertEqual((await self.listed())["categories"], ["Scans"])
+        self.assertEqual((await self.delete([self.target("a.png", 1)]))["cleared"], 1)
+
+    async def test_video_files_can_have_a_category(self):
+        (self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42")
+        data = await self.put(["clip.mp4"], "Clips")
+        self.assertEqual(data["updated"], 1)
+
+    async def test_windows_case_folds_the_key(self):
+        make_png(self.output_dir / "Sub" / "A.png")
+        with patch("py.api.images._path_flavor", ntpath):
+            data = await self.put(["Sub/A.png"], "Scans")
+            self.assertEqual(data["updated"], 1)
+            data = await self.put(["sub/a.PNG"], "Scans")
+            self.assertEqual(data["unchanged"], 1)
+            self.assertEqual(self.rows()[0][1], "sub/a.png")
+            self.assertEqual((await self.delete(["SUB/A.png"]))["cleared"], 1)
+
+    async def test_a_key_beyond_the_stored_bound_is_403(self):
+        make_png(self.output_dir / "a.png")
+        too_long = "/" + "r" * MAX_KEY_LENGTH
+        with patch("py.api.images._favorite_root_key", return_value=too_long):
+            data = await self.put(["a.png"], "Scans", status=403)
+            self.assertNotIn(too_long, json.dumps(data))
+            await self.delete(["a.png"], status=403)
+        self.assertEqual(self.rows(), [])
+
+    async def test_an_unexpected_failure_is_a_logged_500(self):
+        make_png(self.output_dir / "a.png")
+
+        def broken(*_args, **_kwargs):
+            raise sqlite3.OperationalError("secret detail")
+
+        self.api.db.set_image_categories = broken
+        self.api.db.clear_image_categories = broken
+        with self.assertLogs("prompt_manager", level="ERROR"):
+            data = await self.put(["a.png"], "Scans", status=500)
+        self.assertEqual(data["error"], "The category could not be changed.")
+        self.assertNotIn("secret", json.dumps(data))
+        with self.assertLogs("prompt_manager", level="ERROR"):
+            data = await self.delete(["a.png"], status=500)
+        self.assertNotIn("secret", json.dumps(data))
+
+    async def test_delete_category_never_reaches_delete_image(self):
+        make_png(self.output_dir / "a.png")
+        calls = []
+
+        async def recorded(request):
+            calls.append(request.match_info)
+            raise AssertionError("delete_image reached")
+
+        self.api.delete_image = recorded
+        await self.put(["a.png"], "Scans")
+        data = await self.delete(["a.png"])
+        self.assertEqual(data["cleared"], 1)
+        self.assertEqual(calls, [])
+
+
+class TestRefusals(ImageCategoryTestCase):
+    async def test_malformed_bodies_are_400(self):
+        make_png(self.output_dir / "a.png")
+        item = self.target("a.png")
+        cases = [
+            {"items": [item]},
+            {"category": "x"},
+            {"items": [item], "category": "x", "extra": 1},
+            {"items": item, "category": "x"},
+            {"items": [], "category": "x"},
+            {"items": [item] * 501, "category": "x"},
+            {"items": [item], "category": 3},
+            {"items": [item], "category": None},
+            {"items": [item], "category": " \t "},
+            {"items": [item], "category": "x" * 101},
+            {"items": [item], "category": "a\x07b"},
+            {"items": [{**item, "root_index": "0"}], "category": "x"},
+            {"items": [{**item, "root_scope": "nothex"}], "category": "x"},
+            {"items": [{"root_index": 0, "path": "a.png"}], "category": "x"},
+            {"items": [{**item, "root_index": 1}], "category": "x"},
+        ]
+        for body in cases:
+            with self.subTest(body=str(body)[:80]):
+                data = await self.send("PUT", body, status=400)
+                self.assertFalse(data["success"])
+        await self.send("DELETE", {"items": [item], "category": "x"}, status=400)
+        await self.send("DELETE", {"items": []}, status=400)
+        response = await self.client.request(
+            "PUT", CATEGORY, data="not json", headers={"Content-Type": "text/plain"}
+        )
+        self.assertEqual(response.status, 400)
+        self.assertEqual(self.rows(), [])
+
+    async def test_unsafe_paths_are_403(self):
+        (self.output_dir / "notes.txt").write_text("not media")
+        for path in ("../outside.png", "/a.png", "C:/a.png", "a\\b.png", "a/./b.png"):
+            with self.subTest(path=path):
+                await self.put([self.target(path)], "x", status=403)
+        data = await self.put([self.target("notes.txt")], "x", status=403)
+        self.assertEqual(data["error"], "Only media files can have a category")
+        await self.delete([self.target("notes.txt")], status=403)
+
+    async def test_a_stale_scope_is_409(self):
+        make_png(self.output_dir / "a.png")
+        item = {**self.target("a.png"), "root_scope": "0" * 64}
+        data = await self.put([item], "x", status=409)
+        self.assertEqual(data["code"], "source-changed")
+        data = await self.delete([item], status=409)
+        self.assertEqual(data["code"], "source-changed")
+
+    async def test_an_oversized_body_is_413(self):
+        response = await self.client.request(
+            "PUT",
+            CATEGORY,
+            data=b'{"items": [], "category": "' + b"x" * (2 * 1024 * 1024) + b'"}',
+            headers={"Content-Type": "application/json"},
+        )
+        self.assertEqual(response.status, 413)
+
+
+class TestCategoryList(ImageCategoryTestCase):
+    async def test_lists_only_values_whose_file_is_listed(self):
+        make_png(self.output_dir / "a.png")
+        make_png(self.output_dir / "c.png")
+        await self.put(["a.png"], "B")
+        await self.put(["c.png"], "A")
+        self.api.db.set_image_categories([self.key("gone.png")], "Z")
+        self.assertEqual(
+            await self.listed(), {"success": True, "categories": ["A", "B"]}
+        )
+
+    async def test_empty_library_lists_nothing(self):
+        self.assertEqual((await self.listed())["categories"], [])
+
+    async def test_a_failed_read_is_500(self):
+        make_png(self.output_dir / "a.png")
+
+        def broken(*_args, **_kwargs):
+            raise sqlite3.OperationalError("secret detail")
+
+        self.api.db.own_categories_of = broken
+        with self.assertLogs("prompt_manager", level="ERROR"):
+            data = await self.listed(status=500)
+        self.assertEqual(data["error"], "Image categories could not be read.")
+        self.assertNotIn("secret", json.dumps(data))