FIX bound OpenAI Realtime completion grace

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(
         *,

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", ""):

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"),
     [

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."""