alex.a.alves/voiceagent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: finaliza STT sem bloquear comandos da sessão

Browse Source
This commit is contained in:
alex alves
2026-11-03 11:12:40 -03:00
parent 153b2986b7
commit 4523910077
2 changed files with 271 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
src/voice_assistant/audio_gateway/app/main.py
View File

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

178
tests/test_gateway.py
View File

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