Merge remote-tracking branch 'origin/main' into feat/247-foreign-workflows-pr2

py/api/__init__.py

@@ -74,6 +74,7 @@ from .autotag_routes import AutotagRoutesMixin
 from .lora_integration import LoraIntegrationMixin
 from .lora_import_routes import LoraImportRoutesMixin
 from .image_categories import ImageCategoryRoutesMixin
+from .source_rescope import SourceRescopeMixin
 
 try:
     from ...database.operations import PromptDatabase
@@ -342,6 +343,7 @@ class PromptManagerAPI(
     LoraIntegrationMixin,
     LoraImportRoutesMixin,
     ImageCategoryRoutesMixin,
+    SourceRescopeMixin,
 ):
     """REST API handler for PromptManager operations and web interface.
 
@@ -708,6 +710,7 @@ class PromptManagerAPI(
         self._register_lora_routes(routes)
         self._register_lora_import_routes(routes)
         self._register_image_category_routes(routes)
+        self._register_rescope_routes(routes)
 
         # Register gzip compression middleware (once)
         global _gzip_registered

py/api/images.py

@@ -13,6 +13,7 @@ import re
 import stat
 import subprocess  # nosec B404 - ffmpeg frames through run_encoder, fixed argv
 import sys
+import threading
 import time as _time
 import urllib.parse
 from bisect import bisect_left, bisect_right, insort
@@ -81,6 +82,7 @@ try:
         reveal_executor,
         reveal_in_file_manager,
     )
+    from ...utils.workspace_keys import gallery_root_scope, workspace_keys_for
 except ImportError:
     from database.image_collections import CollectionMissing
     from database.image_favorites import MAX_KEY_LENGTH
@@ -105,6 +107,7 @@ except ImportError:
         reveal_executor,
         reveal_in_file_manager,
     )
+    from utils.workspace_keys import gallery_root_scope, workspace_keys_for
 
 
 IMAGE_EXTENSIONS = _IMAGE_EXTENSIONS
@@ -518,7 +521,8 @@ MAX_JSON_BODY_BYTES = 1024 * 1024
 _READ_CHUNK = 64 * 1024
 
 
-_WORKSPACE_CURSOR_SECRET = os.urandom(32)
+# Guards the first load of an API's workspace keys (#430, _workspace_keys).
+_WORKSPACE_KEYS_LOCK = threading.Lock()
 
 
 def _workspace_root(query: dict, roots: list[Path]) -> int | None:
@@ -618,15 +622,37 @@ def _rated_key(entry: _Listed, rank: int) -> tuple:
     return (-rank, *_workspace_key(entry, "rating"))
 
 
-def _workspace_cursor(key: tuple, scope: str) -> str:
+# Bump whenever the cursor's key tuple changes shape: a cursor now outlives a
+# restart and an upgrade (#430), so an old shape must be refused, not parsed.
+_CURSOR_VERSION = 2
+_CURSOR_STALE = "The image collection changed; refresh the gallery"
+
+
+class CursorStale(ValueError):
+    """A well-formed cursor this listing cannot verify (#430).
+
+    Another key, another root list or other filters: the page refreshes its
+    collection instead of retrying the same cursor.
+    """
+
+
+def _cursor_mac(secret: bytes, scope: str, payload: bytes) -> bytes:
+    # The listing scope (it hashes every absolute root) is bound in the MAC
+    # only, never sent: the scope is 64 hex, so the "\n" join is unambiguous.
+    return hmac.new(secret, scope.encode() + b"\n" + payload, hashlib.sha256).digest()
+
+
+def _workspace_cursor(key: tuple, scope: str, secret: bytes) -> str:
+    """The signed keyset cursor of *key*; *secret* is the API's cursor key."""
     payload = json.dumps(
-        {"v": 1, "scope": scope, "key": key}, separators=(",", ":"), ensure_ascii=False
+        {"v": _CURSOR_VERSION, "key": key}, separators=(",", ":"), ensure_ascii=False
     ).encode()
-    signature = hmac.new(_WORKSPACE_CURSOR_SECRET, payload, hashlib.sha256).digest()
+    signature = _cursor_mac(secret, scope, payload)
     return base64.urlsafe_b64encode(signature + payload).decode().rstrip("=")
 
 
-def _workspace_cursor_key(token: str, scope: str, order: str) -> tuple:
+def _workspace_cursor_key(token: str, scope: str, order: str, secret: bytes) -> tuple:
+    """The keyset key of *token*; CursorStale when it is not this listing's."""
     try:
         if len(token) > 4096:
             raise ValueError()
@@ -634,13 +660,14 @@ def _workspace_cursor_key(token: str, scope: str, order: str) -> tuple:
             token + "=" * (-len(token) % 4), altchars=b"-_", validate=True
         )
         signature, payload = raw[:32], raw[32:]
-        expected = hmac.new(_WORKSPACE_CURSOR_SECRET, payload, hashlib.sha256).digest()
-        if not hmac.compare_digest(signature, expected):
+        if len(signature) != 32 or not payload:
             raise ValueError()
+        if not hmac.compare_digest(signature, _cursor_mac(secret, scope, payload)):
+            raise CursorStale(_CURSOR_STALE)
         data = json.loads(payload)
         key = data["key"]
-        if data["v"] != 1 or data["scope"] != scope or not isinstance(key, list):
-            raise ValueError()
+        if data["v"] != _CURSOR_VERSION or not isinstance(key, list):
+            raise CursorStale(_CURSOR_STALE)
         types = (
             (str, str, int, (float, int))
             if order in ("name", "name-desc")
@@ -654,6 +681,8 @@ def _workspace_cursor_key(token: str, scope: str, order: str) -> tuple:
         ):
             raise ValueError()
         return tuple(key)
+    except CursorStale:
+        raise
     except (ValueError, KeyError, TypeError, UnicodeError, binascii.Error):
         raise ValueError("Invalid gallery cursor; refresh the gallery") from None
 
@@ -1338,13 +1367,27 @@ def _in_subfolder(rel_dir, subfolder):
 class ImageRoutesMixin:
     """Mixin providing image and gallery-related API endpoints."""
 
+    def _workspace_keys(self):
+        """This API's scope and cursor keys, loaded once (#430).
+
+        They derive from the secret persisted beside the database, so a restart
+        keeps every scope and cursor an open page holds. The first call reads
+        (or writes) one small file; later calls return the cached keys, which
+        a live library move or restore does not change before a restart.
+        """
+        keys = self.__dict__.get("_workspace_keys_cache")
+        if keys is None:
+            with _WORKSPACE_KEYS_LOCK:
+                keys = self.__dict__.get("_workspace_keys_cache")
+                if keys is None:
+                    db = getattr(self, "db", None)
+                    model = getattr(db, "model", None)
+                    keys = workspace_keys_for(getattr(model, "db_path", None))
+                    self.__dict__["_workspace_keys_cache"] = keys
+        return keys
+
     def _gallery_root_scope(self, root: Path, index: int) -> str:
-        payload = json.dumps(
-            ["prompt-manager-gallery-root-v1", index, str(_canonical(root))],
-            separators=(",", ":"),
-            ensure_ascii=False,
-        ).encode("utf-8")
-        return hmac.new(_WORKSPACE_CURSOR_SECRET, payload, hashlib.sha256).hexdigest()
+        return gallery_root_scope(self._workspace_keys(), str(_canonical(root)), index)
 
     def _gallery_scope_error(
         self, query: Mapping, roots: list[Path]
@@ -2008,6 +2051,8 @@ class ImageRoutesMixin:
             limit, _offset = parse_page_params(request.query, default_limit=100)
         except ValueError:
             return bad_request("limit and offset must be integers")
+        # #430: the first use may create the key file; never on the event loop.
+        cursor_secret = (await self._run_in_executor(self._workspace_keys)).cursor
         try:
             root_index = _workspace_root(request.query, output_dirs)
             filters = _workspace_filters(request.query)
@@ -2031,7 +2076,15 @@ class ImageRoutesMixin:
             ).hexdigest()
             token = request.query.get("cursor")
             after = (
-                _workspace_cursor_key(token, scope, filters["order"]) if token else None
+                _workspace_cursor_key(token, scope, filters["order"], cursor_secret)
+                if token
+                else None
+            )
+        except CursorStale as exc:
+            return web.json_response(
+                {"success": False, "code": "cursor-stale", "error": str(exc)},
+                status=400,
+                headers={"Cache-Control": "no-store"},
             )
         except ValueError as exc:
             return bad_request(str(exc))
@@ -2472,7 +2525,11 @@ class ImageRoutesMixin:
             positive_count = _COUNT_OMITTED
         has_more = len(page) > limit
         page = page[:limit]
-        next_cursor = _workspace_cursor(key(page[-1]), scope) if has_more else None
+        next_cursor = (
+            _workspace_cursor(key(page[-1]), scope, self._workspace_keys().cursor)
+            if has_more
+            else None
+        )
         return (
             [(entry.path, entry.mtime, entry.root, entry.index) for entry in page],
             total,

py/api/prompts.py

@@ -8,8 +8,10 @@ import sqlite3
 from aiohttp import web
 
 try:
+    from ...database.library_commands import LibraryCommandError
     from ...database.pairs import PairDependencyError
 except ImportError:
+    from database.library_commands import LibraryCommandError
     from database.pairs import PairDependencyError
 from urllib.parse import quote
 
@@ -1128,7 +1130,14 @@ class PromptRoutesMixin:
     async def bulk_add_tags(self, request):
         """Bulk add tags to prompts."""
         try:
-            data = await request.json()
+            try:
+                data = await request.json()
+            except ValueError:
+                return bad_request("Invalid JSON body")
+            if not isinstance(data, dict):
+                return bad_request("Expected a JSON object")
+            if "expected_epoch" in data:
+                return await self._bulk_library_tags(data)
             prompt_ids = data.get("prompt_ids", [])
             new_tags = data.get("tags", [])
 
@@ -1180,6 +1189,8 @@ class PromptRoutesMixin:
                 )
             if "images_filter" in data:
                 return await self._bulk_category_by_filter(data)
+            if "expected_epoch" in data:
+                return await self._bulk_library_category(data)
             if not data.get("prompt_ids"):
                 return web.json_response(
                     {"success": False, "error": "No prompt IDs provided"}, status=400
@@ -1216,6 +1227,54 @@ class PromptRoutesMixin:
                 status=500,
             )
 
+    async def _bulk_library_tags(self, data: dict) -> web.Response:
+        """The workspace Library's fenced form (#442).
+
+        ``{prompt_ids, tags, expected_epoch}``: prompt tags only, never image tags.
+        """
+        return await self._fenced_library_write(
+            data, "tags", self.db.bulk_add_library_tags
+        )
+
+    async def _fenced_library_write(
+        self, data: dict, field: str, write
+    ) -> web.Response:
+        """One fenced Library bulk write: exact keys, checked ids, 409 on a restore."""
+        fields = {"prompt_ids", field, "expected_epoch"}
+        if set(data) != fields:
+            return bad_request(
+                "Request must contain exactly " + ", ".join(sorted(fields))
+            )
+        try:
+            prompt_ids = parse_bulk_prompt_ids(data["prompt_ids"])
+            result = await self._run_in_executor(
+                write, prompt_ids, data[field], expected_epoch=data["expected_epoch"]
+            )
+        except LibraryCommandError as exc:
+            return web.json_response(
+                {
+                    "success": False,
+                    "code": exc.code,
+                    "library_epoch": exc.epoch,
+                    "error": "Prompt storage was restored; refresh the library",
+                },
+                status=409,
+            )
+        except ValueError as exc:
+            return bad_request(str(exc))
+        return web.json_response({"success": True, **result})
+
+    async def _bulk_library_category(self, data: dict) -> web.Response:
+        """The workspace Library's fenced form (#437).
+
+        ``{prompt_ids, category, expected_epoch}``: a non-blank category on
+        the selected Library parents, refused with 409 ``storage-changed``
+        when prompt storage was restored since the page read its epoch.
+        """
+        return await self._fenced_library_write(
+            data, "category", self.db.bulk_set_library_category
+        )
+
     async def _bulk_category_by_filter(self, data: dict) -> web.Response:
         """Set the category on every linked positive a listing filter admits.
 

py/api/source_rescope.py

@@ -0,0 +1,100 @@
+"""Where a stale Images target is now (#430): ``POST /prompt_manager/images/rescope``.
+
+A page addresses a listed file as ``{root_index, root_scope, path}``. When the
+sources are renumbered its scope no longer names the root at that index, and
+every per-file action answers 409 ``source-changed``. Before it gives up, the
+page asks here. The target is mapped only when its scope is, at its old
+index, the scope of one of the current roots: the server issued that scope
+for exactly that folder, so the answer (the file's current listing
+descriptor) names nothing the page could not already address. A removed or
+replaced root, a scope of another key and a foreign scope stay refused.
+
+Only the Images roots (``_get_all_output_dirs``) are mapped; a LoRA-linked
+root's target is never re-scoped here.
+"""
+
+import hmac
+from pathlib import Path
+
+from aiohttp import web
+
+from .images import (
+    _forbidden,
+    _is_thumbnail_cache_path,
+    _json_body_or_error,
+    _media_access_error,
+)
+from .media_target import media_target
+from .prompts import bad_request
+
+# Far above any real source count; keeps the index out of huge-int territory.
+MAX_ROOT_INDEX = 4095
+SOURCE_GONE = (
+    "The source of this image changed or was removed. Refresh the image "
+    "collection to find it again."
+)
+
+
+def _coded(status: int, code: str, message: str) -> web.Response:
+    return web.json_response(
+        {"success": False, "code": code, "error": message}, status=status
+    )
+
+
+def _no_store(response: web.Response) -> web.Response:
+    response.headers["Cache-Control"] = "no-store"
+    return response
+
+
+class SourceRescopeMixin:
+    """The rescope route; composed into ``PromptManagerAPI``."""
+
+    def _register_rescope_routes(self, routes):
+        @routes.post("/prompt_manager/images/rescope")
+        async def rescope_route(request):
+            return _no_store(await self.rescope_image(request))
+
+    def _current_root_for(self, scope: str, old_index: int, roots: list) -> int | None:
+        """The index of the first current root *scope* was issued for at *old_index*."""
+        for index, root in enumerate(roots):
+            if hmac.compare_digest(scope, self._gallery_root_scope(root, old_index)):
+                return index
+        return None
+
+    def _rescope_blocking(self, scope: str, old_index: int, relative: str):
+        """The answer to one rescope (blocking: it resolves paths)."""
+        roots = list(self._get_all_output_dirs())
+        index = self._current_root_for(scope, old_index, roots)
+        if index is None:
+            return _coded(409, "source-changed", SOURCE_GONE)
+        root = Path(roots[index])
+        path = root / relative
+        denied = _media_access_error(path, [root])
+        if denied is not None:
+            return denied
+        try:
+            mtime = path.stat().st_mtime if path.is_file() else None
+        except OSError:
+            mtime = None
+        if mtime is None:
+            return _coded(404, "file-missing", "Image file not found")
+        [image] = self._format_workspace_page([(path, mtime, root, index)])
+        return web.json_response({"success": True, "image": image})
+
+    async def rescope_image(self, request) -> web.Response:
+        data, error = await _json_body_or_error(request)
+        if error is not None:
+            return error
+        if set(data) != {"target"}:
+            return bad_request("Send exactly one target")
+        target, error = media_target(data["target"])
+        if error is not None:
+            return error
+        old_index, scope, relative = target
+        if old_index > MAX_ROOT_INDEX:
+            return bad_request("Invalid gallery root")
+        if _is_thumbnail_cache_path(relative):
+            return _forbidden("Previews are not re-scoped; use the original")
+        return await self._run_in_executor(
+            self._rescope_blocking, scope, old_index, relative
+        )