BUG: Cancelling video frame extraction leaves unreferenced temporary images #3700

import asyncio
import functools
import sys
import tempfile
import threading
import types
from pathlib import Path
from unittest.mock import MagicMock, patch

from pyrit.models import MessagePiece, ScoringExpectation
from pyrit.score.scorer import Scorer
from pyrit.score.video_scorer import VideoHelper

async def main_async():
    loop = asyncio.get_running_loop()
    started, finished = asyncio.Event(), asyncio.Event()
    proceed = threading.Event()
    written = []
    cv2 = types.ModuleType("cv2")
    cv2.CAP_PROP_FRAME_COUNT = 1
    cv2.CAP_PROP_POS_FRAMES = 2

    class Capture:
        def __init__(self, path):
            pass
        def get(self, property_id):
            return 2
        def set(self, property_id, value):
            pass
        def read(self):
            return True, object()
        def release(self):
            loop.call_soon_threadsafe(finished.set)

    def write(path, frame):
        Path(path).write_bytes(b"synthetic-frame")
        written.append(Path(path))
        if len(written) == 1:
            loop.call_soon_threadsafe(started.set)
            if not proceed.wait(30):
                raise TimeoutError("Worker rendezvous timed out")
        return True

    cv2.VideoCapture = Capture
    cv2.imencode = lambda extension, frame: (True, b"synthetic")
    cv2.imwrite = write
    scorer = MagicMock(spec=Scorer)
    scorer._select_expectation.return_value = ScoringExpectation()
    helper = VideoHelper(image_capable_scorer=scorer, num_sampled_frames=2)
    original_tempfile = tempfile.NamedTemporaryFile
    with tempfile.TemporaryDirectory() as directory:
        root = Path(directory)
        video = root / "synthetic.mp4"
        video.touch()
        piece = MessagePiece(
            role="assistant", original_value=str(video), original_value_data_type="video_path",
        )
        with patch.dict(sys.modules, {"cv2": cv2}), patch(
            "pyrit.score.video_scorer.tempfile.NamedTemporaryFile",
            functools.partial(original_tempfile, dir=root),
        ):
            task = asyncio.create_task(helper._score_frames_async(
                message_piece=piece, expectation=None,
            ))
            try:
                await asyncio.wait_for(started.wait(), 30)
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    print("cancelled_before_worker_finished:", not finished.is_set())
                print("files_at_cancellation:", len(list(root.glob("*_frame_*.png"))))
            finally:
                proceed.set()
                await asyncio.wait_for(finished.wait(), 30)
                await loop.shutdown_default_executor()
            outputs = list(root.glob("*_frame_*.png"))
            print("files_after_worker_finished:", len(outputs))
            print("frame_scorer_calls:", scorer._score_batch_nested_async.call_count)

asyncio.run(main_async())