From a1aa576e0088d274db29112bf9a3cef872d96f49 Mon Sep 17 00:00:00 2001
From: wanjinhao1 <394959433@qq.com>
Date: Mon, 2 Nov 2026 09:12:50 +0800
Subject: [PATCH] Fix #3700: clean up temporary frame files when video scoring
 is cancelled

---
 pyrit/score/video_scorer.py           | 32 ++++++++++++++++++--
 tests/unit/score/test_video_scorer.py | 43 +++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/pyrit/score/video_scorer.py b/pyrit/score/video_scorer.py
index 7cac6a1f06..ec4f66f27d 100644
--- a/pyrit/score/video_scorer.py
+++ b/pyrit/score/video_scorer.py
@@ -168,8 +168,8 @@ async def _score_frames_async(
         if not Path(video_path).exists():
             raise FileNotFoundError(f"Video file not found: {video_path}")
 
-        # Extract frames from video
-        frames = await asyncio.to_thread(self._extract_frames, video_path)
+        # Extract frames from video (and clean up temporary files if cancelled)
+        frames = await self._extract_frames_async(video_path)
         if not frames:
             raise ValueError("No frames extracted from video for scoring.")
 
@@ -210,6 +210,34 @@ async def _score_frames_async(
 
         return frame_scores, len(frames)
 
+    async def _extract_frames_async(self, video_path: str) -> list[str]:
+        """
+        Extract frames off-loop and clean up temporary frame files before propagating cancellation.
+
+        Returns:
+            list[str]: Paths of the extracted temporary frame files.
+
+        Raises:
+            asyncio.CancelledError: After extraction finishes and temporary frame files are removed.
+        """
+        frame_paths: list[str] = []
+
+        def extract_frames() -> list[str]:
+            frame_paths.extend(self._extract_frames(video_path))
+            return frame_paths
+
+        try:
+            return await run_legacy_sync_async(extract_frames)
+        except asyncio.CancelledError:
+            if frame_paths:
+                await run_legacy_sync_async(self._unlink_frame_files, frame_paths)
+            raise
+
+    @staticmethod
+    def _unlink_frame_files(frame_paths: list[str]) -> None:
+        for frame_path in frame_paths:
+            Path(frame_path).unlink(missing_ok=True)
+
     def _extract_frames(self, video_path: str) -> list[str]:
         """
         Extract a specified number of image frames from a video file.
diff --git a/tests/unit/score/test_video_scorer.py b/tests/unit/score/test_video_scorer.py
index 99148e5fdf..bf5848cf87 100644
--- a/tests/unit/score/test_video_scorer.py
+++ b/tests/unit/score/test_video_scorer.py
@@ -654,6 +654,49 @@ def extract_audio(video_path: str) -> str | None:
     score_audio.assert_not_awaited()
 
 
+@pytest.mark.parametrize("cancel_twice", [False, True])
+async def test_video_frame_scorer_cancellation_during_extraction_async(
+    tmp_path: Path, video_converter_sample_video: MessagePiece, cancel_twice: bool
+) -> None:
+    """Cancelling _score_frames_async removes temporary frame files once extraction finishes."""
+    loop = asyncio.get_running_loop()
+    started = asyncio.Event()
+    release = threading.Event()
+    source_path = Path(video_converter_sample_video.converted_value)
+    source_path.touch(exist_ok=True)
+    frame_paths = [tmp_path / "cancelled_frame_0.png", tmp_path / "cancelled_frame_1.png"]
+    scorer = VideoTrueFalseScorer(image_capable_scorer=MockTrueFalseScorer())
+
+    def extract_frames(video_path: str) -> list[str]:
+        loop.call_soon_threadsafe(started.set)
+        if not release.wait(TEST_WATCHDOG_SECONDS):
+            raise TimeoutError("The test did not release frame extraction.")
+        for frame_path in frame_paths:
+            frame_path.write_bytes(b"frame")
+        return [str(frame_path) for frame_path in frame_paths]
+
+    with patch.object(type(scorer._video_helper), "_extract_frames", side_effect=extract_frames):
+        task = asyncio.create_task(scorer._video_helper._score_frames_async(message_piece=video_converter_sample_video, expectation=None))
+        try:
+            await wait_for_event_async(event=started, tasks=[task])
+            task.cancel()
+            await asyncio.sleep(0)
+            assert not task.done()
+            if cancel_twice:
+                task.cancel()
+                await asyncio.sleep(0)
+                assert not task.done()
+            release.set()
+            with pytest.raises(asyncio.CancelledError):
+                await wait_for_completion_async(future=task)
+        finally:
+            release.set()
+            await settle_tasks_async(task)
+
+    assert all(not frame_path.exists() for frame_path in frame_paths)
+    assert source_path.exists()
+
+
 @pytest.mark.skipif(not is_opencv_installed(), reason="opencv is not installed")
 async def test_video_audio_scorer_keeps_source_when_extraction_returns_it(video_converter_sample_video):
     """_score_video_audio_async must not delete the video when extraction hands back the input path.