TEST Harden media cleanup test task ownership

doc/contributing/4_running_tests.md

@@ -78,6 +78,13 @@ CI workers. These are deadlock backstops, not expected runtimes. Production dead
 are unchanged. Deadline and cancellation tests use explicit controls and keep their
 separate timing assertions.
 
+Async correctness tests can use `TEST_WATCHDOG_SECONDS` from `tests/unit/async_utils.py`
+for a finite 300-second deadlock backstop. Tests that start background tasks must
+release any test-controlled waits and join their tasks before removing mocks or
+disposing fixtures, including when a watchdog expires. When waiting for a task to
+signal readiness, also watch for that task finishing so an early exception is not
+hidden by the readiness timeout.
+
 The discovery provider-state test and Copilot cancellation-test startup use the
 shared `TEST_WATCHDOG_SECONDS` backstop in `tests/unit/async_utils.py`. Discovery
 timeout coverage expires the real cancellation path only after the request starts.

doc/contributing/5_unit_tests.md

@@ -29,6 +29,15 @@ When observing an operation's cancellation or cleanup, use `wait_for_completion_
 `unit.async_utils`. Unlike `asyncio.wait_for`, its watchdog does not send another cancellation
 request to the operation when the wait expires. Release blocked workers and drain owned tasks
 in `finally` so a failed assertion does not leave background work behind.
+Use `settle_tasks_async` for the final join: it retains the gather, cancels pending work
+if its watchdog expires, and drains cancellation cleanup before raising `TimeoutError`.
+That final drain intentionally waits for cooperative cleanup rather than abandoning a task
+at a second deadline. An outer test timeout can interrupt this drain; it is a deadlock
+backstop, not a guarantee that an uncooperative task settles. A whole-test timeout equal
+to an inner watchdog can fire before that watchdog's expiry handling runs.
+Use `wait_for_event_async` to watch readiness alongside owned tasks, so a task failure
+before readiness is reported immediately instead of waiting for an event. Once readiness
+is set, the caller observes the task's outcome separately, even if it has already failed.
 
 For deadline tests, expire a real `asyncio.Timeout` with `reschedule` once the operation reaches
 the intended pending await. Check the configured timeout arguments, cancellation, cleanup, and

tests/unit/async_utils.py

@@ -2,7 +2,9 @@
 # Licensed under the MIT license.
 
 import asyncio
-from typing import TypeVar
+from collections.abc import Sequence
+from contextlib import suppress
+from typing import Any, TypeVar
 
 T = TypeVar("T")
 
@@ -26,3 +28,47 @@ async def wait_for_completion_async(*, future: asyncio.Future[T], timeout: float
     if not done:
         raise TimeoutError("The operation under test did not complete before the test watchdog expired.")
     return future.result()
+
+
+async def settle_tasks_async(*tasks: asyncio.Task[Any], timeout: float = TEST_WATCHDOG_SECONDS) -> None:
+    """Join owned tasks, cancelling and draining stragglers before reporting watchdog expiry.
+
+    Release test-controlled gates first. Unlike an observation wait, teardown must retain
+    ownership until cancellation cleanup finishes, even after the watchdog has expired.
+    """
+    joined = asyncio.gather(*tasks, return_exceptions=True)
+    try:
+        await wait_for_completion_async(future=joined, timeout=timeout)
+    except TimeoutError:
+        joined.cancel()
+        with suppress(asyncio.CancelledError):
+            await joined
+        raise
+
+
+async def wait_for_event_async(
+    *,
+    event: asyncio.Event,
+    tasks: Sequence[asyncio.Future[Any]],
+    timeout: float = TEST_WATCHDOG_SECONDS,
+) -> None:
+    """Report task completion or failure only when readiness has not been signalled.
+
+    Once readiness is set, the caller owns observing the operation's outcome, including
+    failures that complete before this readiness observer resumes.
+    """
+    ready = asyncio.create_task(event.wait())
+    try:
+        done, _ = await asyncio.wait({ready, *tasks}, timeout=timeout, return_when=asyncio.FIRST_COMPLETED)
+        if event.is_set():
+            return
+        if not done:
+            raise TimeoutError("The operation did not signal readiness before the test watchdog expired.")
+        for task in tasks:
+            if task.done():
+                task.result()
+        raise AssertionError("The operation completed without signalling readiness.")
+    finally:
+        ready.cancel()
+        with suppress(asyncio.CancelledError):
+            await ready

tests/unit/backend/test_attack_service.py

@@ -78,6 +78,7 @@
 )
 from pyrit.models.conversation_stats import ConversationStats
 from pyrit.prompt_normalizer import PromptNormalizer
+from unit.async_utils import TEST_WATCHDOG_SECONDS, settle_tasks_async, wait_for_event_async
 from unit.backend.mocks import _make_matching_target_mock, make_attack_result, make_mock_memory
 from unit.mocks import MockPromptTarget
 
@@ -1566,7 +1567,7 @@ async def test_add_message_propagates_service_errors(
             await attack_service.add_message_async(attack_result_id="attack", request=request)
         reader.assert_not_called()
 
-    @pytest.mark.timeout(10)
+    @pytest.mark.timeout(TEST_WATCHDOG_SECONDS)
     @pytest.mark.parametrize("read_method", ["get_attack_async", "get_conversation_messages_async"])
     @pytest.mark.parametrize("incoming_send", [False, True])
     async def test_response_assembly_keeps_conversation_reserved_async(
@@ -1602,7 +1603,7 @@ async def pause_after_read_async(**kwargs: Any) -> AttackSummary | ConversationM
                 service.add_message_async(attack_result_id=ar.attack_result_id, request=request)
             )
             try:
-                await started.wait()
+                await wait_for_event_async(event=started, tasks=[active])
                 with pytest.raises(ManualSendConflictError):
                     await peer.add_message_async(
                         attack_result_id=ar.attack_result_id,
@@ -1615,7 +1616,8 @@ async def pause_after_read_async(**kwargs: Any) -> AttackSummary | ConversationM
                     )
             finally:
                 release.set()
-                result = await active
+                await settle_tasks_async(active)
+            result = active.result()
 
         assert result.attack.message_count == 2
         assert [message.message_pieces[0].original_value for message in result.messages.messages] == [
@@ -1634,7 +1636,7 @@ async def pause_after_read_async(**kwargs: Any) -> AttackSummary | ConversationM
             ),
         )
 
-    @pytest.mark.timeout(10)
+    @pytest.mark.timeout(TEST_WATCHDOG_SECONDS)
     @pytest.mark.parametrize("read_method", ["get_attack_async", "get_conversation_messages_async"])
     @pytest.mark.parametrize("cancel", [False, True])
     async def test_response_failure_releases_conversation_async(
@@ -1659,14 +1661,15 @@ async def fail_read_async(**_: Any) -> None:
         ):
             active = asyncio.create_task(attack_service.add_message_async(attack_result_id="attack", request=request))
             try:
-                await started.wait()
+                await wait_for_event_async(event=started, tasks=[active])
                 assert sender._scheduler._conversations == {"conversation"}
                 if cancel:
                     active.cancel()
             finally:
                 release.set()
-                with pytest.raises(asyncio.CancelledError if cancel else RuntimeError):
-                    await active
+                await settle_tasks_async(active)
+            with pytest.raises(asyncio.CancelledError if cancel else RuntimeError):
+                active.result()
             assert not sender._scheduler._conversations
             await sender.add_message_async(attack_result_id="attack", request=request)
 

tests/unit/backend/test_conversation_editor.py

@@ -50,6 +50,7 @@
 from pyrit.models.target.target_capabilities import TargetCapabilities
 from pyrit.prompt_target import OpenAIResponseTarget, TargetConfiguration
 from pyrit.score.scorer_prompt_validator import ScorerPromptValidator
+from unit.async_utils import TEST_WATCHDOG_SECONDS, settle_tasks_async, wait_for_completion_async, wait_for_event_async
 from unit.backend.mocks import _settle_send_async, message_send_lifecycle_async
 from unit.mocks import MockPromptTarget, get_mock_target_identifier, run_memory_session_async
 
@@ -1275,6 +1276,49 @@ async def test_media_rollback_deletes_only_staged_files_async(self, sqlite_insta
         assert not list(Path(sqlite_instance.results_path).rglob("*.txt"))
         assert (await sqlite_instance.get_attack_results_async()) == []
 
+    @pytest.mark.timeout(TEST_WATCHDOG_SECONDS)
+    async def test_media_cleanup_watchdog_settles_save_before_teardown_async(
+        self, *, sqlite_instance: SQLiteMemory
+    ) -> None:
+        request = draft()
+        request.messages[0].pieces = [
+            ConversationPieceRequest(data_type="binary_path", original_value="data:text/plain;base64,bWVkaWE=")
+        ]
+        storage = sqlite_instance.results_storage_io
+        write, delete = storage.write_file_async, storage.delete_file_async
+        cleanup_started, cleanup_finished = asyncio.Event(), asyncio.Event()
+
+        async def failed_write_async(*args: Any, **kwargs: Any) -> None:
+            await write(*args, **kwargs)
+            raise OSError("write interrupted")
+
+        async def blocked_delete_async(path: str) -> None:
+            try:
+                cleanup_started.set()
+                await asyncio.Event().wait()
+            finally:
+                await delete(path)
+                cleanup_finished.set()
+
+        with (
+            patch.object(storage, "write_file_async", side_effect=failed_write_async),
+            patch.object(storage, "delete_file_async", side_effect=blocked_delete_async) as delete_file,
+        ):
+            task = asyncio.create_task(AttackService().save_conversation_async(request=request))
+            try:
+                await wait_for_event_async(event=cleanup_started, tasks=[task])
+                with pytest.raises(TimeoutError, match="test watchdog expired"):
+                    await settle_tasks_async(task, timeout=0)
+                assert task.cancelled()
+                assert cleanup_finished.is_set()
+                delete_file.assert_awaited_once()
+            finally:
+                if not task.done():
+                    task.cancel()
+                await settle_tasks_async(task)
+        assert not list(Path(sqlite_instance.results_path).rglob("*.txt"))
+        assert (await sqlite_instance.get_attack_results_async()) == []
+
     @pytest.mark.parametrize("rollback_fails", [False, True])
     async def test_media_rollback_preserves_save_error_when_target_close_fails_async(
         self, sqlite_instance: SQLiteMemory, caplog: pytest.LogCaptureFixture, rollback_fails: bool
@@ -1615,18 +1659,21 @@ async def test_transaction_failure_rolls_back_attack_and_conversation_async(
         assert (await sqlite_instance.get_conversation_metadata_async(conversation_id=str(request.save_id))) is None
         assert (await sqlite_instance.get_attack_results_async()) == []
 
+    @pytest.mark.timeout(TEST_WATCHDOG_SECONDS)
     async def test_concurrent_retries_leave_one_conversation_async(self, sqlite_instance: SQLiteMemory) -> None:
         service = AttackService()
         request = draft()
-        first, second = await asyncio.gather(
-            service.save_conversation_async(request=request),
-            service.save_conversation_async(request=request),
-        )
+        tasks = [asyncio.create_task(service.save_conversation_async(request=request)) for _ in range(2)]
+        try:
+            first, second = await wait_for_completion_async(future=asyncio.gather(*tasks))
+        finally:
+            await settle_tasks_async(*tasks)
         assert first.messages == second.messages
         assert len(await sqlite_instance.get_attack_results_async()) == 1
         assert len(await sqlite_instance.get_message_pieces_async(conversation_id=str(request.save_id))) == 1
 
     @pytest.mark.parametrize("cancel", [False, True])
+    @pytest.mark.timeout(TEST_WATCHDOG_SECONDS)
     async def test_failed_media_write_cleans_partial_file_async(
         self, *, sqlite_instance: SQLiteMemory, cancel: bool
     ) -> None:
@@ -1652,11 +1699,18 @@ async def interrupted_write_async(*args: Any, **kwargs: Any) -> None:
 
         with patch.object(storage, "write_file_async", side_effect=interrupted_write_async):
             task = asyncio.create_task(AttackService().save_conversation_async(request=request))
-            await asyncio.wait_for(started.wait(), timeout=5)
-            if cancel:
-                task.cancel()
+            try:
+                await wait_for_event_async(event=started, tasks=[task])
+                if cancel:
+                    task.cancel()
+                    finish.set()
+                with pytest.raises(asyncio.CancelledError if cancel else OSError):
+                    await wait_for_completion_async(future=task)
+            finally:
+                if not task.done() and not task.cancelling():
+                    task.cancel()
+                # The shielded writer must finish before the save task can clean up.
                 finish.set()
-            with pytest.raises(asyncio.CancelledError if cancel else OSError):
-                await task
+                await settle_tasks_async(task)
         assert not list(Path(sqlite_instance.results_path).rglob("*.txt"))
         assert (await sqlite_instance.get_attack_results_async()) == []

tests/unit/backend/test_media_persistence.py

@@ -3,13 +3,22 @@
 
 """Tests for shared backend media persistence."""
 
+import asyncio
+import uuid
 from pathlib import Path
+from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 from urllib.parse import urlencode
 
 import pytest
 
-from pyrit.backend.models.attacks import MessagePieceRequest
+from pyrit.backend.models.attacks import (
+    ConversationMessageRequest,
+    ConversationPieceRequest,
+    MessagePieceRequest,
+    SaveConversationRequest,
+)
+from pyrit.backend.services.attack_service import AttackService
 from pyrit.backend.services.media_persistence import (
     MediaAccessDeniedError,
     MediaOrigin,
@@ -20,6 +29,7 @@
     validate_media_path,
 )
 from pyrit.memory import AzureBlobStorageIO, CentralMemory, DiskStorageIO, SQLiteMemory, data_serializer_factory
+from unit.async_utils import TEST_WATCHDOG_SECONDS, settle_tasks_async, wait_for_completion_async, wait_for_event_async
 
 
 def _serializer(*, value: str = "/saved/media.bin") -> MagicMock:
@@ -480,3 +490,76 @@ async def test_persistence_failure_returns_no_partial_result() -> None:
             data_type="binary_path",
             serializer_factory=MagicMock(return_value=serializer),
         )
+
+
+@pytest.mark.usefixtures("patch_central_database")
+@pytest.mark.timeout(TEST_WATCHDOG_SECONDS)
+async def test_cancelled_save_finishes_writer_before_deleting_media_async(sqlite_instance: SQLiteMemory) -> None:
+    request = SaveConversationRequest(
+        save_id=uuid.uuid4(),
+        destination="new_attack",
+        objective="Save media",
+        operator="owner",
+        messages=[
+            ConversationMessageRequest(
+                role="user",
+                pieces=[
+                    ConversationPieceRequest(data_type="binary_path", original_value="data:text/plain;base64,bWVkaWE=")
+                ],
+            )
+        ],
+    )
+    started, release_writer, writer_finished, cancellation_observed = (
+        asyncio.Event(),
+        asyncio.Event(),
+        asyncio.Event(),
+        asyncio.Event(),
+    )
+    storage = sqlite_instance.results_storage_io
+    write, delete, shield = storage.write_file_async, storage.delete_file_async, asyncio.shield
+    writers: list[asyncio.Task[Any]] = []
+
+    async def blocked_write_async(*args: Any, **kwargs: Any) -> None:
+        writer = asyncio.current_task()
+        assert writer is not None
+        writers.append(writer)
+        started.set()
+        await release_writer.wait()
+        await write(*args, **kwargs)
+        writer_finished.set()
+
+    async def checked_delete_async(path: str) -> None:
+        assert writer_finished.is_set()
+        await delete(path)
+
+    def observe_shield(future: asyncio.Future[Any]) -> asyncio.Future[Any]:
+        protected = shield(future)
+        protected.add_done_callback(lambda done: cancellation_observed.set() if done.cancelled() else None)
+        return protected
+
+    with (
+        patch.object(storage, "write_file_async", side_effect=blocked_write_async),
+        patch.object(storage, "delete_file_async", side_effect=checked_delete_async) as delete_file,
+        patch.object(asyncio, "shield", side_effect=observe_shield),
+    ):
+        task = asyncio.create_task(AttackService().save_conversation_async(request=request))
+        try:
+            await wait_for_event_async(event=started, tasks=[task])
+            task.cancel()
+            # Observe delivery to the shield, not merely issuance of a cancellation request.
+            await wait_for_event_async(event=cancellation_observed, tasks=[task])
+            assert not task.done()
+            assert not writers[0].done()
+            assert not writer_finished.is_set()
+            delete_file.assert_not_awaited()
+            release_writer.set()
+            with pytest.raises(asyncio.CancelledError):
+                await wait_for_completion_async(future=task)
+            assert writer_finished.is_set()
+            delete_file.assert_awaited_once()
+        finally:
+            release_writer.set()
+            if not task.done() and not task.cancelling():
+                task.cancel()
+            await settle_tasks_async(task, *writers)
+    assert not list(Path(sqlite_instance.results_path).rglob("*.txt"))