diff --git a/py/api/annotations.py b/py/api/annotations.py index ad602b4..8c799ff 100644 --- a/py/api/annotations.py +++ b/py/api/annotations.py @@ -284,14 +284,16 @@ class AnnotationRoutesMixin: if store is not None: store.invalidate_all() - def close_annotation_captures(self) -> None: + def close_annotation_captures(self, *, final: bool = False) -> 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. + Taken under the store lock. ``final`` (the shutdown cleanup) also makes + ``_capture_store`` refuse to create another store, so a recording + still in flight cannot leak a new temp directory after it. """ with _CAPTURE_STORE_LOCK: - self._captures_closed = True + if final: + self._captures_closed = True unsubscribe, self._captures_unsubscribe = self._captures_unsubscribe, None store, self._captures = self._captures, None if unsubscribe is not None: @@ -300,7 +302,7 @@ class AnnotationRoutesMixin: store.close() async def cleanup_annotation_captures(self, app: web.Application) -> None: - self.close_annotation_captures() + self.close_annotation_captures(final=True) # -- shared request steps -------------------------------------------------- @@ -346,7 +348,10 @@ class AnnotationRoutesMixin: root, error = self._scoped_root(index, scope) if error is not None: return error - store = self._capture_store() + try: + store = self._capture_store() + except CaptureError as exc: # the server is shutting down + return _capture_error(exc) try: epoch = await self._run_in_executor(self.db.get_execution_epoch) capture = await _settled( diff --git a/tests/test_api_annotations.py b/tests/test_api_annotations.py index e59a65f..6476ea7 100644 --- a/tests/test_api_annotations.py +++ b/tests/test_api_annotations.py @@ -1045,7 +1045,15 @@ class TestFailures(AnnotationAPITestCase): 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. + # The shutdown cleanup is final: a later observation (or a monitor + # recording still in flight) creates no new store or directory. + status, body = await self.observe({"target": self.target("sub/a.png")}) + self.assertEqual((status, body["code"]), (500, "capture-error"), body) + self.assertIsNone(self.api._captures) + self.assertNotIn(identity, observers._restore_observers) + # An explicit, non-final close still lets a later observation start + # a fresh store (the tests' own store swaps rely on it). + self.api._captures_closed = False await self.observe_ok() self.assertNotEqual(self.api._captures.directory, directory) diff --git a/tests/test_review.py b/tests/test_review.py index 16ebf6a..bcafb50 100644 --- a/tests/test_review.py +++ b/tests/test_review.py @@ -901,7 +901,7 @@ class TestCaptureStoreShutdown(ReviewTestCase): self._roots() path = self.output_dir / "two.png" write_png(path, prompt=ambiguous_graph("cat", "dog")) - self.api.close_annotation_captures() + await self.api.cleanup_annotation_captures(None) # the shutdown hook self.assertIsNone(self.api._captures) with patch("utils.media_capture.tempfile.mkdtemp") as mkdtemp: with self.assertRaises(CaptureError):