From 06bd642611d6d54aae8a47485e4b7a2b65a7149c Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Thu, 22 Oct 2026 08:46:45 -0700 Subject: [PATCH] feat(annotations): add scoped observe, capture preview and tag save routes POST /prompt_manager/image-annotations/observe captures the file named by the listing identity {root_index, root_scope, path} (scope checked against the current roots before any file access) and answers the saved tag set of exactly those bytes, refreshing a matching record's verified fingerprint. GET/DELETE .../capture/{token} serve the private copy (no-store, nosniff, pinned with acquire/release_use while it streams) and release it. PUT .../tags re-hashes the current file and writes only when it still holds the captured bytes, under the storage epoch and tag revision (file-changed, storage-changed, tag-conflict, tag-limit, replayed). Bodies are capped at 64 KiB with unknown fields refused; envelopes never carry server paths or locators. Restores invalidate every capture and the host's cleanup removes the capture directory. --- __init__.py | 1 + database/annotations.py | 10 +- py/api/__init__.py | 7 + py/api/annotations.py | 512 ++++++++++++++++ tests/test_api_annotations.py | 1027 +++++++++++++++++++++++++++++++++ tests/test_api_core.py | 14 + 6 files changed, 1568 insertions(+), 3 deletions(-) create mode 100644 py/api/annotations.py create mode 100644 tests/test_api_annotations.py diff --git a/__init__.py b/__init__.py index f4945892..140f7ea7 100644 --- a/__init__.py +++ b/__init__.py @@ -90,6 +90,7 @@ try: host_app = getattr(config.server_instance, "app", None) if host_app is not None and hasattr(host_app, "on_cleanup"): host_app.on_cleanup.append(api.cleanup_history_reads) + host_app.on_cleanup.append(api.cleanup_annotation_captures) except Exception as e: # Log to internal logging system without console spam diff --git a/database/annotations.py b/database/annotations.py index 0a39be2c..219a5039 100644 --- a/database/annotations.py +++ b/database/annotations.py @@ -119,12 +119,16 @@ class AnnotationCommandError(LibraryCommandError): self.current = current +class ImageTagLimitError(ValueError): + """More than ``MAX_IMAGE_TAGS`` distinct names in one set.""" + + def normalize_image_tags(names: object) -> list[str]: """Trim and validate each name, fold case-insensitive repeats, cap at 100. The first occurrence keeps its spelling and position. Raises ``ValueError`` - for anything but a list of valid vocabulary names or for more than - ``MAX_IMAGE_TAGS`` distinct names. + for anything but a list of valid vocabulary names, and its subclass + ``ImageTagLimitError`` for more than ``MAX_IMAGE_TAGS`` distinct names. """ if not isinstance(names, list): raise ValueError("tags must be a list of strings") @@ -138,7 +142,7 @@ def normalize_image_tags(names: object) -> list[str]: seen.add(key) result.append(name) if len(result) > MAX_IMAGE_TAGS: - raise ValueError(f"An image can carry at most {MAX_IMAGE_TAGS} tags") + raise ImageTagLimitError(f"An image can carry at most {MAX_IMAGE_TAGS} tags") return result diff --git a/py/api/__init__.py b/py/api/__init__.py index ce206999..b45b3d58 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -26,6 +26,7 @@ from .prompts import PromptRoutesMixin from .pairs import PairRoutesMixin from .library import LibraryRoutesMixin from .vocabulary import VocabularyRoutesMixin +from .annotations import AnnotationRoutesMixin from .saved_views import SavedViewRoutesMixin from .history import HistoryRoutesMixin from .images import ImageRoutesMixin @@ -252,6 +253,7 @@ class PromptManagerAPI( PairRoutesMixin, LibraryRoutesMixin, VocabularyRoutesMixin, + AnnotationRoutesMixin, SavedViewRoutesMixin, HistoryRoutesMixin, PromptRoutesMixin, @@ -285,6 +287,10 @@ class PromptManagerAPI( self.db = db if db is not None else PromptDatabase() self._history_read_fence = None self._history_reads_closed = False + # Private media captures of the annotation routes (py/api/annotations.py), + # created on first use with their restore subscription. + self._captures = None + self._captures_unsubscribe = None self._cached_output_dir = None # Lazy-cached by _find_comfyui_output_dir() self._html_cache = {} # Cached HTML file contents keyed by path self._gallery_cache = {} # dict: path_str -> (files, timestamp) @@ -441,6 +447,7 @@ class PromptManagerAPI( self._register_pair_routes(routes) self._register_library_routes(routes) self._register_vocabulary_routes(routes) + self._register_annotation_routes(routes) self._register_saved_view_routes(routes) self._register_history_routes(routes) self._register_prompt_routes(routes) diff --git a/py/api/annotations.py b/py/api/annotations.py new file mode 100644 index 00000000..03e647de --- /dev/null +++ b/py/api/annotations.py @@ -0,0 +1,512 @@ +"""Image annotation routes: observe, capture preview and byte-exact tag saves. + +Every route names its file through the listing identity ``{root_index, +root_scope, path}`` and checks the scope against the current gallery roots +before any file access. Observe copies the file into a private capture (the +token the preview and Save refer to), Save re-hashes the current file and +writes only when the bytes are still the captured ones. Responses carry no +server path and no locator. +""" + +import asyncio +import datetime +import os +import re +from pathlib import Path + +from aiohttp import web + +from .images import _forbidden, _has_traversal, _unavailable_link, _verified_title +from .prompts import bad_request + +try: + from ...database.annotations import ( + AnnotationCommandError, + ImageTagLimitError, + normalize_image_tags, + ) + from ...database.execution_history import subscribe_history_restore + from ...database.library_commands import LibraryCommandError + from ...utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore +except ImportError: + from database.annotations import ( + AnnotationCommandError, + ImageTagLimitError, + normalize_image_tags, + ) + from database.execution_history import subscribe_history_restore + from database.library_commands import LibraryCommandError + from utils.media_capture import CAPTURE_EXTENSIONS, 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( + { + "target", + "capture_token", + "sha256", + "expected_epoch", + "expected_tag_revision", + "tags", + } +) +_HEX64 = re.compile(r"[0-9a-f]{64}") +_UUID = re.compile(r"[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12}") +# secrets.token_urlsafe(32) is 43 characters; the bound only rejects junk. +_TOKEN = re.compile(r"[A-Za-z0-9_-]{1,128}") + +_PREVIEW_TYPES = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".webp": "image/webp", + ".gif": "image/gif", + ".bmp": "image/bmp", + ".tiff": "image/tiff", + ".tif": "image/tiff", + ".mp4": "video/mp4", + ".m4v": "video/mp4", + ".webm": "video/webm", + ".mov": "video/quicktime", + ".mkv": "video/x-matroska", + ".avi": "video/x-msvideo", + ".wmv": "video/x-ms-wmv", +} + +_STATUS = { + "review-capacity": 429, + "unsupported-media": 415, + "file-changed": 409, + "source-changed": 409, + "storage-changed": 409, + "tag-conflict": 409, + "capture-expired": 409, +} +_MESSAGES = { + "source-changed": ( + "The source changed. Refresh the image collection before tagging this file." + ), + "file-changed": "The file changed since it was previewed. Reopen the Tag dialog.", + "review-capacity": "Too many previews are open. Close a Tag dialog and try again.", + "storage-changed": "Prompt storage was restored; reopen the Tag dialog.", + "tag-conflict": ( + "The tags of this image changed since the dialog opened; review them first." + ), + "capture-expired": "The preview expired. Open the Tag dialog again.", +} +_OTHER_FILE = "The preview belongs to another file. Open the Tag dialog again." + + +def preview_content_type(extension: str) -> str: + """The media type the preview is served as; never sniffed from bytes.""" + return _PREVIEW_TYPES.get(extension.lower(), "application/octet-stream") + + +def _error(status: int, code: str, message: str, **extra) -> web.Response: + return web.json_response( + {"success": False, "code": code, "error": message, **extra}, + status=status, + headers={"Cache-Control": "no-store"}, + ) + + +def _coded(code: str, message: str | None = None, **extra) -> web.Response: + return _error(_STATUS[code], code, message or _MESSAGES[code], **extra) + + +def _capture_error(exc: CaptureError) -> web.Response: + """The envelope for a capture outcome; messages never name a path.""" + if exc.code not in _STATUS: + return _error(500, exc.code, "The file could not be captured.") + 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) + ) + return moment.replace(microsecond=0).isoformat().replace("+00:00", "Z") + + +def _capture_info(store: CaptureStore, capture) -> dict: + return { + "token": capture.token, + "expires_at": _expires_at(store.expires_in(capture.token)), + "preview_url": CAPTURE_URL + capture.token, + "sha256": capture.sha256, + "size": capture.size, + "width": capture.width, + "height": capture.height, + "media_type": capture.media_type, + "generate": { + "supported": capture.inferable, + "reason": capture.inferable_reason, + }, + } + + +class AnnotationRoutesMixin: + """Observe, capture preview and tag save routes over PR A's storage.""" + + def _register_annotation_routes(self, routes: web.RouteTableDef) -> None: + routes.post("/prompt_manager/image-annotations/observe")(self.observe_image) + routes.get(CAPTURE_URL + "{token}")(self.get_capture_preview) + routes.delete(CAPTURE_URL + "{token}")(self.release_capture) + routes.put("/prompt_manager/image-annotations/tags")(self.save_image_tags) + + # -- capture store lifetime ------------------------------------------------ + + def _capture_store(self) -> CaptureStore: + """The process's capture store, created on first use. + + Created lazily so an API object that never observes a file (most tests, + a server that never opens the Tag dialog) owns no temp directory. A + restore of the database dooms every capture (``invalidate_all``). + """ + if self._captures is None: + self._captures = CaptureStore() + if self._captures_unsubscribe is None: + self._captures_unsubscribe = subscribe_history_restore( + self.db.model.db_path, self._invalidate_captures + ) + return self._captures + + def _invalidate_captures(self) -> None: + store = self._captures + if store is not None: + store.invalidate_all() + + def close_annotation_captures(self) -> None: + """Remove every capture file and the owned directory. Idempotent.""" + unsubscribe, self._captures_unsubscribe = self._captures_unsubscribe, None + if unsubscribe is not None: + unsubscribe() + store, self._captures = self._captures, None + if store is not None: + store.close() + + async def cleanup_annotation_captures(self, app: web.Application) -> None: + self.close_annotation_captures() + + # -- shared request steps -------------------------------------------------- + + async def _annotation_body(self, request: web.Request, fields: frozenset): + data, error = await self._read_json_body(request, max_bytes=MAX_BODY_BYTES) + if error is not None: + return None, error + if set(data) != fields: + return None, bad_request( + "Request must contain exactly " + ", ".join(sorted(fields)) + ) + return data, None + + def _scoped_root(self, index: int, scope: str): + """``(root, None)`` when the scope names a current root, else an error.""" + roots = list(self._get_all_output_dirs()) + error = self._gallery_scope_error( + {"root_scope": scope, "root": str(index)}, roots + ) + if error is not None: + if error.status == 409: + return None, _coded("source-changed") + return None, error + return roots[index], None + + # -- observe --------------------------------------------------------------- + + async def observe_image(self, request: web.Request) -> web.Response: + """POST ``{target}``: capture the file and answer its saved tag set. + + The capture is the byte-exact copy the preview and Save refer to; the + tags are those of exactly these bytes (an empty set at revision 0 for + unseen bytes). Observing never writes tags; equal bytes refresh the + record's verified fingerprint (S1). + """ + data, error = await self._annotation_body(request, _OBSERVE_FIELDS) + if error is not None: + return error + target, error = _target(data["target"]) + if error is not None: + return error + index, scope, relative = target + root, error = self._scoped_root(index, scope) + if error is not None: + return error + store = self._capture_store() + try: + epoch = await self._run_in_executor(self.db.get_execution_epoch) + capture = await self._run_in_executor( + store.capture, + root, + relative, + epoch=epoch, + scope=scope, + root_index=index, + ) + except CaptureError as exc: + return _capture_error(exc) + except Exception: + self.logger.exception("Image observation failed") + return _error(500, "observe-failed", "The image could not be observed.") + try: + stored, title = await self._run_in_executor( + self._observe_capture, store, capture, root, relative + ) + except Exception: + store.release(capture.token) + self.logger.exception("Image observation failed") + return _error(500, "observe-failed", "The image could not be observed.") + return web.json_response( + { + "success": True, + "library_epoch": capture.epoch, + "target": {"root_index": index, "root_scope": scope, "path": relative}, + "capture": _capture_info(store, capture), + "title": title, + "filename": relative.rsplit("/", 1)[-1], + "tags": stored["tags"], + "tag_revision": stored["tag_revision"], + }, + headers={"Cache-Control": "no-store"}, + ) + + def _observe_capture(self, store: CaptureStore, capture, root, relative: str): + """Blocking part of observe: fingerprint refresh, tag read, title.""" + self.db.refresh_image_fingerprint( + capture.locator, capture.sha256, capture.size, capture.mtime_ns + ) + stored = self.db.read_image_tags(capture.locator, capture.sha256) + return stored, self._observed_title(store, capture, root, relative) + + def _observed_title(self, store: CaptureStore, capture, root, relative: str): + """The dialog's title under the card and inspector rule (``_verified_title``). + + A stored Library title always shows; a first-line title only when the + captured bytes' own records confirm it. An ambiguous association has no + title. Videos read no card metadata, so only a stored title shows. + """ + path = str(Path(root) / relative) + summary = self._workspace_link_summaries([path]).get(path) + if summary is None: + summary = _unavailable_link() + positive_title = None + if ( + capture.media_type != "video" + and summary.get("title_source") == "text" + and store.acquire(capture.token) is not None + ): + try: + facts = self._read_card_facts(Path(capture.path), parse=True) + finally: + store.release_use(capture.token) + positive_title = facts.positive_title if facts is not None else None + return _verified_title(summary, positive_title) + + # -- capture preview ------------------------------------------------------- + + async def get_capture_preview(self, request: web.Request) -> web.StreamResponse: + """GET the captured bytes; the token is pinned while they stream.""" + token = request.match_info["token"] + store = self._captures + if store is None or _TOKEN.fullmatch(token) is None: + return _error(404, "capture-expired", _MESSAGES["capture-expired"]) + capture = store.acquire(token) + if capture is None: + return _error(404, "capture-expired", _MESSAGES["capture-expired"]) + try: + _root, error = self._scoped_root(capture.root_index, capture.scope) + if error is not None: + # The roots changed under the token: it is no longer usable. + store.release(token) + return error + try: + handle = await self._run_in_executor(open, capture.path, "rb") + except OSError: + store.release(token) + return _error(404, "capture-expired", _MESSAGES["capture-expired"]) + try: + return await self._stream_capture(request, capture, handle) + finally: + handle.close() + finally: + store.release_use(token) + + @staticmethod + async def _stream_capture(request, capture, handle) -> web.StreamResponse: + response = web.StreamResponse( + status=200, + headers={ + "Content-Type": preview_content_type(Path(capture.path).suffix), + "Content-Length": str(capture.size), + "Cache-Control": "no-store", + "X-Content-Type-Options": "nosniff", + }, + ) + await response.prepare(request) + loop = asyncio.get_running_loop() + while True: + chunk = await loop.run_in_executor(None, handle.read, _PREVIEW_CHUNK) + if not chunk: + break + await response.write(chunk) + await response.write_eof() + return response + + async def release_capture(self, request: web.Request) -> web.Response: + """DELETE a capture token; releasing an unknown token is a success.""" + token = request.match_info["token"] + store = self._captures + if store is not None and _TOKEN.fullmatch(token) is not None: + store.release(token) + return web.json_response( + {"success": True}, headers={"Cache-Control": "no-store"} + ) + + # -- save ------------------------------------------------------------------ + + async def save_image_tags(self, request: web.Request) -> web.Response: + """PUT the tag set of the previewed bytes, re-hashing the current file. + + Writes only when the file still holds the captured bytes + (``file-changed`` with zero writes otherwise), under the storage epoch + and the tag revision the dialog read; an identical retry is + ``replayed``. The hash never runs inside a database transaction. + """ + data, error = await self._annotation_body(request, _SAVE_FIELDS) + if error is not None: + return error + target, error = _target(data["target"]) + if error is not None: + return error + guards, error = _save_guards(data) + if error is not None: + return error + token, sha256, expected_epoch, revision, tags = guards + index, scope, relative = target + root, error = self._scoped_root(index, scope) + if error is not None: + return error + store = self._captures + capture = store.get(token) if store is not None else None + if capture is None: + return _coded("capture-expired") + if capture.root_index != index or capture.scope != scope: + return _coded("capture-expired", _OTHER_FILE) + if sha256 != capture.sha256: + return bad_request("sha256 does not match the preview") + if expected_epoch != capture.epoch: + current = await self._run_in_executor(self.db.get_execution_epoch) + return _coded("storage-changed", library_epoch=current) + try: + facts = await self._run_in_executor( + _verify_captured_file, store, capture, root, relative + ) + except CaptureError as exc: + return _capture_error(exc) + if facts is None or facts.locator != capture.locator: + return _coded("capture-expired", _OTHER_FILE) + if facts.sha256 != capture.sha256: + return _coded("file-changed") + try: + receipt = await self._run_in_executor( + self.db.replace_image_tags, + locator=facts.locator, + sha256=facts.sha256, + size=facts.size, + mtime_ns=facts.mtime_ns, + expected_epoch=expected_epoch, + expected_tag_revision=revision, + tags=tags, + ) + except AnnotationCommandError as exc: + current = exc.current or {} + return _coded( + exc.code, + library_epoch=exc.epoch, + tags=current.get("tags", []), + tag_revision=current.get("tag_revision", 0), + ) + except LibraryCommandError as exc: + return _coded(exc.code, library_epoch=exc.epoch) + except ValueError as exc: + return bad_request(str(exc)) + except Exception: + self.logger.exception("Image tag save failed") + return _error(500, "save-failed", "The tags could not be saved.") + return web.json_response( + {"success": True, **receipt}, headers={"Cache-Control": "no-store"} + ) + + +def _verify_captured_file(store: CaptureStore, capture, root, relative: str): + """Re-hash the current file, or None when the token names another file. + + The cheap locator comparison runs first, so a token sent for another + target never costs a whole-file hash. + """ + expected = os.path.normcase( + os.path.join(os.path.realpath(root), *relative.split("/")) + ) + if expected != capture.locator: + return None + return store.verify(root, relative) + + +def _save_guards(data: dict): + """Validated ``(token, sha256, epoch, revision, tags)`` or ``(None, error)``.""" + token = data["capture_token"] + if not isinstance(token, str) or _TOKEN.fullmatch(token) is None: + return None, bad_request("Invalid capture token") + sha256 = data["sha256"] + if not isinstance(sha256, str) or _HEX64.fullmatch(sha256) is None: + return None, bad_request("sha256 must be 64 lowercase hexadecimal characters") + epoch = data["expected_epoch"] + if not isinstance(epoch, str) or _UUID.fullmatch(epoch) is None: + return None, bad_request("expected_epoch must be a canonical UUID") + revision = data["expected_tag_revision"] + if type(revision) is not int or revision < 0: + return None, bad_request("expected_tag_revision must be a non-negative integer") + try: + tags = normalize_image_tags(data["tags"]) + except ImageTagLimitError as exc: + return None, _error(400, "tag-limit", str(exc)) + except ValueError as exc: + return None, bad_request(str(exc)) + return (token, sha256, epoch, revision, tags), None diff --git a/tests/test_api_annotations.py b/tests/test_api_annotations.py new file mode 100644 index 00000000..2df315fa --- /dev/null +++ b/tests/test_api_annotations.py @@ -0,0 +1,1027 @@ +"""Image annotation routes: observe, capture preview and tag save. + +Every request names its file through the listing identity +``{root_index, root_scope, path}``; the scope is checked before any file +access, tag reads and saves are byte-exact (the capture's digest, re-hashed +at Save) and no response ever carries a server path or a locator. +""" + +import hashlib +import json +import os +import sqlite3 +import sys +import tempfile +import types +import unittest +import uuid +from pathlib import Path +from unittest.mock import patch + +# Discovery imports this module before tests/test_api_autotag.py, whose +# ``folder_paths`` stub (a real module, so ``patch(..., create=True)`` can add +# and remove attributes) the later modules rely on. Install the same stub +# before test_api_images would install a MagicMock in its place. +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 = tempfile.gettempdir + sys.modules["folder_paths"] = _stub + +try: + from test_api_images import ImageAPITestCase, make_png + from open_handles import assert_closed, open_handles_under +except ImportError: + from tests.test_api_images import ImageAPITestCase, make_png + from tests.open_handles import assert_closed, open_handles_under + +from database.execution_history_schema import rotate_history_epoch +from utils import media_capture +from utils.media_capture import CaptureStore + +try: + from py.api import annotations as annotations_module +except ImportError: # the module under test is not there yet (RED) + annotations_module = None + +OBSERVE = "/prompt_manager/image-annotations/observe" +TAGS = "/prompt_manager/image-annotations/tags" +CAPTURE = "/prompt_manager/image-annotations/capture/" +MAX_BODY = 64 * 1024 + +# Columns that an observe may legitimately refresh (S1) or that never matter +# for "nothing changed" comparisons. +_SNAPSHOT_TABLES = { + "prompts": "*", + "generated_images": "*", + "tags": "*", + "prompt_tags": "*", + "prompt_pair_tags": "*", + "image_favorites": "*", + "media_annotation_files": "id, locator", + "media_annotation_bytes": "id, file_id, sha256, first_seen_at, tag_revision", + "media_image_tags": "*", +} + + +class FakeClock: + def __init__(self, start=1000.0): + self.now = start + + def __call__(self): + return self.now + + +class AnnotationAPITestCase(ImageAPITestCase): + """The image API app plus helpers for the annotation routes.""" + + async def get_application(self): + app = await super().get_application() + self.assertIsNotNone( + annotations_module, "py.api.annotations is not implemented" + ) + self.epoch = self.api.db.get_execution_epoch() + return app + + async def tearDownAsync(self): + store = getattr(self.api, "_captures", None) + if store is not None: + # Every preview handle must already be closed (Windows unlink). + assert_closed(self, store.directory) + self.api.close_annotation_captures() + if store is not None: + self.assertFalse(os.path.exists(store.directory)) + await super().tearDownAsync() + + # -- helpers --------------------------------------------------------- + + def scope(self, index=0, roots=None): + roots = roots if roots is not None else self.api._get_all_output_dirs() + return self.api._gallery_root_scope(roots[index], index) + + def target(self, path="sub/a.png", **changes): + return {"root_index": 0, "root_scope": self.scope(), "path": path, **changes} + + async def json_of(self, response): + """The JSON envelope; an unrouted path (no JSON) fails as an assertion.""" + self.assertEqual( + response.content_type, + "application/json", + f"{response.method} {response.url.path} answered {response.status}", + ) + return await response.json() + + async def observe(self, body): + response = await self.client.post(OBSERVE, json=body) + return response.status, await self.json_of(response) + + async def observe_ok(self, path="sub/a.png"): + status, body = await self.observe({"target": self.target(path)}) + self.assertEqual(status, 200, body) + return body + + async def put_tags(self, body): + response = await self.client.put(TAGS, json=body) + return response.status, await self.json_of(response) + + def save_body(self, observed, tags, **changes): + return { + "target": dict(observed["target"]), + "capture_token": observed["capture"]["token"], + "sha256": observed["capture"]["sha256"], + "expected_epoch": observed["library_epoch"], + "expected_tag_revision": observed["tag_revision"], + "tags": list(tags), + **changes, + } + + def locator(self, path): + return os.path.normcase(os.path.realpath(str(self.output_dir / path))) + + def write_bytes(self, path, data, *, keep_mtime=False): + """Rewrite *path* in place, optionally restoring its ``mtime_ns``.""" + full = self.output_dir / path + before = full.stat() + full.write_bytes(data) + if keep_mtime: + os.utime(full, ns=(before.st_atime_ns, before.st_mtime_ns)) + return full + + def db_snapshot(self): + conn = sqlite3.connect(self._temp_db.name) + try: + snapshot = {} + for table, columns in _SNAPSHOT_TABLES.items(): + exists = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", + (table,), + ).fetchone() + if exists is None: + continue + rows = conn.execute(f"SELECT {columns} FROM {table}").fetchall() + snapshot[table] = sorted(map(repr, rows)) + return snapshot + finally: + conn.close() + + def bytes_rows(self): + conn = sqlite3.connect(self._temp_db.name) + try: + return conn.execute( + "SELECT sha256, size, mtime_ns, tag_revision FROM " + "media_annotation_bytes ORDER BY first_seen_at" + ).fetchall() + finally: + conn.close() + + def saved_set(self, path, data, tags): + """Store *tags* for the exact *data* bytes of *path*; return the receipt.""" + full = self.output_dir / path + stat = full.stat() + return self.api.db.replace_image_tags( + locator=self.locator(path), + sha256=hashlib.sha256(data).hexdigest(), + size=stat.st_size, + mtime_ns=stat.st_mtime_ns, + expected_epoch=self.epoch, + expected_tag_revision=0, + tags=tags, + ) + + def use_fake_clock(self): + clock = FakeClock() + self.api.close_annotation_captures() + self.api._captures = CaptureStore(clock=clock) + return clock + + def assert_no_paths(self, body): + text = json.dumps(body) + self.assertNotIn(str(self.comfy_dir), text) + self.assertNotIn(str(self.output_dir), text) + self.assertNotIn(os.path.realpath(self.output_dir), text) + store = getattr(self.api, "_captures", None) + if store is not None: + self.assertNotIn(store.directory, text) + self.assertNotIn("\\\\", text.replace('\\"', "")) + self.assertNotIn(os.sep + "output" + os.sep, text) + + +class TestObserve(AnnotationAPITestCase): + async def test_observe_requires_valid_scope_before_any_file_access(self): + make_png(self.output_dir / "sub" / "a.png") + other = self.comfy_dir / "other" + other.mkdir() + make_png(other / "sub" / "a.png") + stale_scope = self.api._gallery_root_scope(other, 0) + reordered = self.api._gallery_root_scope(self.output_dir, 1) + with patch.object( + media_capture.CaptureStore, "capture", autospec=True + ) as capture: + capture.side_effect = AssertionError("capture must not run") + for name, changes, status in ( + ("stale root", {"root_scope": stale_scope}, 409), + ("reordered root", {"root_scope": reordered}, 409), + ("removed root", {"root_index": 1, "root_scope": reordered}, 409), + ("wrong scope", {"root_scope": "0" * 64}, 409), + ("uppercase scope", {"root_scope": "A" * 64}, 400), + ("short scope", {"root_scope": "abc"}, 400), + ("bool index", {"root_index": True}, 400), + ("negative index", {"root_index": -1}, 400), + ("string index", {"root_index": "0"}, 400), + ("traversal", {"path": "../a.png"}, 403), + ("dot segment", {"path": "./a.png"}, 403), + ("empty segment", {"path": "sub//a.png"}, 403), + ("absolute", {"path": "/etc/a.png"}, 403), + ("backslash", {"path": "sub\\a.png"}, 403), + ("drive", {"path": "C:/a.png"}, 403), + ("drive segment", {"path": "sub/C:a.png"}, 403), + ("nul", {"path": "a\x00.png"}, 403), + ("non-media", {"path": "notes.txt"}, 403), + ("empty path", {"path": ""}, 400), + ("long path", {"path": "a/" * 600 + "a.png"}, 400), + ("int path", {"path": 7}, 400), + ): + with self.subTest(name): + status_, body = await self.observe( + {"target": self.target(**changes)} + ) + self.assertEqual(status_, status, body) + self.assertFalse(body["success"]) + if status == 409: + self.assertEqual(body["code"], "source-changed") + self.assert_no_paths(body) + for name, body in ( + ("missing target", {}), + ("extra field", {"target": self.target(), "extra": 1}), + ("target not an object", {"target": "sub/a.png"}), + ("target missing path", {"target": {"root_index": 0}}), + ("target extra field", {"target": self.target(sha256="x")}), + ): + with self.subTest(name): + status_, body_ = await self.observe(body) + self.assertEqual(status_, 400, body_) + self.assertFalse(body_["success"]) + with self.subTest("body over 64 KiB"): + padding = "x" * MAX_BODY + response = await self.client.post( + OBSERVE, + data=json.dumps({"target": self.target(), "pad": padding}), + headers={"Content-Type": "application/json"}, + ) + self.assertEqual(response.status, 413) + with self.subTest("not json"): + response = await self.client.post( + OBSERVE, data=b"{oops", headers={"Content-Type": "application/json"} + ) + self.assertEqual(response.status, 400) + capture.assert_not_called() + + async def test_observe_returns_saved_tags_for_exact_bytes_only(self): + path = make_png(self.output_dir / "sub" / "a.png") + original = path.read_bytes() + receipt = self.saved_set("sub/a.png", original, ["Forest", "mist"]) + self.assertEqual(receipt["tag_revision"], 1) + + body = await self.observe_ok() + self.assertEqual(body["tags"], ["Forest", "mist"]) + self.assertEqual(body["tag_revision"], 1) + self.assertEqual(body["library_epoch"], self.epoch) + self.assertEqual(body["filename"], "a.png") + self.assertIsNone(body["title"]) + self.assertEqual( + body["target"], + {"root_index": 0, "root_scope": self.scope(), "path": "sub/a.png"}, + ) + capture = body["capture"] + self.assertEqual( + sorted(capture), + sorted( + [ + "token", + "expires_at", + "preview_url", + "sha256", + "size", + "width", + "height", + "media_type", + "generate", + ] + ), + ) + self.assertEqual(capture["sha256"], hashlib.sha256(original).hexdigest()) + self.assertEqual(capture["size"], len(original)) + self.assertEqual((capture["width"], capture["height"]), (4, 4)) + self.assertEqual(capture["media_type"], "image") + self.assertEqual(capture["generate"], {"supported": True, "reason": None}) + self.assertEqual(capture["preview_url"], CAPTURE + capture["token"]) + self.assertTrue(capture["expires_at"].endswith("Z")) + + # Different bytes at the same path: an empty set at revision 0. + changed = make_png(path, color=(0, 255, 0)).read_bytes() + self.assertNotEqual(changed, original) + body = await self.observe_ok() + self.assertEqual((body["tags"], body["tag_revision"]), ([], 0)) + self.assertEqual(body["capture"]["sha256"], hashlib.sha256(changed).hexdigest()) + + # The original bytes again (new mtime): the saved set, and S1 refreshes + # the record's verified fingerprint so cheap reads see it as current. + self.write_bytes("sub/a.png", original) + stat = path.stat() + current = self.api.db.current_image_tags( + {self.locator("sub/a.png"): (stat.st_size, stat.st_mtime_ns)} + ) + body = await self.observe_ok() + self.assertEqual((body["tags"], body["tag_revision"]), (["Forest", "mist"], 1)) + refreshed = self.api.db.current_image_tags( + {self.locator("sub/a.png"): (stat.st_size, stat.st_mtime_ns)} + ) + if stat.st_mtime_ns != self.bytes_rows()[0][2]: + self.assertEqual(current, {}) + self.assertEqual(refreshed, {self.locator("sub/a.png"): ["Forest", "mist"]}) + + async def test_observe_changes_no_tags_revisions_prompts_or_links(self): + path = make_png(self.output_dir / "sub" / "a.png") + data = path.read_bytes() + self.saved_set("sub/a.png", data, ["forest"]) + make_png(self.output_dir / "sub" / "b.png", color=(1, 2, 3)) + self._link_image(self.output_dir / "sub" / "a.png", "a prompt") + prompt_id = self._save_prompt("tagged prompt", tags=["forest", "night"]) + self.assertTrue(prompt_id) + before = self.db_snapshot() + + for _ in range(2): + await self.observe_ok("sub/a.png") + await self.observe_ok("sub/b.png") + + self.assertEqual(self.db_snapshot(), before) + # Unseen bytes (b.png) never create a record. + self.assertEqual(len(self.bytes_rows()), 1) + self.assertEqual(self.bytes_rows()[0][3], 1) + + async def test_observe_title_follows_the_card_rule(self): + make_png(self.output_dir / "sub" / "a.png") + prompt_id = self._save_prompt("Same words\nrest") + self.api.db.link_image_to_prompt(prompt_id, str(self.output_dir / "sub/a.png")) + body = await self.observe_ok() + # A first-line title is withheld unless the file's own records confirm + # it (the card rule); this PNG carries no prompt metadata. + self.assertIsNone(body["title"]) + with self.api.db.model.get_connection() as connection: + connection.execute( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + (prompt_id, "Same words"), + ) + # A stored title is user data and always shows. + body = await self.observe_ok() + self.assertEqual(body["title"], "Same words") + + async def test_observe_video_is_manual_only(self): + clip = self.output_dir / "clip.mp4" + clip.parent.mkdir(exist_ok=True) + clip.write_bytes(b"\x00\x00\x00\x18ftypmp42" + b"\x00" * 64) + body = await self.observe_ok("clip.mp4") + capture = body["capture"] + self.assertEqual(capture["media_type"], "video") + self.assertEqual((capture["width"], capture["height"]), (None, None)) + self.assertEqual( + capture["generate"], + {"supported": False, "reason": media_capture.VIDEO_ONLY_REASON}, + ) + self.assertIsNone(body["title"]) + response = await self.client.get(capture["preview_url"]) + self.assertEqual(response.status, 200) + self.assertEqual(response.headers["Content-Type"], "video/mp4") + self.assertEqual(await response.read(), clip.read_bytes()) + + async def test_observe_maps_capture_outcomes_to_codes(self): + make_png(self.output_dir / "sub" / "a.png") + outside = make_png(self.comfy_dir / "outside.png") + self._make_symlink(outside, self.output_dir / "sub" / "alias.png") + status, body = await self.observe({"target": self.target("sub/alias.png")}) + self.assertEqual((status, body["code"]), (409, "source-changed")) + self.assert_no_paths(body) + + status, body = await self.observe({"target": self.target("sub/missing.png")}) + self.assertEqual((status, body["code"]), (409, "file-changed")) + self.assert_no_paths(body) + + self.api.close_annotation_captures() + self.api._captures = CaptureStore(max_file=16) + status, body = await self.observe({"target": self.target("sub/a.png")}) + self.assertEqual((status, body["code"]), (415, "unsupported-media")) + self.assert_no_paths(body) + + self.api.close_annotation_captures() + self.api._captures = CaptureStore(max_captures=1) + await self.observe_ok("sub/a.png") + status, body = await self.observe({"target": self.target("sub/a.png")}) + self.assertEqual((status, body["code"]), (429, "review-capacity")) + self.assert_no_paths(body) + + +class TestCaptureRoutes(AnnotationAPITestCase): + async def test_capture_route_serves_exact_bytes_no_store_nosniff(self): + path = make_png(self.output_dir / "sub" / "a.png") + source_bytes = path.read_bytes() + body = await self.observe_ok() + token = body["capture"]["token"] + # The source may change after the capture; the preview never does. + make_png(path, color=(9, 9, 9)) + + resp = await self.client.get(body["capture"]["preview_url"]) + self.assertEqual(resp.status, 200) + self.assertEqual(resp.headers["Cache-Control"], "no-store") + self.assertEqual(resp.headers["X-Content-Type-Options"], "nosniff") + self.assertEqual(resp.headers["Content-Type"], "image/png") + self.assertEqual(resp.headers["Content-Length"], str(len(source_bytes))) + self.assertEqual(await resp.read(), source_bytes) + + store = self.api._captures + entry = store._entries[token] + self.assertEqual(entry.uses, 0) + self.assertEqual(open_handles_under(store.directory), []) + self.assertIsNotNone(store.get(token)) + + async def test_capture_route_rejects_unknown_expired_and_rescoped_tokens(self): + make_png(self.output_dir / "sub" / "a.png") + clock = self.use_fake_clock() + for token in ("nope", "A" * 300, "x.y", "%2e%2e%2fx", "a b"): + with self.subTest(token=token): + resp = await self.client.get(CAPTURE + token) + self.assertEqual(resp.status, 404) + data = await self.json_of(resp) + self.assertEqual(data["code"], "capture-expired") + self.assertEqual(resp.headers["Cache-Control"], "no-store") + + body = await self.observe_ok() + url = body["capture"]["preview_url"] + self.assertEqual((await self.client.get(url)).status, 200) + + # Released explicitly: gone, and releasing again stays successful. + resp = await self.client.delete(url) + self.assertEqual( + (resp.status, await self.json_of(resp)), (200, {"success": True}) + ) + self.assertEqual((await self.client.get(url)).status, 404) + resp = await self.client.delete(url) + self.assertEqual( + (resp.status, await self.json_of(resp)), (200, {"success": True}) + ) + + # Idle expiry. + body = await self.observe_ok() + url = body["capture"]["preview_url"] + clock.now += media_capture.DEFAULT_IDLE_SECONDS + 1 + resp = await self.client.get(url) + self.assertEqual(resp.status, 404) + self.assertEqual((await self.json_of(resp))["code"], "capture-expired") + + # The roots changed under a live token: denied and released. + body = await self.observe_ok() + url = body["capture"]["preview_url"] + moved = self.comfy_dir / "moved" + moved.mkdir() + original_roots = self.api._get_all_output_dirs + self.api._get_all_output_dirs = lambda: [moved] + try: + resp = await self.client.get(url) + self.assertEqual(resp.status, 409) + data = await self.json_of(resp) + self.assertEqual(data["code"], "source-changed") + self.assert_no_paths(data) + finally: + self.api._get_all_output_dirs = original_roots + resp = await self.client.get(url) + self.assertEqual(resp.status, 404) + self.assertFalse(any(Path(self.api._captures.directory).iterdir())) + + async def test_preview_pins_the_capture_while_it_streams(self): + path = make_png(self.output_dir / "sub" / "a.png", size=(64, 64)) + source_bytes = path.read_bytes() + body = await self.observe_ok() + token = body["capture"]["token"] + store = self.api._captures + capture_path = store.get(token).path + stream = annotations_module.AnnotationRoutesMixin._stream_capture + seen = {} + + async def release_mid_stream(request, capture, handle): + # A concurrent DELETE (or sweep) lands while the bytes stream. + self.assertEqual(store._entries[token].uses, 1) + self.assertTrue(store.release(token)) + seen["kept"] = os.path.exists(capture_path) + return await stream(request, capture, handle) + + with ( + patch.object( + annotations_module.AnnotationRoutesMixin, + "_stream_capture", + staticmethod(release_mid_stream), + ), + patch.object(annotations_module, "_PREVIEW_CHUNK", 16), + ): + resp = await self.client.get(body["capture"]["preview_url"]) + self.assertEqual(resp.status, 200) + self.assertEqual(await resp.read(), source_bytes) + self.assertTrue(seen["kept"]) + # Removed with the last use, after its handle closed. + self.assertFalse(os.path.exists(capture_path)) + self.assertEqual( + (await self.client.get(body["capture"]["preview_url"])).status, 404 + ) + + async def test_preview_of_a_vanished_capture_file_is_expired(self): + make_png(self.output_dir / "sub" / "a.png") + body = await self.observe_ok() + os.unlink(self.api._captures.get(body["capture"]["token"]).path) + resp = await self.client.get(body["capture"]["preview_url"]) + self.assertEqual(resp.status, 404) + data = await self.json_of(resp) + self.assertEqual(data["code"], "capture-expired") + self.assert_no_paths(data) + self.assertIsNone(self.api._captures.get(body["capture"]["token"])) + + async def test_restore_invalidates_captures(self): + make_png(self.output_dir / "sub" / "a.png") + body = await self.observe_ok() + url = body["capture"]["preview_url"] + self.assertEqual((await self.client.get(url)).status, 200) + backup = str(self.comfy_dir / "backup.db") + self.assertTrue(self.api.db.model.backup_database(backup)) + safety = self.api.db.model.restore_from_file(backup) + try: + resp = await self.client.get(url) + self.assertEqual(resp.status, 404) + self.assertEqual((await self.json_of(resp))["code"], "capture-expired") + status, saved = await self.put_tags(self.save_body(body, ["forest"])) + self.assertEqual((status, saved["code"]), (409, "capture-expired")) + self.assertEqual(self.bytes_rows(), []) + finally: + for path in (safety, backup): + for suffix in ("", "-wal", "-shm"): + if path and os.path.exists(path + suffix): + os.unlink(path + suffix) + + +class TestSaveTags(AnnotationAPITestCase): + async def test_put_saves_then_rereads_the_same_set(self): + make_png(self.output_dir / "sub" / "a.png") + observed = await self.observe_ok() + status, receipt = await self.put_tags( + self.save_body(observed, ["Forest", "mist", "forest"]) + ) + self.assertEqual(status, 200, receipt) + self.assertEqual( + receipt, + { + "success": True, + "library_epoch": self.epoch, + "sha256": observed["capture"]["sha256"], + "tags": ["Forest", "mist"], + "tag_revision": 1, + "replayed": False, + }, + ) + body = await self.observe_ok() + self.assertEqual((body["tags"], body["tag_revision"]), (["Forest", "mist"], 1)) + + async def test_put_rehashes_current_file_and_rejects_changed_bytes(self): + path = make_png(self.output_dir / "sub" / "a.png") + original = path.read_bytes() + observed = await self.observe_ok() + before = self.db_snapshot() + # Same size, same mtime, other bytes: only the digest tells. + replacement = bytearray(original) + replacement[-1] ^= 0xFF + self.write_bytes("sub/a.png", bytes(replacement), keep_mtime=True) + self.assertEqual(path.stat().st_size, len(original)) + + status, body = await self.put_tags(self.save_body(observed, ["forest"])) + self.assertEqual((status, body["code"]), (409, "file-changed")) + self.assertEqual(self.db_snapshot(), before) + self.assert_no_paths(body) + + # Reopening shows the new bytes as an unseen, empty set. + reopened = await self.observe_ok() + self.assertEqual((reopened["tags"], reopened["tag_revision"]), ([], 0)) + self.assertNotEqual( + reopened["capture"]["sha256"], observed["capture"]["sha256"] + ) + + # A vanished file is also file-changed with zero writes. + path.unlink() + status, body = await self.put_tags(self.save_body(reopened, ["forest"])) + self.assertEqual((status, body["code"]), (409, "file-changed")) + self.assertEqual(self.db_snapshot(), before) + + async def test_put_rejects_capture_for_other_target_or_epoch(self): + make_png(self.output_dir / "sub" / "a.png") + make_png(self.output_dir / "sub" / "b.png", color=(0, 0, 255)) + observed_a = await self.observe_ok("sub/a.png") + observed_b = await self.observe_ok("sub/b.png") + before = self.db_snapshot() + + # b's token for a's target: refused before any whole-file hash. + body = self.save_body(observed_a, ["forest"]) + body["capture_token"] = observed_b["capture"]["token"] + body["sha256"] = observed_b["capture"]["sha256"] + with patch.object( + media_capture.CaptureStore, "verify", autospec=True + ) as verify: + status, result = await self.put_tags(body) + verify.assert_not_called() + self.assertEqual((status, result["code"]), (409, "capture-expired")) + + # A digest that is not the capture's. + body = self.save_body(observed_a, ["forest"], sha256="0" * 64) + status, result = await self.put_tags(body) + self.assertEqual(status, 400) + + # An unknown token. + body = self.save_body(observed_a, ["forest"], capture_token="unknown-token") + status, result = await self.put_tags(body) + self.assertEqual((status, result["code"]), (409, "capture-expired")) + + # Another epoch than the capture's. + body = self.save_body(observed_a, ["forest"], expected_epoch=str(uuid.uuid4())) + status, result = await self.put_tags(body) + self.assertEqual((status, result["code"]), (409, "storage-changed")) + self.assertEqual(result["library_epoch"], self.epoch) + + # A target the token was not captured for (scope of another root). + other = self.comfy_dir / "other" + other.mkdir() + make_png(other / "sub" / "a.png") + roots = [self.output_dir, other] + self.api._get_all_output_dirs = lambda: roots + body = self.save_body(observed_a, ["forest"]) + body["target"] = { + "root_index": 1, + "root_scope": self.scope(1, roots), + "path": "sub/a.png", + } + status, result = await self.put_tags(body) + self.assertEqual((status, result["code"]), (409, "capture-expired")) + + self.assertEqual(self.db_snapshot(), before) + for result_ in (result,): + self.assert_no_paths(result_) + + async def test_put_conflict_storage_changed_and_limit_codes(self): + make_png(self.output_dir / "sub" / "a.png") + observed = await self.observe_ok() + status, first = await self.put_tags(self.save_body(observed, ["forest"])) + self.assertEqual(status, 200, first) + + # A stale revision with a different set: the current set comes back. + status, body = await self.put_tags(self.save_body(observed, ["woods"])) + self.assertEqual((status, body["code"]), (409, "tag-conflict")) + self.assertEqual((body["tags"], body["tag_revision"]), (["forest"], 1)) + self.assertEqual(body["library_epoch"], self.epoch) + + # Validation at the boundary, before any file or database work. + with patch.object( + media_capture.CaptureStore, "verify", autospec=True + ) as verify: + verify.side_effect = AssertionError("verify must not run") + cases = ( + ("too many tags", {"tags": [f"t{i}" for i in range(101)]}, 400), + ("tag not a string", {"tags": ["ok", 3]}, 400), + ("tags not a list", {"tags": "forest"}, 400), + ("blank tag", {"tags": [" "]}, 400), + ("bool revision", {"expected_tag_revision": True}, 400), + ("negative revision", {"expected_tag_revision": -1}, 400), + ("bad epoch", {"expected_epoch": "not-a-uuid"}, 400), + ("bad digest", {"sha256": "xyz"}, 400), + ("bad token type", {"capture_token": 5}, 400), + ("traversal", {"target": self.target("../a.png")}, 403), + ("wrong scope", {"target": self.target(root_scope="f" * 64)}, 409), + ) + for name, changes, expected in cases: + with self.subTest(name): + body_ = {**self.save_body(observed, ["forest"]), **changes} + status, result = await self.put_tags(body_) + self.assertEqual(status, expected, result) + self.assertFalse(result["success"]) + if name == "too many tags": + self.assertEqual(result["code"], "tag-limit") + with self.subTest("extra field"): + body_ = self.save_body(observed, ["forest"]) + body_["extra"] = 1 + status, result = await self.put_tags(body_) + self.assertEqual(status, 400) + with self.subTest("missing field"): + body_ = self.save_body(observed, ["forest"]) + del body_["tags"] + status, result = await self.put_tags(body_) + self.assertEqual(status, 400) + with self.subTest("body over 64 KiB"): + body_ = self.save_body(observed, ["forest"]) + body_["tags"] = ["x" * 60] * 1200 + response = await self.client.put( + TAGS, + data=json.dumps(body_), + headers={"Content-Type": "application/json"}, + ) + self.assertEqual(response.status, 413) + verify.assert_not_called() + + # Storage restored under the dialog: storage-changed, zero writes. + before = self.db_snapshot() + fresh = await self.observe_ok() + rotate_history_epoch(self.api.db.model.get_connection()) + status, body = await self.put_tags(self.save_body(fresh, ["woods"])) + self.assertEqual((status, body["code"]), (409, "storage-changed")) + self.assertEqual(body["library_epoch"], self.api.db.get_execution_epoch()) + self.assertNotEqual(body["library_epoch"], self.epoch) + self.assertEqual(self.db_snapshot(), before) + + async def test_put_replay_returns_replayed_true(self): + make_png(self.output_dir / "sub" / "a.png") + observed = await self.observe_ok() + body = self.save_body(observed, ["forest", "mist"]) + status, first = await self.put_tags(body) + self.assertEqual( + (status, first["replayed"], first["tag_revision"]), (200, False, 1) + ) + before = self.db_snapshot() + + status, again = await self.put_tags(body) + self.assertEqual(status, 200, again) + self.assertEqual(again, {**first, "replayed": True}) + self.assertEqual(self.db_snapshot(), before) + + # The same set at the new revision is a replay as well. + body["expected_tag_revision"] = 1 + status, third = await self.put_tags(body) + self.assertEqual( + (status, third["replayed"], third["tag_revision"]), (200, True, 1) + ) + self.assertEqual(self.db_snapshot(), before) + + async def test_put_saves_only_the_target_byte_set(self): + """Tagging one file never touches a sibling that shares its prompt.""" + a = make_png(self.output_dir / "sub" / "a.png") + make_png(self.output_dir / "sub" / "b.png", color=(2, 2, 2)) + twin = self.output_dir / "sub" / "twin.png" + twin.write_bytes(a.read_bytes()) + prompt_id = self._save_prompt("shared prompt", tags=["shared"]) + for name in ("a.png", "b.png", "twin.png"): + self.api.db.link_image_to_prompt( + prompt_id, str(self.output_dir / "sub" / name) + ) + observed_b = await self.observe_ok("sub/b.png") + self.saved_set( + "sub/b.png", (self.output_dir / "sub" / "b.png").read_bytes(), ["b"] + ) + legacy_before = { + k: v + for k, v in self.db_snapshot().items() + if k + not in ( + "media_annotation_files", + "media_annotation_bytes", + "media_image_tags", + ) + } + + observed = await self.observe_ok("sub/a.png") + status, receipt = await self.put_tags(self.save_body(observed, ["forest"])) + self.assertEqual(status, 200, receipt) + + after = self.db_snapshot() + legacy_after = { + k: v + for k, v in after.items() + if k + not in ( + "media_annotation_files", + "media_annotation_bytes", + "media_image_tags", + ) + } + # The tag registry gained "forest"; prompts, links and prompt tags are equal. + self.assertEqual( + {k: v for k, v in legacy_after.items() if k != "tags"}, + {k: v for k, v in legacy_before.items() if k != "tags"}, + ) + self.assertEqual((await self.observe_ok("sub/b.png"))["tags"], ["b"]) + self.assertEqual((await self.observe_ok("sub/twin.png"))["tags"], []) + self.assertEqual(observed_b["tags"], []) + + +class TestFailures(AnnotationAPITestCase): + async def test_unexpected_failures_answer_generic_envelopes(self): + make_png(self.output_dir / "sub" / "a.png") + secret = str(self.output_dir / "sub" / "a.png") + with patch.object( + media_capture.CaptureStore, + "capture", + autospec=True, + side_effect=media_capture.CaptureError(secret), + ): + status, body = await self.observe({"target": self.target()}) + self.assertEqual((status, body["code"]), (500, "capture-error")) + self.assert_no_paths(body) + + with ( + patch.object( + media_capture.CaptureStore, + "capture", + autospec=True, + side_effect=OSError(5, "I/O error", secret), + ), + self.assertLogs("prompt_manager.api", "ERROR"), + ): + status, body = await self.observe({"target": self.target()}) + self.assertEqual((status, body["code"]), (500, "observe-failed")) + self.assert_no_paths(body) + + with ( + patch.object( + self.api.db, "read_image_tags", side_effect=RuntimeError(secret) + ), + self.assertLogs("prompt_manager.api", "ERROR"), + ): + status, body = await self.observe({"target": self.target()}) + self.assertEqual((status, body["code"]), (500, "observe-failed")) + self.assert_no_paths(body) + # The capture of the failed observation is released at once. + self.assertEqual(self.api._captures.stats()["captures"], 0) + + observed = await self.observe_ok() + with ( + patch.object( + self.api.db, "replace_image_tags", side_effect=RuntimeError(secret) + ), + self.assertLogs("prompt_manager.api", "ERROR"), + ): + status, body = await self.put_tags(self.save_body(observed, ["forest"])) + self.assertEqual((status, body["code"]), (500, "save-failed")) + self.assert_no_paths(body) + self.assertEqual(self.bytes_rows(), []) + + async def test_title_lookup_failure_still_observes(self): + make_png(self.output_dir / "sub" / "a.png") + with patch.object( + self.api.db, "get_image_link_summaries", side_effect=RuntimeError("db") + ): + body = await self.observe_ok() + self.assertIsNone(body["title"]) + self.assertEqual((body["tags"], body["tag_revision"]), ([], 0)) + + async def test_save_of_a_released_capture_is_expired(self): + make_png(self.output_dir / "sub" / "a.png") + observed = await self.observe_ok() + resp = await self.client.delete(observed["capture"]["preview_url"]) + self.assertEqual(resp.status, 200) + status, body = await self.put_tags(self.save_body(observed, ["forest"])) + self.assertEqual((status, body["code"]), (409, "capture-expired")) + self.assertEqual(self.bytes_rows(), []) + + async def test_cleanup_removes_the_capture_directory_and_subscription(self): + import database.execution_history as observers + + make_png(self.output_dir / "sub" / "a.png") + await self.observe_ok() + directory = self.api._captures.directory + identity = os.path.normcase(os.path.realpath(self.api.db.model.db_path)) + self.assertEqual(len(observers._restore_observers[identity]), 1) + await self.api.cleanup_annotation_captures(None) + await self.api.cleanup_annotation_captures(None) + self.assertFalse(os.path.exists(directory)) + self.assertNotIn(identity, observers._restore_observers) + # A later observation starts a fresh store. + await self.observe_ok() + self.assertNotEqual(self.api._captures.directory, directory) + + +class TestEnvelopes(AnnotationAPITestCase): + async def test_responses_never_contain_absolute_paths_or_locators(self): + make_png(self.output_dir / "sub" / "a.png") + bodies = [] + observed = await self.observe_ok() + bodies.append(observed) + status, body = await self.observe({"target": self.target("sub/missing.png")}) + bodies.append(body) + status, body = await self.observe({"target": self.target("../a.png")}) + bodies.append(body) + status, body = await self.put_tags(self.save_body(observed, ["forest"])) + bodies.append(body) + status, body = await self.put_tags(self.save_body(observed, ["woods"])) + bodies.append(body) + resp = await self.client.get(CAPTURE + "missing") + bodies.append(await self.json_of(resp)) + resp = await self.client.delete(observed["capture"]["preview_url"]) + bodies.append(await self.json_of(resp)) + status, body = await self.put_tags(self.save_body(observed, ["forest"])) + bodies.append(body) + for body in bodies: + self.assert_no_paths(body) + self.assertNotIn(self.locator("sub/a.png"), json.dumps(body)) + self.assertNotIn("locator", json.dumps(body)) + + async def test_legacy_inspect_and_autotag_shapes_unchanged(self): + make_png(self.output_dir / "image.png") + resp = await self.client.get( + "/prompt_manager/images/inspect/image.png", params={"root": "0"} + ) + self.assertEqual(resp.status, 200) + data = await self.json_of(resp) + self.assertEqual( + sorted(data), + [ + "file", + "image", + "inspection_version", + "library", + "metadata", + "prompts", + "resolution", + "root_scope", + "success", + ], + ) + self.assertEqual(sorted(data["image"]), ["filename", "path", "root_index"]) + self.assertEqual( + sorted(data["library"]), + [ + "category", + "category_scope", + "pair_id", + "positive_id", + "scope", + "status", + "tags", + "tags_scope", + "title", + ], + ) + + class FakeService: + custom_prompt = "" + default_prompt = "p" + wd14_general_threshold = 0.35 + wd14_character_threshold = 0.85 + + def get_models_status(self): + return {"wd14-vit": {"downloaded": False, "model_path": None}} + + def is_model_loaded(self): + return False + + def get_loaded_model_type(self): + return None + + with patch("py.autotag.get_autotag_service", return_value=FakeService()): + resp = await self.client.get("/prompt_manager/autotag/models") + self.assertEqual(resp.status, 200) + self.assertEqual( + sorted(await self.json_of(resp)), + [ + "default_prompt", + "loaded_model_type", + "model_loaded", + "models", + "success", + "wd14_character_threshold", + "wd14_general_threshold", + ], + ) + + def test_preview_types_cover_every_capture_extension(self): + for extension in media_capture.CAPTURE_EXTENSIONS: + kind = annotations_module.preview_content_type(extension) + self.assertTrue(kind.startswith(("image/", "video/")), extension) + self.assertEqual( + annotations_module.preview_content_type(".exe"), "application/octet-stream" + ) + + +class AnnotationHostCleanupTests(unittest.IsolatedAsyncioTestCase): + async def test_root_registers_capture_cleanup_on_the_host_application(self): + from aiohttp import web + + try: + from tests.test_init_package import fake_server_module, load_package + except ImportError: + from test_init_package import fake_server_module, load_package + + with tempfile.TemporaryDirectory() as directory: + server = fake_server_module() + server.PromptServer.instance.app = web.Application() + with load_package( + "annotation_cleanup_pkg", + str(Path(directory) / "db.sqlite"), + directory, + server, + ) as (package, host): + api = package.api + self.assertIn(api.cleanup_annotation_captures, host.app.on_cleanup) + store = api._capture_store() + host.app.freeze() + await host.app.cleanup() + self.assertFalse(os.path.exists(store.directory)) + self.assertIsNone(api._captures) diff --git a/tests/test_api_core.py b/tests/test_api_core.py index 4a5d9e76..14f7c464 100644 --- a/tests/test_api_core.py +++ b/tests/test_api_core.py @@ -448,6 +448,20 @@ class TestUiRoutes(CoreRoutesTestCase): resp = await self.client.request("GET", "/prompt_manager/test") self.assertTrue((await resp.json())["success"]) + async def test_image_annotation_routes_are_registered(self): + table = { + (route.method, route.resource.canonical) + for route in self.app.router.routes() + if route.resource is not None + } + for entry in ( + ("POST", "/prompt_manager/image-annotations/observe"), + ("GET", "/prompt_manager/image-annotations/capture/{token}"), + ("DELETE", "/prompt_manager/image-annotations/capture/{token}"), + ("PUT", "/prompt_manager/image-annotations/tags"), + ): + self.assertIn(entry, table) + class TestStaticRoutes(CoreRoutesTestCase):