diff --git a/pyrit/prompt_target/openai/_openai_realtime_event_router.py b/pyrit/prompt_target/openai/_openai_realtime_event_router.py index 8deb710616..b9e301d9f6 100644 --- a/pyrit/prompt_target/openai/_openai_realtime_event_router.py +++ b/pyrit/prompt_target/openai/_openai_realtime_event_router.py @@ -88,6 +88,17 @@ def is_lifecycle_event(cls, event_kind: _OpenAIRealtimeEventKind) -> bool: """Return whether atomic receiving should log the event as lifecycle-only.""" return event_kind in cls._LIFECYCLE_KINDS + @staticmethod + def get_response_id(*, event: Any) -> str | None: + """Return the response ID carried directly or by a response payload.""" + response_id = getattr(event, "response_id", None) + if isinstance(response_id, str): + return response_id + + response = getattr(event, "response", None) + nested_response_id = getattr(response, "id", None) + return nested_response_id if isinstance(nested_response_id, str) else None + @staticmethod def collect_response_delta( *, diff --git a/pyrit/prompt_target/openai/openai_realtime_target.py b/pyrit/prompt_target/openai/openai_realtime_target.py index a14b5073b3..54e7a6f57a 100644 --- a/pyrit/prompt_target/openai/openai_realtime_target.py +++ b/pyrit/prompt_target/openai/openai_realtime_target.py @@ -580,9 +580,11 @@ async def receive_events_async(self, conversation_id: str) -> RealtimeTargetResu result = RealtimeTargetResult() audio_buffer = bytearray() - audio_done_received = False + audio_done_deadline: float | None = None + current_response_id: str | None = None current_turn_event_count = 0 grace_period_sec = 1.0 # Wait 1 second after audio.done before soft-finishing + loop = asyncio.get_running_loop() try: # Create event iterator @@ -591,13 +593,13 @@ async def receive_events_async(self, conversation_id: str) -> RealtimeTargetResu while True: # If we've seen audio.done, wait with a short timeout for response.done # Otherwise, wait indefinitely for events - timeout = grace_period_sec if audio_done_received else None + timeout = max(0.0, audio_done_deadline - loop.time()) if audio_done_deadline is not None else None try: event = await asyncio.wait_for(event_iter.__anext__(), timeout=timeout) except asyncio.TimeoutError: # Soft-finish: audio.done was received but no response.done after grace period - if audio_done_received: + if audio_done_deadline is not None: logger.warning( f"Soft-finishing: No response.done {grace_period_sec}s after audio.done. " f"Audio bytes: {len(audio_buffer)}" @@ -622,6 +624,16 @@ async def receive_events_async(self, conversation_id: str) -> RealtimeTargetResu event_type = event.type event_kind = _OpenAIRealtimeEventRouter.classify_event(event_type) + event_response_id = _OpenAIRealtimeEventRouter.get_response_id(event=event) + if event_kind is _OpenAIRealtimeEventKind.RESPONSE_CREATED and current_response_id is None: + current_response_id = event_response_id + elif event_response_id is not None and event_response_id != current_response_id: + logger.debug( + f"Skipping event '{event_type}' for response {event_response_id}; " + f"current response is {current_response_id}" + ) + continue + current_turn_event_count += 1 logger.debug(f"Processing event type: {event_type}") audio_size_before = len(audio_buffer) @@ -659,7 +671,8 @@ async def receive_events_async(self, conversation_id: str) -> RealtimeTargetResu elif event_kind is _OpenAIRealtimeEventKind.AUDIO_DONE: logger.debug(f"Received audio.done - will soft-finish in {grace_period_sec}s if no response.done") - audio_done_received = True + if audio_done_deadline is None: + audio_done_deadline = loop.time() + grace_period_sec elif event_kind is _OpenAIRealtimeEventKind.TRANSCRIPT_DELTA: if getattr(event, "delta", ""): diff --git a/tests/integration/targets/test_targets_and_secrets.py b/tests/integration/targets/test_targets_and_secrets.py index e2ec9da733..1d70d42bc4 100644 --- a/tests/integration/targets/test_targets_and_secrets.py +++ b/tests/integration/targets/test_targets_and_secrets.py @@ -95,6 +95,22 @@ def valid_response(resp: str) -> bool: raise AssertionError(f"LLM did not return exactly 'test' after {max_retries} attempts.") +async def _send_realtime_text_async( + *, + target: RealtimeTarget, + conversation_id: str, + text: str, +) -> str: + """Send one text turn through a RealtimeTarget and return its transcript.""" + message = MessagePiece( + role="user", + original_value=text, + conversation_id=conversation_id, + ).to_message() + response = await target.send_prompt_async(message=message) + return str(response[0].get_value()) + + async def _assert_can_send_video_prompt(*, target: PromptTarget) -> None: """Helper function to test video generation targets.""" video_prompt = "A raccoon sailing a pirate ship" @@ -523,6 +539,58 @@ async def test_realtime_target_multi_objective( assert len(result.last_response.converted_value) > 0 +@pytest.mark.parametrize( + ("endpoint", "api_key_env_var", "model_name"), + [ + pytest.param( + "PLATFORM_OPENAI_REALTIME_ENDPOINT", + "PLATFORM_OPENAI_REALTIME_KEY", + "PLATFORM_OPENAI_REALTIME_MODEL", + id="platform-api-key", + ), + pytest.param( + "AZURE_OPENAI_REALTIME_ENDPOINT", + None, + "AZURE_OPENAI_REALTIME_MODEL", + id="azure-entra", + ), + ], +) +@pytest.mark.run_only_if_all_tests +async def test_realtime_target_same_conversation_multi_turn( + sqlite_instance: SQLiteMemory, + endpoint: str, + api_key_env_var: str | None, + model_name: str, +) -> None: + """Test that a RealtimeTarget preserves context across turns on one connection.""" + endpoint_value = _get_required_env_var(endpoint) + target = RealtimeTarget( + endpoint=endpoint_value, + api_key=_get_openai_auth(endpoint=endpoint_value, api_key_env_var=api_key_env_var), + model_name=_get_required_env_var(model_name), + ) + conversation_id = str(uuid.uuid4()) + + try: + first_response = await _send_realtime_text_async( + target=target, + conversation_id=conversation_id, + text="What is the capital of France?", + ) + second_response = await _send_realtime_text_async( + target=target, + conversation_id=conversation_id, + text="What country is the city from your previous answer in?", + ) + + assert "paris" in first_response.lower() + assert "france" in second_response.lower() + assert set(target._existing_conversation) == {conversation_id} + finally: + await target.cleanup_target_async() + + @pytest.mark.parametrize( ("endpoint", "api_key"), [ diff --git a/tests/unit/prompt_target/target/test_realtime_target.py b/tests/unit/prompt_target/target/test_realtime_target.py index f737a0b993..a3a8164a74 100644 --- a/tests/unit/prompt_target/target/test_realtime_target.py +++ b/tests/unit/prompt_target/target/test_realtime_target.py @@ -4,6 +4,7 @@ import asyncio import base64 import wave +from collections.abc import AsyncIterator from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -441,6 +442,147 @@ async def _events(): assert result.transcripts == ["partial"] +async def test_receive_events_audio_done_deadline_is_not_extended_by_stale_or_noisy_events(target): + """Stale, duplicate, and unrelated events consume rather than reset the soft-finish grace period.""" + mock_connection = AsyncMock() + conversation_id = "test_bounded_soft_finish" + target._existing_conversation[conversation_id] = mock_connection + mock_connection.__aiter__.return_value = [ + _scripted_event("response.done", **{"response.status": "success"}), + _scripted_event("response.output_audio.done"), + _scripted_event("provider.noise"), + _scripted_event("response.output_audio.done"), + _scripted_event("session.updated"), + ] + observed_timeouts: list[float | None] = [] + wait_for = asyncio.wait_for + + async def _record_wait_for(awaitable: Any, *, timeout: float | None) -> Any: + observed_timeouts.append(timeout) + return await wait_for(awaitable, timeout=timeout) + + mock_loop = MagicMock() + mock_loop.time.side_effect = [100.0, 100.25, 100.5, 100.75, 101.0] + with ( + patch( + "pyrit.prompt_target.openai.openai_realtime_target.asyncio.get_running_loop", + return_value=mock_loop, + ), + patch( + "pyrit.prompt_target.openai.openai_realtime_target.asyncio.wait_for", + side_effect=_record_wait_for, + ), + ): + result = await target.receive_events_async(conversation_id) + + assert result.audio_bytes == b"" + assert observed_timeouts == [None, None, 0.75, 0.5, 0.25, 0.0] + + +async def test_receive_events_accepts_response_done_before_audio_done_deadline(target): + """A terminal event arriving within the remaining grace period still completes normally.""" + mock_connection = AsyncMock() + conversation_id = "test_late_terminal_event" + target._existing_conversation[conversation_id] = mock_connection + mock_connection.__aiter__.return_value = [ + _scripted_event("response.output_audio.done"), + _scripted_event("response.done", **{"response.status": "success"}), + ] + observed_timeouts: list[float | None] = [] + wait_for = asyncio.wait_for + + async def _record_wait_for(awaitable: Any, *, timeout: float | None) -> Any: + observed_timeouts.append(timeout) + return await wait_for(awaitable, timeout=timeout) + + mock_loop = MagicMock() + mock_loop.time.side_effect = [100.0, 100.99] + with ( + patch( + "pyrit.prompt_target.openai.openai_realtime_target.asyncio.get_running_loop", + return_value=mock_loop, + ), + patch( + "pyrit.prompt_target.openai.openai_realtime_target.asyncio.wait_for", + side_effect=_record_wait_for, + ), + ): + result = await target.receive_events_async(conversation_id) + + assert result.audio_bytes == b"" + assert observed_timeouts[0] is None + assert observed_timeouts[1] == pytest.approx(0.01) + + +async def test_receive_events_cancellation_during_audio_done_grace_propagates(target): + """Cancellation while waiting within the grace period is never converted into a soft finish.""" + mock_connection = AsyncMock() + conversation_id = "test_grace_cancellation" + target._existing_conversation[conversation_id] = mock_connection + + async def _events() -> AsyncIterator[Any]: + yield _scripted_event("response.output_audio.done") + raise asyncio.CancelledError + + mock_connection.__aiter__.side_effect = _events + mock_loop = MagicMock() + mock_loop.time.side_effect = [100.0, 100.25] + with patch( + "pyrit.prompt_target.openai.openai_realtime_target.asyncio.get_running_loop", + return_value=mock_loop, + ): + with pytest.raises(asyncio.CancelledError): + await target.receive_events_async(conversation_id) + + +async def test_receive_events_ignores_late_events_from_soft_finished_response(target): + """Late prior-turn deltas and completion events must not contaminate the next response.""" + mock_connection = AsyncMock() + conversation_id = "test_response_ownership" + target._existing_conversation[conversation_id] = mock_connection + + async def _first_response_events() -> AsyncIterator[Any]: + yield _scripted_event("response.created", **{"response.id": "response-1"}) + yield _scripted_event( + "response.audio.delta", + response_id="response-1", + delta=base64.b64encode(b"first").decode("ascii"), + ) + yield _scripted_event("response.audio.done", response_id="response-1") + raise asyncio.TimeoutError + + async def _second_response_events() -> AsyncIterator[Any]: + yield _scripted_event( + "response.audio_transcript.delta", + response_id="response-1", + delta="late first transcript", + ) + yield _scripted_event("response.done", **{"response.id": "response-1", "response.status": "success"}) + yield _scripted_event("response.created", **{"response.id": "response-2"}) + yield _scripted_event( + "response.audio.delta", + response_id="response-2", + delta=base64.b64encode(b"second").decode("ascii"), + ) + yield _scripted_event( + "response.audio_transcript.delta", + response_id="response-2", + delta="second transcript", + ) + yield _scripted_event("response.audio.done", response_id="response-2") + yield _scripted_event("response.done", **{"response.id": "response-2", "response.status": "success"}) + + event_streams = iter([_first_response_events(), _second_response_events()]) + mock_connection.__aiter__.side_effect = lambda: next(event_streams) + + first_result = await target.receive_events_async(conversation_id) + second_result = await target.receive_events_async(conversation_id) + + assert first_result.audio_bytes == b"first" + assert second_result.audio_bytes == b"second" + assert second_result.transcripts == ["second transcript"] + + async def test_receive_events_connection_close_soft_finishes_with_audio(target): """Atomic receiving returns accumulated audio when the provider closes before response.done."""