FIX Harden prompt target cancellation cleanup

pyrit/prompt_target/hugging_face/hugging_face_chat_target.py

@@ -172,6 +172,7 @@ def __init__(
             raise RuntimeError("CUDA requested but not available.")
 
         self.load_model_and_tokenizer_task = asyncio.create_task(self.load_model_and_tokenizer_async())
+        self._model_load_lock = asyncio.Lock()
 
     def _build_identifier(self) -> ComponentIdentifier:
         """
@@ -347,7 +348,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me
         Raises:
             EmptyResponseException: If the model generates an empty response.
         """
-        await self.load_model_and_tokenizer_task
+        await self._wait_for_model_and_tokenizer_async()
 
         request = normalized_conversation[-1].message_pieces[0]
 
@@ -405,6 +406,14 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me
             logger.error(f"Error occurred during inference: {e}")
             raise
 
+    async def _wait_for_model_and_tokenizer_async(self) -> None:
+        """Wait for shared model loading without allowing a send cancellation to cancel it."""
+        async with self._model_load_lock:
+            if self.load_model_and_tokenizer_task.cancelled():
+                self.load_model_and_tokenizer_task = asyncio.create_task(self.load_model_and_tokenizer_async())
+            load_task = self.load_model_and_tokenizer_task
+        await asyncio.shield(load_task)
+
     def _build_chat_messages(self, *, normalized_conversation: list[Message]) -> list[dict[str, str]]:
         """
         Build a list of chat message dicts from the full normalized conversation.

pyrit/prompt_target/openai/openai_realtime_target.py

@@ -419,10 +419,14 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me
             connection = await self._connect_async(conversation_id=conversation_id)
             self._existing_conversation[conversation_id] = connection
 
-            # Only send config when creating a new connection
-            await self.send_config_async(conversation_id=conversation_id, conversation=normalized_conversation)
-            # Give the server a moment to process the session update
-            await asyncio.sleep(0.5)
+            try:
+                # Only send config when creating a new connection
+                await self.send_config_async(conversation_id=conversation_id, conversation=normalized_conversation)
+                # Give the server a moment to process the session update
+                await asyncio.sleep(0.5)
+            except BaseException:
+                await self.reset_conversation_async(conversation_id=conversation_id)
+                raise
 
         response_type = request.converted_value_data_type
 
@@ -491,18 +495,27 @@ async def reset_conversation_async(self, *, conversation_id: str) -> None:
 
         Args:
             conversation_id (str): The conversation ID to disconnect from.
+
+        Raises:
+            asyncio.CancelledError: If cleanup is cancelled, after the connection finishes closing.
         """
         connection = self._existing_conversation.pop(conversation_id, None)
         if not connection:
             return
 
+        close_task = asyncio.ensure_future(connection.close())
         try:
-            await connection.close()
-        except Exception as error:  # noqa: BLE001 - cleanup must not replace the attack outcome
-            logger.warning(f"Error closing connection for {conversation_id}: {error}")
-            return
-
-        logger.info(f"Disconnected from {self._endpoint} with conversation ID: {conversation_id}")
+            await asyncio.shield(close_task)
+        except asyncio.CancelledError as cancellation_error:
+            try:
+                await close_task
+            except BaseException as close_error:
+                raise cancellation_error from close_error
+            raise
+        except Exception as e:
+            logger.warning(f"Error closing connection for {conversation_id}: {e}")
+        else:
+            logger.info(f"Disconnected from {self._endpoint} with conversation ID: {conversation_id}")
 
     async def cleanup_conversation_async(self, conversation_id: str) -> None:
         """
@@ -815,24 +828,22 @@ async def send_text_async(
         connection = self._get_connection(conversation_id=conversation_id)
 
         # Start listening for responses
-        receive_tasks = asyncio.create_task(self.receive_events_async(conversation_id=conversation_id))
+        receive_task = asyncio.create_task(self.receive_events_async(conversation_id=conversation_id))
 
         logger.info(f"Sending text message: {text}")
 
-        # Send conversation item
-        await connection.conversation.item.create(
-            item={
-                "type": "message",
-                "role": "user",
-                "content": [{"type": "input_text", "text": text}],
-            }
-        )
-
-        # Request response from model
-        await self.send_response_create_async(conversation_id=conversation_id)
-
-        # Wait for response - receive_events has its own soft-finish logic
-        result = await receive_tasks
+        try:
+            await connection.conversation.item.create(
+                item={
+                    "type": "message",
+                    "role": "user",
+                    "content": [{"type": "input_text", "text": text}],
+                }
+            )
+            await self.send_response_create_async(conversation_id=conversation_id)
+            result = await receive_task
+        finally:
+            await self._cancel_receive_task_async(receive_task=receive_task)
 
         if not result.audio_bytes:
             raise RuntimeError("No audio received from the server.")
@@ -865,7 +876,7 @@ async def send_audio_async(
 
         audio_content, num_channels, sample_width, frame_rate = await asyncio.to_thread(self._read_wav_file, filename)
 
-        receive_tasks = asyncio.create_task(self.receive_events_async(conversation_id=conversation_id))
+        receive_task = asyncio.create_task(self.receive_events_async(conversation_id=conversation_id))
 
         try:
             audio_base64 = base64.b64encode(audio_content).decode("utf-8")
@@ -879,23 +890,28 @@ async def send_audio_async(
                     "content": [{"type": "input_audio", "audio": audio_base64}],
                 }
             )
-
+            logger.debug("Sending response.create")
+            await self.send_response_create_async(conversation_id=conversation_id)
+            logger.debug("Waiting for response events...")
+            result = await receive_task
         except Exception as e:
             logger.error(f"Error sending audio: {e}")
             raise
+        finally:
+            await self._cancel_receive_task_async(receive_task=receive_task)
 
-        logger.debug("Sending response.create")
-        await self.send_response_create_async(conversation_id=conversation_id)
-
-        logger.debug("Waiting for response events...")
-        # Wait for response - receive_events has its own soft-finish logic
-        result = await receive_tasks
         if not result.audio_bytes:
             raise RuntimeError("No audio received from the server.")
 
         output_audio_path = await self.save_audio_async(result.audio_bytes, num_channels, sample_width, frame_rate)
         return output_audio_path, result
 
+    async def _cancel_receive_task_async(self, *, receive_task: asyncio.Task[RealtimeTargetResult]) -> None:
+        """Cancel and retrieve a Realtime receive task."""
+        if not receive_task.done():
+            receive_task.cancel()
+        await asyncio.gather(receive_task, return_exceptions=True)
+
     async def _construct_message_from_response_async(self, response: Any, request: Any) -> Message:
         """
         Not used in RealtimeTarget - message construction handled by receive_events.

tests/unit/prompt_target/target/test_huggingface_chat_target.py

@@ -27,6 +27,31 @@ def is_torch_installed():
         return False
 
 
+@pytest.mark.skipif(not is_torch_installed(), reason="torch is not installed")
+async def test_send_cancellation_does_not_cancel_shared_model_load(patch_central_database):
+    target = HuggingFaceChatTarget(model_id="test_model", use_cuda=False)
+    load_started = asyncio.Event()
+    load_release = asyncio.Event()
+
+    async def load_model_async() -> None:
+        load_started.set()
+        await load_release.wait()
+
+    shared_load = asyncio.ensure_future(load_model_async())
+    target.load_model_and_tokenizer_task = shared_load
+    wait_task = asyncio.ensure_future(target._wait_for_model_and_tokenizer_async())
+    await load_started.wait()
+
+    wait_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await wait_task
+
+    assert not shared_load.cancelled()
+    load_release.set()
+    await shared_load
+    await target._wait_for_model_and_tokenizer_async()
+
+
 # Fixture to mock get_required_value
 @pytest.fixture(autouse=True)
 def mock_get_required_value(request):

tests/unit/prompt_target/target/test_realtime_target.py

@@ -86,6 +86,71 @@ async def test_send_prompt_async(target):
     await target.cleanup_target_async()
 
 
+async def test_cancellation_during_session_config_discards_connection(target):
+    connection = AsyncMock()
+    target._connect_async = AsyncMock(return_value=connection)
+    config_started = asyncio.Event()
+
+    async def wait_in_config_async(*, conversation_id: str, conversation: list[Message]) -> None:
+        config_started.set()
+        await asyncio.Event().wait()
+
+    target.send_config_async = AsyncMock(side_effect=wait_in_config_async)
+    message = Message.from_prompt(prompt="Hello", role="user")
+    message.get_piece().conversation_id = "cancelled-config"
+
+    send_task = asyncio.create_task(target.send_prompt_async(message=message))
+    await config_started.wait()
+    send_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await send_task
+
+    connection.close.assert_awaited_once_with()
+    assert "cancelled-config" not in target._existing_conversation
+
+
+async def test_response_create_failure_cancels_receive_task(target):
+    connection = AsyncMock()
+    target._existing_conversation["response-failure"] = connection
+    receive_started = asyncio.Event()
+    receive_cancelled = asyncio.Event()
+
+    async def receive_events_async(*, conversation_id: str) -> RealtimeTargetResult:
+        receive_started.set()
+        try:
+            await asyncio.Event().wait()
+        finally:
+            receive_cancelled.set()
+        raise AssertionError("unreachable")
+
+    async def fail_response_create_async(*, conversation_id: str) -> None:
+        await receive_started.wait()
+        raise RuntimeError("response create failed")
+
+    target.receive_events_async = receive_events_async
+    target.send_response_create_async = AsyncMock(side_effect=fail_response_create_async)
+
+    with pytest.raises(RuntimeError, match="response create failed"):
+        await target.send_text_async(text="Hello", conversation_id="response-failure")
+
+    assert receive_started.is_set()
+    assert receive_cancelled.is_set()
+
+
+async def test_cancel_receive_task_async_gathers_completed_failure(target):
+    async def fail_receive_async() -> RealtimeTargetResult:
+        raise RuntimeError("receive failed")
+
+    receive_task = asyncio.create_task(fail_receive_async())
+    with pytest.raises(RuntimeError, match="receive failed"):
+        await receive_task
+
+    with patch.object(asyncio, "gather", new_callable=AsyncMock) as gather:
+        await target._cancel_receive_task_async(receive_task=receive_task)
+
+    gather.assert_awaited_once_with(receive_task, return_exceptions=True)
+
+
 async def test_send_prompt_async_propagates_interrupted_to_metadata(target):
     """When a turn result carries interrupted=True, both response pieces' metadata must reflect it."""
     target._connect_async = AsyncMock(return_value=AsyncMock())
@@ -1433,6 +1498,34 @@ async def test_reset_conversation_async_swallows_close_error(target):
     assert "conv" not in target._existing_conversation
 
 
+async def test_reset_conversation_async_finishes_close_before_propagating_cancellation(target):
+    close_started = asyncio.Event()
+    close_release = asyncio.Event()
+    close_finished = asyncio.Event()
+
+    async def close_async() -> None:
+        close_started.set()
+        await close_release.wait()
+        close_finished.set()
+
+    mock_connection = AsyncMock()
+    mock_connection.close.side_effect = close_async
+    target._existing_conversation["conv"] = mock_connection
+
+    cleanup_task = asyncio.create_task(target.reset_conversation_async(conversation_id="conv"))
+    await close_started.wait()
+    cleanup_task.cancel()
+    await asyncio.sleep(0)
+
+    assert not cleanup_task.done()
+    close_release.set()
+    with pytest.raises(asyncio.CancelledError):
+        await cleanup_task
+
+    assert close_finished.is_set()
+    assert "conv" not in target._existing_conversation
+
+
 async def test_reset_conversation_async_propagates_cancellation(target):
     mock_connection = AsyncMock()
     mock_connection.close.side_effect = asyncio.CancelledError