fix(workspace): keep open pages valid across a restart and recover a changed source (#430)

The gallery root scope and the Images listing cursor were keyed with a
per-process random key, so every ComfyUI restart made an open tab's
scopes answer 409 source-changed (blank Tag image preview, a Generate
that did nothing) and its cursors fail Load more.

- utils/workspace_keys.py: one random secret in workspace-scope.key
  beside the database (owner-only, never config.json or the database),
  two derived keys; per-process fallback with a warning. Move library
  copies it.
- Cursor v2 binds the listing scope in its MAC instead of sending an
  unkeyed hash of the root paths; an unverifiable cursor answers 400
  cursor-stale and Images offers Refresh instead of a dead Retry.
- POST /prompt_manager/images/rescope maps a stale target to its
  current root only when the held scope was issued for it.
- Images finds a renumbered file again from the inspector, the viewer
  and the Tag image dialog (the dialog stays open, draft kept); nothing
  else is kept by an old root_index:path key.
- Tag image dialog: a state that blocks Save shows its reason above the
  actions with Refresh, disabled actions look disabled, and an empty
  preview says "Preview unavailable" (D147).

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
@@ -340,6 +341,7 @@ class PromptManagerAPI(
     LoraIntegrationMixin,
     LoraImportRoutesMixin,
     ImageCategoryRoutesMixin,
+    SourceRescopeMixin,
 ):
     """REST API handler for PromptManager operations and web interface.
 
@@ -706,6 +708,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/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
+        )