diff --git a/src/agents/voice/models/openai_stt.py b/src/agents/voice/models/openai_stt.py index 939af0be62..f14db755c7 100644 --- a/src/agents/voice/models/openai_stt.py +++ b/src/agents/voice/models/openai_stt.py @@ -40,6 +40,10 @@ class WebsocketDoneSentinel: pass +class _ListenerError(Exception): + pass + + def _audio_to_base64(audio_data: list[npt.NDArray[np.int16 | np.float32]]) -> str: return _audio_buffer_to_base64(np.concatenate(audio_data)) @@ -55,7 +59,9 @@ def _audio_buffer_to_base64(buffer: npt.NDArray[np.int16 | np.float32]) -> str: async def _wait_for_event( - event_queue: asyncio.Queue[dict[str, Any]], expected_types: list[str], timeout: float + event_queue: asyncio.Queue[dict[str, Any] | ErrorSentinel], + expected_types: list[str], + timeout: float, ): """ Wait for an event from event_queue whose type is in expected_types within the specified timeout. @@ -66,6 +72,8 @@ async def _wait_for_event( if remaining <= 0: raise TimeoutError(f"Timeout waiting for event(s): {expected_types}") evt = await asyncio.wait_for(event_queue.get(), timeout=remaining) + if isinstance(evt, ErrorSentinel): + raise _ListenerError("Websocket listener failed") from evt.error evt_type = evt.get("type", "") if evt_type in expected_types: return evt @@ -98,8 +106,10 @@ def __init__( asyncio.Queue() ) self._websocket: websockets.ClientConnection | None = None - self._event_queue: asyncio.Queue[dict[str, Any] | WebsocketDoneSentinel] = asyncio.Queue() - self._state_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue() + self._event_queue: asyncio.Queue[dict[str, Any] | ErrorSentinel | WebsocketDoneSentinel] = ( + asyncio.Queue() + ) + self._state_queue: asyncio.Queue[dict[str, Any] | ErrorSentinel] = asyncio.Queue() self._turn_audio_buffer: list[npt.NDArray[np.int16 | np.float32]] = [] self._tracing_span: Span[TranscriptionSpanData] | None = None @@ -140,8 +150,8 @@ def _end_turn(self, _transcript: str) -> None: async def _event_listener(self) -> None: assert self._websocket is not None, "Websocket not initialized" - async for message in self._websocket: - try: + try: + async for message in self._websocket: event = json.loads(message) if event.get("type") == "error": @@ -156,10 +166,12 @@ async def _event_listener(self) -> None: await self._state_queue.put(event) await self._event_queue.put(event) - except Exception as e: - await self._output_queue.put(ErrorSentinel(e)) - raise STTWebsocketConnectionError("Error parsing events") from e - await self._event_queue.put(WebsocketDoneSentinel()) + except Exception as e: + error = ErrorSentinel(e) + await self._event_queue.put(error) + await self._state_queue.put(error) + finally: + await self._event_queue.put(WebsocketDoneSentinel()) async def _configure_session(self) -> None: assert self._websocket is not None, "Websocket not initialized" @@ -191,6 +203,8 @@ async def _setup_connection(self, ws: websockets.ClientConnection) -> None: ["session.created", "transcription_session.created"], SESSION_CREATION_TIMEOUT, ) + except _ListenerError: + raise except TimeoutError as e: wrapped_err = STTWebsocketConnectionError( "Timeout waiting for transcription_session.created event" @@ -213,6 +227,8 @@ async def _setup_connection(self, ws: websockets.ClientConnection) -> None: logger.debug("Session updated") else: logger.debug("Session updated: %s", event) + except _ListenerError: + raise except TimeoutError as e: wrapped_err = STTWebsocketConnectionError( "Timeout waiting for transcription_session.updated event" @@ -232,6 +248,8 @@ async def _handle_events(self) -> None: if isinstance(event, WebsocketDoneSentinel): # processed all events and websocket is done break + if isinstance(event, ErrorSentinel): + raise STTWebsocketConnectionError("Error parsing events") from event.error event_type = event.get("type", "unknown") if event_type in [ @@ -298,6 +316,11 @@ async def _process_websocket_connection(self) -> None: else: logger.error("Listener task not initialized") raise AgentsException("Listener task not initialized") + except _ListenerError as e: + if self._process_events_task is None: + self._process_events_task = asyncio.create_task(self._handle_events()) + await self._process_events_task + raise STTWebsocketConnectionError("Error parsing events") from e.__cause__ except Exception as e: await self._output_queue.put(ErrorSentinel(e)) raise diff --git a/tests/voice/test_openai_stt.py b/tests/voice/test_openai_stt.py index 1514676bbe..e75e27c37a 100644 --- a/tests/voice/test_openai_stt.py +++ b/tests/voice/test_openai_stt.py @@ -30,6 +30,7 @@ from agents.voice.models.openai_stt import ( EVENT_INACTIVITY_TIMEOUT, ErrorSentinel, + WebsocketDoneSentinel, _audio_buffer_to_base64, ) @@ -612,19 +613,29 @@ def fake_time_func(): @pytest.mark.asyncio -async def test_session_error_event(): +async def test_session_error_event(monkeypatch: pytest.MonkeyPatch): """ - If the session receives an event with "type": "error", it should propagate an exception - and put an ErrorSentinel in the output queue. + If the session receives an event with "type": "error", it should emit preceding transcripts, + drain the event processor, and then propagate an exception. """ mock_ws = create_mock_websocket( [ json.dumps({"type": "transcription_session.created"}), json.dumps({"type": "transcription_session.updated"}), + json.dumps( + { + "type": "conversation.item.input_audio_transcription.completed", + "transcript": "Transcript before error", + } + ), # Then an error from the server json.dumps({"type": "error", "error": "Simulated server error!"}), ] ) + monkeypatch.setattr( + "agents.voice.models.openai_stt.EVENT_INACTIVITY_TIMEOUT", + 0.1, + ) with patch("websockets.connect", return_value=mock_ws): audio_input = await FakeStreamedAudioInput.get(count=2) @@ -638,13 +649,95 @@ async def test_session_error_event(): trace_include_sensitive_data=False, trace_include_sensitive_audio_data=False, ) + event_queue_put = AsyncMock(wraps=session._event_queue.put) + monkeypatch.setattr(session._event_queue, "put", event_queue_put) + collected_turns: list[str] = [] with pytest.raises(STTWebsocketConnectionError): turns = session.transcribe_turns() - async for _ in turns: - pass + async for turn in turns: + collected_turns.append(turn) + assert collected_turns == ["Transcript before error"] + assert any( + isinstance(call.args[0], WebsocketDoneSentinel) + for call in event_queue_put.await_args_list + ) await session.close() + assert session._process_events_task is not None + assert session._process_events_task.done() + assert not session._process_events_task.cancelled() + + +@pytest.mark.asyncio +async def test_session_error_event_before_session_created(): + mock_ws = create_mock_websocket( + [json.dumps({"type": "error", "error": "Simulated setup error!"})] + ) + + with patch("websockets.connect", return_value=mock_ws): + audio_input = await FakeStreamedAudioInput.get(count=2) + session = OpenAISTTTranscriptionSession( + input=audio_input, + client=AsyncMock(api_key="FAKE_KEY"), + model="whisper-1", + settings=STTModelSettings(), + trace_include_sensitive_data=False, + trace_include_sensitive_audio_data=False, + ) + + async def consume_turns() -> None: + async for _ in session.transcribe_turns(): + pass + + with pytest.raises(STTWebsocketConnectionError): + await asyncio.wait_for(consume_turns(), timeout=1) + + assert session._process_events_task is not None + assert session._process_events_task.done() + assert not session._process_events_task.cancelled() + + +@pytest.mark.asyncio +async def test_listener_timeout_drains_buffered_transcript_before_setup(): + messages = [ + json.dumps( + { + "type": "conversation.item.input_audio_transcription.completed", + "transcript": "Transcript before listener timeout", + } + ) + ] + + async def messages_then_timeout() -> AsyncGenerator[str, None]: + for message in messages: + yield message + raise TimeoutError("Simulated listener timeout") + + mock_ws = AsyncMock() + mock_ws.__aenter__.return_value = mock_ws + mock_ws.__aiter__.side_effect = messages_then_timeout + + with patch("websockets.connect", return_value=mock_ws): + audio_input = await FakeStreamedAudioInput.get(count=2) + session = OpenAISTTTranscriptionSession( + input=audio_input, + client=AsyncMock(api_key="FAKE_KEY"), + model="whisper-1", + settings=STTModelSettings(), + trace_include_sensitive_data=False, + trace_include_sensitive_audio_data=False, + ) + + collected_turns: list[str] = [] + with pytest.raises(STTWebsocketConnectionError): + async for turn in session.transcribe_turns(): + collected_turns.append(turn) + + assert collected_turns == ["Transcript before listener timeout"] + assert session._process_events_task is not None + assert session._process_events_task.done() + assert not session._process_events_task.cancelled() @pytest.mark.asyncio