diff --git a/src/voice_assistant/audio_gateway/app/main.py b/src/voice_assistant/audio_gateway/app/main.py index 61c4c7b..2bc9c6c 100644 --- a/src/voice_assistant/audio_gateway/app/main.py +++ b/src/voice_assistant/audio_gateway/app/main.py @@ -113,6 +113,8 @@ app_config = load_config() provider_catalog = load_provider_catalog() MAX_AUDIO_CHUNK_BYTES = 64 * 1024 # 64 KB +STT_STOP_TIMEOUT_SECONDS = 10.0 +FINALIZATION_CANCEL_TIMEOUT_SECONDS = 5.0 WEB_DIR = Path(__file__).resolve().parents[1] / "web" @@ -298,6 +300,7 @@ async def audio_socket(websocket: WebSocket) -> None: finalization_task: asyncio.Task[None] | None = None + audio_input_closed = False turn_transition_lock = asyncio.Lock() interrupting_turns: set[int] = set() completed_turns: set[int] = set() @@ -310,16 +313,40 @@ async def audio_socket(websocket: WebSocket) -> None: async def finalize_turn( turn_id: int, + turn_stt: STTAdapter, ) -> None: """ - Complete a turn without blocking WebSocket message reception. - - Normal completion must not race with interruption. + Finalize STT and complete the turn without blocking reception. """ if pipeline is None or session is None: return try: + try: + async with asyncio.timeout(STT_STOP_TIMEOUT_SECONDS): + await turn_stt.stop_session() + except TimeoutError: + await handle_stt_stop_failure( + turn_id, + turn_stt, + code="STT_STOP_TIMEOUT", + message="STT finalization exceeded the time limit.", + ) + return + except Exception: + logger.exception( + "STT finalization failed. session_id=%s turn_id=%s", + session.session_id, + turn_id, + ) + await handle_stt_stop_failure( + turn_id, + turn_stt, + code="STT_STOP_FAILED", + message="STT could not finalize audio input.", + ) + return + await pipeline.wait_turn_task() async with turn_transition_lock: @@ -369,6 +396,50 @@ async def audio_socket(websocket: WebSocket) -> None: "Failed to report turn execution failure." ) + async def handle_stt_stop_failure( + turn_id: int, + turn_stt: STTAdapter, + *, + code: str, + message: str, + ) -> None: + if pipeline is None or session is None: + return + + if not pipeline.is_active_turn(turn_id): + return + + cleanup_complete = True + + for operation_name, operation in ( + ("turn task cancellation", pipeline.cancel_turn_task), + ("STT abort", turn_stt.abort_session), + ): + try: + await operation() + except Exception: + cleanup_complete = False + logger.exception( + "Failed during %s. session_id=%s turn_id=%s", + operation_name, + session.session_id, + turn_id, + ) + + if cleanup_complete: + pipeline.interrupt_if_active(turn_id) + + await send_json( + websocket, + error_event( + code, + message, + session_id=session.session_id, + turn_id=turn_id, + retryable=cleanup_complete, + ), + ) + async def cancel_finalization_task() -> None: """ Cancel and await the Gateway-owned finalization task. @@ -394,7 +465,7 @@ async def audio_socket(websocket: WebSocket) -> None: try: await asyncio.wait_for( asyncio.shield(task), - timeout=5.0, + timeout=FINALIZATION_CANCEL_TIMEOUT_SECONDS, ) except asyncio.CancelledError: current = asyncio.current_task() @@ -405,12 +476,10 @@ async def audio_socket(websocket: WebSocket) -> None: if not task.cancelled(): raise - except TimeoutError: - logger.error( - "Finalization task cancellation timed out. " - "session_id=%s", - session.session_id if session else None, - ) + except TimeoutError as exc: + raise TimeoutError( + "Turn finalization cancellation timed out." + ) from exc finally: if task.done() and finalization_task is task: @@ -808,6 +877,7 @@ async def audio_socket(websocket: WebSocket) -> None: continue # Startup completed successfully. + audio_input_closed = False await send_json( websocket, TurnStarted( @@ -824,6 +894,11 @@ async def audio_socket(websocket: WebSocket) -> None: try: request = AudioChunk.model_validate(payload) + if audio_input_closed: + raise ValueError( + "Audio input is closed for this turn." + ) + if session.active_turn_id != request.turn_id: raise ValueError( f"Unexpected turn_id={request.turn_id}; " @@ -895,12 +970,9 @@ async def audio_socket(websocket: WebSocket) -> None: "Turn finalization is already running." ) - await stt.stop_session() - - # Do not await OCI/TTS here. - # Keep the WebSocket receive loop available. + audio_input_closed = True finalization_task = asyncio.create_task( - finalize_turn(request.turn_id) + finalize_turn(request.turn_id, stt) ) except (ValidationError, RuntimeError, ValueError) as exc: @@ -1038,7 +1110,12 @@ async def audio_socket(websocket: WebSocket) -> None: stt = None pending_audio_chunk = None - except (ValidationError, RuntimeError, ValueError) as exc: + except ( + ValidationError, + RuntimeError, + ValueError, + TimeoutError, + ) as exc: await send_json( websocket, error_event( diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 93a4c5a..a439345 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -1,4 +1,6 @@ import asyncio +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager import threading import pytest from unittest.mock import patch, AsyncMock, Mock @@ -2528,3 +2530,179 @@ def test_websocket_duplicate_audio_start_preserves_active_turn(): assert websocket.receive_json()["type"] == "interrupted" assert stt._started is False assert created_turn_managers[0].active_turn is None + + +class ControlledStopSTTAdapter(MockSTTAdapter): + def __init__(self, *, ignore_cancellation=False): + super().__init__() + self.ignore_cancellation = ignore_cancellation + self.stop_started = threading.Event() + self.stop_cancelled = threading.Event() + self.stop_release = threading.Event() + self.stop_calls = 0 + + async def stop_session(self): + self.stop_calls += 1 + self.stop_started.set() + + try: + while not self.stop_release.is_set(): + await asyncio.sleep(0.01) + except asyncio.CancelledError: + self.stop_cancelled.set() + if not self.ignore_cancellation: + raise + while not self.stop_release.is_set(): + await asyncio.sleep(0.01) + + await super().stop_session() + + +@contextmanager +def open_turn_with_stt(stt): + created_turn_managers.clear() + with patch.object( + main, + "create_stt_adapter", + side_effect=[stt, MockSTTAdapter()], + ): + try: + with client.websocket_connect("/ws/audio") as websocket: + websocket.send_json({"type": "session.start"}) + assert websocket.receive_json()["type"] == "session.ready" + websocket.send_json({"type": "audio.start"}) + started = websocket.receive_json() + assert started["type"] == "turn.started" + yield websocket, started["turn_id"] + finally: + stt.stop_release.set() + + +def receive_while_stt_waits(websocket, stt): + with ThreadPoolExecutor(max_workers=1) as executor: + response = executor.submit(websocket.receive_json) + try: + return response.result(timeout=2.0) + except TimeoutError: + stt.stop_release.set() + raise + + +@pytest.mark.parametrize( + ("event_type", "expected"), + [("interrupt", "interrupted"), ("session.end", "session.ended")], +) +def test_websocket_accepts_control_while_stt_stop_waits(event_type, expected): + stt = ControlledStopSTTAdapter() + with open_turn_with_stt(stt) as (websocket, turn_id): + websocket.send_json({"type": "audio.stop", "turn_id": turn_id}) + assert stt.stop_started.wait(timeout=2.0) + websocket.send_json({"type": event_type, "turn_id": turn_id}) + + response = receive_while_stt_waits(websocket, stt) + assert response["type"] == expected + assert stt.stop_cancelled.is_set() + assert stt._started is False + assert created_turn_managers[0].active_turn is None + + if event_type == "interrupt": + websocket.send_json({"type": "audio.start"}) + assert websocket.receive_json()["type"] == "turn.started" + + +def test_websocket_disconnect_cancels_stt_stop(): + stt = ControlledStopSTTAdapter() + with open_turn_with_stt(stt) as (websocket, turn_id): + websocket.send_json({"type": "audio.stop", "turn_id": turn_id}) + assert stt.stop_started.wait(timeout=2.0) + + assert stt.stop_cancelled.is_set() + assert stt._started is False + assert created_turn_managers[0].active_turn is None + + +@pytest.mark.parametrize("failure", ["timeout", "provider"]) +@pytest.mark.parametrize("abort_fails", [False, True]) +def test_websocket_stt_stop_failure_releases_only_safe_turns( + monkeypatch, failure, abort_fails, +): + stt = ControlledStopSTTAdapter() + monkeypatch.setattr(main, "STT_STOP_TIMEOUT_SECONDS", 0.05) + if failure == "provider": + monkeypatch.setattr( + stt, "stop_session", + AsyncMock(side_effect=RuntimeError("STT stop failed")), + ) + if abort_fails: + monkeypatch.setattr( + stt, "abort_session", + AsyncMock(side_effect=RuntimeError("STT abort failed")), + ) + + with open_turn_with_stt(stt) as (websocket, turn_id): + websocket.send_json({"type": "audio.stop", "turn_id": turn_id}) + response = receive_while_stt_waits(websocket, stt) + + assert response["type"] == "error" + assert response["code"] == ( + "STT_STOP_TIMEOUT" if failure == "timeout" else "STT_STOP_FAILED" + ) + assert response["turn_id"] == turn_id + assert response["retryable"] is (not abort_fails) + + if abort_fails: + assert created_turn_managers[0].active_turn.turn_id == turn_id + else: + assert created_turn_managers[0].active_turn is None + assert stt._started is False + + websocket.send_json({"type": "audio.start"}) + next_response = websocket.receive_json() + if abort_fails: + assert next_response["code"] == "INVALID_TURN" + else: + assert next_response["type"] == "turn.started" + + +@pytest.mark.parametrize( + ("event_type", "expected"), + [("audio.stop", "INVALID_TURN"), ("audio.chunk", "AUDIO_SEQUENCE_ERROR")], +) +def test_websocket_rejects_input_while_stt_stop_waits(event_type, expected): + stt = ControlledStopSTTAdapter() + with open_turn_with_stt(stt) as (websocket, turn_id): + websocket.send_json({"type": "audio.stop", "turn_id": turn_id}) + assert stt.stop_started.wait(timeout=2.0) + websocket.send_json({ + "type": event_type, "turn_id": turn_id, "sequence_number": 0, + }) + + response = receive_while_stt_waits(websocket, stt) + assert response["code"] == expected + assert stt.stop_calls == 1 + websocket.send_json({"type": "interrupt", "turn_id": turn_id}) + assert receive_while_stt_waits(websocket, stt)["type"] == "interrupted" + + +@pytest.mark.parametrize( + ("event_type", "expected"), + [("interrupt", "TURN_CANCEL_FAILED"), ("session.end", "SESSION_END_FAILED")], +) +def test_websocket_stt_stop_cancel_timeout_preserves_turn( + monkeypatch, event_type, expected, +): + stt = ControlledStopSTTAdapter(ignore_cancellation=True) + monkeypatch.setattr(main, "FINALIZATION_CANCEL_TIMEOUT_SECONDS", 0.05) + with open_turn_with_stt(stt) as (websocket, turn_id): + websocket.send_json({"type": "audio.stop", "turn_id": turn_id}) + assert stt.stop_started.wait(timeout=2.0) + websocket.send_json({"type": event_type, "turn_id": turn_id}) + + response = receive_while_stt_waits(websocket, stt) + assert response["code"] == expected + assert stt.stop_cancelled.is_set() + assert stt._started is True + assert created_turn_managers[0].active_turn.turn_id == turn_id + websocket.send_json({"type": "audio.start"}) + assert websocket.receive_json()["code"] == "INVALID_TURN" + stt.stop_release.set()