fix: finaliza STT sem bloquear comandos da sessão

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

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(