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 a14b5073b3..a772acb39d 100644 --- a/pyrit/prompt_target/openai/openai_realtime_target.py +++ b/pyrit/prompt_target/openai/openai_realtime_target.py @@ -418,10 +418,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.cleanup_conversation_async(conversation_id) + raise response_type = request.converted_value_data_type @@ -486,14 +490,13 @@ async def cleanup_conversation_async(self, conversation_id: str) -> None: Args: conversation_id (str): The conversation ID to disconnect from. """ - connection = self._existing_conversation.get(conversation_id) + connection = self._existing_conversation.pop(conversation_id, None) if connection: try: await connection.close() logger.info(f"Disconnected from {self._endpoint} with conversation ID: {conversation_id}") except Exception as e: logger.warning(f"Error closing connection for {conversation_id}: {e}") - del self._existing_conversation[conversation_id] async def _connect_async(self, *, conversation_id: str) -> Any: """ @@ -777,24 +780,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.") @@ -827,7 +828,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") @@ -841,23 +842,29 @@ 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 an unfinished Realtime receive task.""" + if receive_task.done(): + return + 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 f737a0b993..51a810dbcd 100644 --- a/tests/unit/prompt_target/target/test_realtime_target.py +++ b/tests/unit/prompt_target/target/test_realtime_target.py @@ -85,6 +85,57 @@ 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_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())