diff --git a/py/api/__init__.py b/py/api/__init__.py index 3eaed5e..4e60272 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -294,6 +294,8 @@ class PromptManagerAPI( # created on first use with their restore subscription. self._captures = None self._captures_unsubscribe = None + # Set by close_annotation_captures (shutdown): no store is recreated. + self._captures_closed = False 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) diff --git a/py/api/annotations.py b/py/api/annotations.py index caadb24..ad602b4 100644 --- a/py/api/annotations.py +++ b/py/api/annotations.py @@ -264,8 +264,13 @@ class AnnotationRoutesMixin: a server that never opens the Tag dialog) owns no temp directory. A restore of the database dooms every capture (``invalidate_all``). Thread-safe: Needs review recording calls it off the event loop. + Raises ``CaptureError`` after ``close_annotation_captures``. """ with _CAPTURE_STORE_LOCK: + if getattr(self, "_captures_closed", False): + # A recording still in flight at shutdown must not create a + # new temp directory after the cleanup removed the old one. + raise CaptureError("The capture store is closed") if self._captures is None: self._captures = CaptureStore() if self._captures_unsubscribe is None: @@ -280,11 +285,17 @@ class AnnotationRoutesMixin: 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 + """Remove every capture file and the owned directory. Idempotent. + + Taken under the store lock, and final: ``_capture_store`` refuses to + create another store afterwards. + """ + with _CAPTURE_STORE_LOCK: + self._captures_closed = True + unsubscribe, self._captures_unsubscribe = self._captures_unsubscribe, None + store, self._captures = self._captures, None if unsubscribe is not None: unsubscribe() - store, self._captures = self._captures, None if store is not None: store.close() diff --git a/py/api/review.py b/py/api/review.py index 5980dd6..2d3e425 100644 --- a/py/api/review.py +++ b/py/api/review.py @@ -149,6 +149,10 @@ def format_review_target(spelled_path, locator, real_roots, *, pathmod=os.path): return None +class _StaleCursor(ValueError): + """A well-formed cursor this process cannot verify (e.g. after a restart).""" + + def _encode_cursor(epoch: str, revision: int, key: tuple) -> str: payload = json.dumps( {"v": 1, "e": epoch, "r": revision, "k": list(key)}, @@ -160,7 +164,12 @@ def _encode_cursor(epoch: str, revision: int, key: tuple) -> str: def _decode_cursor(token: str) -> tuple: - """``(epoch, revision, (created_at, id))`` of a cursor; ValueError if forged.""" + """``(epoch, revision, (created_at, id))`` of a cursor. + + ``ValueError`` when malformed; :class:`_StaleCursor` when well-formed but + signed by another secret (the per-process secret changes on every start, + so an open page's cursor is stale after a restart, not forged). + """ try: if len(token) > MAX_CURSOR: raise ValueError() @@ -168,9 +177,11 @@ def _decode_cursor(token: str) -> tuple: token + "=" * (-len(token) % 4), altchars=b"-_", validate=True ) signature, payload = raw[:32], raw[32:] + if len(signature) != 32 or not payload: + raise ValueError() expected = hmac.new(_CURSOR_SECRET, payload, hashlib.sha256).digest() if not hmac.compare_digest(signature, expected): - raise ValueError() + raise _StaleCursor("This page of the queue is out of date") data = json.loads(payload) key = data["k"] if ( @@ -183,6 +194,8 @@ def _decode_cursor(token: str) -> tuple: ): raise ValueError() return data["e"], data["r"], tuple(key) + except _StaleCursor: + raise except (ValueError, KeyError, TypeError, UnicodeError, binascii.Error): raise ValueError("Invalid review cursor; reload the queue") from None @@ -280,6 +293,8 @@ class ReviewRoutesMixin: if cursor: try: bound = _decode_cursor(cursor) + except _StaleCursor: + return await self._stale_cursor() except ValueError as exc: return bad_request(str(exc)) try: @@ -311,6 +326,17 @@ class ReviewRoutesMixin: headers=_NO_STORE, ) + async def _stale_cursor(self) -> web.Response: + """409 ``review-conflict`` with the current total: reload the queue.""" + try: + page = await self._run_in_executor(self.db.review_queue_page, limit=1) + except Exception: + self.logger.exception("Review queue read failed") + return _failed() + return _coded( + "review-conflict", total=page["total"], library_epoch=page["epoch"] + ) + def _review_page(self, limit: int, after): """Worker thread: one queue page and the public views of its items.""" page = self.db.review_queue_page(limit=limit, after=after) diff --git a/tests/test_api_review.py b/tests/test_api_review.py index 44b8877..1232041 100644 --- a/tests/test_api_review.py +++ b/tests/test_api_review.py @@ -231,8 +231,19 @@ class TestQueueReads(ReviewRoutesTestCase): self.assertTrue(cursor) # Tampered: one character changed, or a cursor of another shape. broken = ("A" if cursor[5] != "A" else "B").join((cursor[:5], cursor[6:])) + for token in ("!!!", "QUJD", cursor[:40]): + status, data = await self._json("GET", f"{ITEMS_URL}?cursor={token}") + self.assertEqual(status, 400, (token, data)) + # Well-formed but not this process's signature (tampered, or opened + # before a ComfyUI restart): a recoverable conflict (M3). status, data = await self._json("GET", f"{ITEMS_URL}?cursor={broken}") - self.assertEqual(status, 400, data) + self.assertEqual( + (status, data["code"], data["total"]), (409, "review-conflict", 3) + ) + with patch.object(review_api, "_CURSOR_SECRET", b"\x01" * 32): + status, data = await self._json("GET", f"{ITEMS_URL}?cursor={cursor}") + self.assertEqual((status, data["code"]), (409, "review-conflict"), data) + self.assertEqual(data["library_epoch"], first["library_epoch"]) # A decision moves the review index revision: the old cursor is stale. item, detail, observed = await self._open() status, data = await self._decide( diff --git a/tests/test_review.py b/tests/test_review.py index cf1d277..16ebf6a 100644 --- a/tests/test_review.py +++ b/tests/test_review.py @@ -33,7 +33,7 @@ from py.api import PromptManagerAPI from py.config import GalleryConfig from utils.hashing import generate_prompt_hash from utils.image_monitor import ImageGenerationHandler -from utils.media_capture import CaptureCapacity, CaptureStore +from utils.media_capture import CaptureCapacity, CaptureError, CaptureStore from utils.prompt_tracker import PromptTracker try: # guarded so a missing module fails assertions, not test loading @@ -892,3 +892,26 @@ class TestRecordedSpelling(ReviewTestCase): self.assertEqual(self._review_revision(), revision + 1) with self.assertRaises(ValueError): db.refresh_review_spelling(locator, os.path.join(root, "other.png")) + + +class TestCaptureStoreShutdown(ReviewTestCase): + async def test_recording_after_close_creates_no_store(self): + """A monitor recording still in flight at shutdown leaks nothing (M2).""" + self._require_module() + self._roots() + path = self.output_dir / "two.png" + write_png(path, prompt=ambiguous_graph("cat", "dog")) + self.api.close_annotation_captures() + self.assertIsNone(self.api._captures) + with patch("utils.media_capture.tempfile.mkdtemp") as mkdtemp: + with self.assertRaises(CaptureError): + self.api.record_review_file(str(path)) + # The monitor logs the failure and keeps its link. + handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db)) + handler.review_recorder = self.api.record_review_file + handler._record_review(str(path)) + mkdtemp.assert_not_called() + self.assertIsNone(self.api._captures) + self.assertIsNone(self.api._captures_unsubscribe) + self.assertEqual(self._items(), []) + self.api.close_annotation_captures() # idempotent