diff --git a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py b/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py index 492f34f21d..552115ac36 100644 --- a/pyrit/prompt_target/hugging_face/hugging_face_chat_target.py +++ b/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. diff --git a/pyrit/prompt_target/openai/openai_realtime_target.py b/pyrit/prompt_target/openai/openai_realtime_target.py index 816a74118e..57127d314d 100644 --- a/pyrit/prompt_target/openai/openai_realtime_target.py +++ b/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. diff --git a/tests/unit/prompt_target/target/test_huggingface_chat_target.py b/tests/unit/prompt_target/target/test_huggingface_chat_target.py index ca6fad9f70..d18053b518 100644 --- a/tests/unit/prompt_target/target/test_huggingface_chat_target.py +++ b/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): diff --git a/tests/unit/prompt_target/target/test_realtime_target.py b/tests/unit/prompt_target/target/test_realtime_target.py index 80f47b864d..b8d499508b 100644 --- a/tests/unit/prompt_target/target/test_realtime_target.py +++ b/tests/unit/prompt_target/target/test_realtime_target.py @@ -3,6 +3,7 @@ import asyncio import base64 +import gc import wave from collections.abc import AsyncIterator from typing import Any @@ -86,6 +87,79 @@ 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_retrieves_completed_failure(target): + async def fail_receive_async() -> RealtimeTargetResult: + raise RuntimeError("receive failed") + + receive_task = asyncio.create_task(fail_receive_async()) + await asyncio.sleep(0) + assert receive_task.done() + + unhandled_exceptions: list[dict[str, Any]] = [] + loop = asyncio.get_running_loop() + previous_exception_handler = loop.get_exception_handler() + loop.set_exception_handler(lambda _loop, context: unhandled_exceptions.append(context)) + try: + await target._cancel_receive_task_async(receive_task=receive_task) + del receive_task + gc.collect() + finally: + loop.set_exception_handler(previous_exception_handler) + + assert not unhandled_exceptions + + 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 +1507,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