fix: finaliza STT sem bloquear comandos da sessão

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.
+        Finalize STT and complete the turn without blocking reception.
-
-        Normal completion must not race with interruption.
         """
         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:
+            except TimeoutError as exc:
-                logger.error(
+                raise TimeoutError(
-                    "Finalization task cancellation timed out. "
+                    "Turn finalization cancellation timed out."
-                    "session_id=%s",
+                ) from exc
-                    session.session_id if session else None,
-                )
 
             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()
+                    audio_input_closed = True
-
-                    # Do not await OCI/TTS here.
-                    # Keep the WebSocket receive loop available.
                     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(

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()