""" Audio Gateway application and WebSocket transport for the Voice Assistant. Purpose: Expose the realtime audio protocol and connect client sessions to the VoicePipeline orchestration layer. Responsibilities: - Accept and validate WebSocket protocol messages. - Receive canonical PCM audio frames. - Create and configure session-scoped components. - Forward STT events and TTS audio to connected clients. - Delegate semantic turn lifecycle to VoicePipeline. - Coordinate transport cleanup on disconnect. - Report protocol and cleanup errors. Non-responsibilities: - Implementing STT, Agent, or TTS provider logic. - Managing semantic turn state directly. - Loading provider implementations dynamically. - Implementing provider-specific cancellation. Architecture: WebSocket Client | v Audio Gateway | v VoicePipeline | | | STT Agent TTS """ import json import logging from pathlib import Path import anyio from fastapi.responses import FileResponse from fastapi.staticfiles import StaticFiles from fastapi import FastAPI, WebSocket, WebSocketDisconnect from pydantic import BaseModel, ValidationError from contextlib import aclosing import asyncio from voice_assistant.voice_pipeline import VoicePipeline from voice_assistant.stt.adapter import STTAdapter from voice_assistant.tts.adapter import TTSAdapter from voice_assistant.turn_manager import ( TurnManager, ) from voice_assistant.agent.adapter import AgentAdapter from voice_assistant.audio.models import ( CANONICAL_AUDIO_FORMAT, ) from voice_assistant.providers import ( create_agent_adapter, create_stt_adapter, create_tts_adapter, ) from voice_assistant.config.loader import ( load_config, load_provider_catalog, ) from voice_assistant.audio_gateway.app.protocol import ( PROTOCOL_VERSION, AUDIO_ENCODING, AUDIO_SAMPLE_RATE, AUDIO_CHANNELS, AudioChunk, AudioStop, Interrupt, SessionEnd, SessionStart, error_event, AgentResponse, Interrupted, SessionEnded, SessionReady, STTFinal, STTPartial, TurnEnded, TurnStarted, AgentDelta, ) from voice_assistant.voice_pipeline import ( PipelineTextDelta, PipelineAgentResponse, PipelineAudioChunk, ) from voice_assistant.domain.session import SessionState # logger configuration logger = logging.getLogger(__name__) app = FastAPI( title="Voice Assistant Audio Gateway", version=PROTOCOL_VERSION, ) 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" app.mount( "/web", StaticFiles(directory=WEB_DIR), name="web", ) @app.get("/ui", include_in_schema=False) async def web_client() -> FileResponse: return FileResponse(WEB_DIR / "index.html") def create_turn_manager() -> TurnManager: """Create the turn manager used by the audio gateway.""" return TurnManager() async def send_json( websocket: WebSocket, payload: BaseModel | dict, ) -> None: """ Serialize and send one protocol event through the WebSocket. Typed Pydantic protocol models are preferred. Dictionary support remains temporarily available for error events until ErrorEvent is migrated to the typed protocol. """ if isinstance(payload, BaseModel): data = payload.model_dump() else: data = payload await websocket.send_text( json.dumps( data, separators=(",", ":"), ) ) async def close_optional_provider(provider: object | None) -> None: """ Close providers exposing an optional asynchronous cleanup hook. Providers without aclose() require no action. """ if provider is None: return close_method = getattr(provider, "aclose", None) if callable(close_method): await close_method() async def forward_stt_events( websocket: WebSocket, stt: STTAdapter, pipeline: VoicePipeline, session_id: str, turn_id: int, ) -> None: async for event in stt.events(): if not pipeline.is_active_turn(turn_id): logger.debug( "Discarding stale STT event. session_id=%s turn_id=%s", session_id, turn_id, ) break if event.type == "partial": pipeline.on_partial( turn_id, event.text, ) await send_json( websocket, STTPartial( session_id=session_id, turn_id=turn_id, text=event.text, ), ) elif event.type == "final": pipeline.on_final( turn_id, event.text, ) await send_json( websocket, STTFinal( session_id=session_id, turn_id=turn_id, text=event.text, ), ) if callable(getattr(pipeline.agent, "stream", None)): async with aclosing( pipeline.stream_turn_response( session_id=session_id, turn_id=turn_id, text=event.text, ) ) as response_events: async for response_event in response_events: if not pipeline.is_active_turn(turn_id): break if isinstance(response_event, PipelineTextDelta): await send_json( websocket, AgentDelta( session_id=session_id, turn_id=turn_id, text=response_event.text, ), ) elif isinstance(response_event, PipelineAgentResponse): await send_json( websocket, AgentResponse( session_id=session_id, turn_id=turn_id, text=response_event.text, ), ) elif isinstance(response_event, PipelineAudioChunk): await websocket.send_bytes(response_event.audio) else: result = await pipeline.process_transcript( session_id=session_id, turn_id=turn_id, text=event.text, ) await send_json( websocket, AgentResponse( session_id=session_id, turn_id=turn_id, text=result.response_text, ), ) async for audio_chunk in pipeline.synthesize_response( result.response_text, ): await websocket.send_bytes(audio_chunk) pipeline.on_tts_complete(turn_id) @app.get("/health") async def health() -> dict[str, str]: return {"status": "ok"} @app.get("/ready") async def ready() -> dict[str, str]: return {"status": "ready"} @app.websocket("/ws/audio") async def audio_socket(websocket: WebSocket) -> None: await websocket.accept() session: SessionState | None = None pending_audio_chunk: AudioChunk | None = None stt: STTAdapter | None = None stt_events_task: asyncio.Task | None = 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() turn_manager = create_turn_manager() agent: AgentAdapter | None = None tts: TTSAdapter | None = None pipeline: VoicePipeline | None = None async def finalize_turn( turn_id: int, turn_stt: STTAdapter, ) -> None: """ 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: if turn_id in interrupting_turns: return if turn_id in completed_turns: return if not pipeline.is_active_turn(turn_id): return pipeline.complete_turn(turn_id) completed_turns.add(turn_id) await send_json( websocket, TurnEnded( session_id=session.session_id, turn_id=turn_id, ), ) except asyncio.CancelledError: raise except Exception as exc: logger.exception( "Turn finalization failed. session_id=%s turn_id=%s", session.session_id, turn_id, ) if turn_id not in interrupting_turns: try: await send_json( websocket, error_event( "TURN_EXECUTION_FAILED", str(exc), session_id=session.session_id, turn_id=turn_id, ), ) except Exception: logger.exception( "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. This task is separate from the execution task owned by VoicePipeline. """ nonlocal finalization_task task = finalization_task if task is None: return if task is asyncio.current_task(): raise RuntimeError( "Finalization task cannot cancel itself." ) if not task.done(): task.cancel() try: await asyncio.wait_for( asyncio.shield(task), timeout=FINALIZATION_CANCEL_TIMEOUT_SECONDS, ) except asyncio.CancelledError: current = asyncio.current_task() if current is not None and current.cancelling(): raise if not task.cancelled(): raise except TimeoutError as exc: raise TimeoutError( "Turn finalization cancellation timed out." ) from exc finally: if task.done() and finalization_task is task: finalization_task = None try: while True: message = await websocket.receive() if message.get("type") == "websocket.disconnect": break # --------------------------------------------------------- # Binary WebSocket frame -> PCM # --------------------------------------------------------- if message.get("bytes") is not None: if session is None: await send_json( websocket, error_event( "SESSION_NOT_FOUND", "session.start is required before audio.", ), ) continue if pending_audio_chunk is None: await send_json( websocket, error_event( "INVALID_AUDIO_FRAME", "Binary PCM frame must be preceded by audio.chunk.", session_id=session.session_id, turn_id=session.active_turn_id, ), ) continue audio_payload = message["bytes"] if len(audio_payload) > MAX_AUDIO_CHUNK_BYTES: await send_json( websocket, error_event( "MESSAGE_TOO_LARGE", ( "Audio chunk exceeds maximum size of " f"{MAX_AUDIO_CHUNK_BYTES} bytes." ), session_id=session.session_id, turn_id=pending_audio_chunk.turn_id, details={ "max_bytes": MAX_AUDIO_CHUNK_BYTES, "received_bytes": len(audio_payload), }, ), ) pending_audio_chunk = None continue try: session.accept_audio_chunk( pending_audio_chunk.turn_id, pending_audio_chunk.sequence_number, len(audio_payload), ) except (RuntimeError, ValueError) as exc: await send_json( websocket, error_event( "AUDIO_SEQUENCE_ERROR", str(exc), session_id=session.session_id, turn_id=pending_audio_chunk.turn_id, details={ "expected_sequence_number": session.expected_audio_sequence, }, ), ) pending_audio_chunk = None continue # echo do PCM para validar o transporte. if stt is None: await send_json( websocket, error_event( "STT_NOT_STARTED", "audio.start is required before audio.", session_id=session.session_id, turn_id=session.active_turn_id, ), ) pending_audio_chunk = None continue try: await stt.send_audio(audio_payload) except (RuntimeError, ValueError) as exc: await send_json( websocket, error_event( "STT_AUDIO_ERROR", str(exc), session_id=session.session_id, turn_id=pending_audio_chunk.turn_id, retryable=True, ), ) pending_audio_chunk = None continue pending_audio_chunk = None continue # --------------------------------------------------------- # Text WebSocket frame -> JSON control event # --------------------------------------------------------- text = message.get("text") if text is None: await send_json( websocket, error_event( "INVALID_MESSAGE", "Unsupported WebSocket frame.", ), ) continue try: payload = json.loads(text) except json.JSONDecodeError: await send_json( websocket, error_event( "INVALID_MESSAGE", "Message must contain valid JSON.", ), ) continue if not isinstance(payload, dict): await send_json( websocket, error_event( "INVALID_MESSAGE", "Message must be a JSON object.", ), ) continue event_type = payload.get("type") if not isinstance(event_type, str) or not event_type.strip(): await send_json( websocket, error_event( "INVALID_MESSAGE", "Field 'type' must be a non-empty string.", ), ) continue # --------------------------------------------------------- # session.start # --------------------------------------------------------- if event_type == "session.start": if session is not None: await send_json( websocket, error_event( "SESSION_ALREADY_STARTED", "This WebSocket already has an active session.", session_id=session.session_id, ), ) continue try: request = SessionStart.model_validate(payload) except ValidationError as exc: await send_json( websocket, error_event( "INVALID_MESSAGE", "Invalid session.start payload.", details={"errors": exc.errors()}, ), ) continue audio = request.audio if ( audio.encoding != AUDIO_ENCODING or audio.sample_rate != AUDIO_SAMPLE_RATE or audio.channels != AUDIO_CHANNELS ): await send_json( websocket, error_event( "INVALID_AUDIO_FORMAT", ( "Only PCM S16LE, mono, 16 kHz " "is supported in the MVP." ), details={ "expected": { "encoding": AUDIO_ENCODING, "sample_rate": AUDIO_SAMPLE_RATE, "channels": AUDIO_CHANNELS, } }, ), ) continue session = SessionState( session_id=request.session_id or SessionState().session_id ) agent = create_agent_adapter( config=app_config, catalog=provider_catalog, ) tts = create_tts_adapter( config=app_config, catalog=provider_catalog, ) pipeline = VoicePipeline( session=session, turn_manager=turn_manager, agent=agent, tts=tts, ) await send_json( websocket, SessionReady( protocol_version=PROTOCOL_VERSION, session_id=session.session_id, audio=CANONICAL_AUDIO_FORMAT, ), ) continue # --------------------------------------------------------- # Todos os eventos abaixo exigem sessão ativa # --------------------------------------------------------- if session is None: await send_json( websocket, error_event( "SESSION_NOT_FOUND", "session.start is required before this event.", ), ) continue if session.closed: await send_json( websocket, error_event( "SESSION_ALREADY_CLOSED", "This session is already closed.", session_id=session.session_id, ), ) continue # --------------------------------------------------------- # audio.start # --------------------------------------------------------- if event_type == "audio.start": if session.active_turn_id is not None: await send_json( websocket, error_event( "INVALID_TURN", "A turn is already active.", session_id=session.session_id, turn_id=session.active_turn_id, ), ) continue turn_id: int | None = None try: if pipeline is None: raise RuntimeError( "Voice pipeline is not initialized." ) turn = pipeline.start_turn() turn_id = turn.turn_id # Each turn receives a fresh STT provider instance. stt = create_stt_adapter( config=app_config, catalog=provider_catalog, ) await stt.start_session() stt_events_task = asyncio.create_task( forward_stt_events( websocket=websocket, stt=stt, pipeline=pipeline, session_id=session.session_id, turn_id=turn_id, ) ) pipeline.register_turn_task(stt_events_task) except (ValidationError, RuntimeError, ValueError) as exc: # Preserve the original startup failure. startup_error = exc cleanup_complete = True # Cancel and await any task created during startup. if stt_events_task is not None: task = stt_events_task try: if ( pipeline is not None and pipeline.turn_task is task ): await pipeline.cancel_turn_task() else: if not task.done(): task.cancel() try: await asyncio.wait_for( asyncio.shield(task), timeout=5.0, ) except asyncio.CancelledError: if not task.cancelled(): raise except Exception: cleanup_complete = False logger.exception( "Failed to cancel task during audio.start rollback. " "session_id=%s turn_id=%s", session.session_id if session else None, turn_id, ) if task.done(): stt_events_task = None else: cleanup_complete = False # Abort the partially initialized STT provider. if stt is not None: try: await stt.abort_session() except Exception: cleanup_complete = False logger.exception( "Failed to abort STT during audio.start rollback. " "session_id=%s turn_id=%s", session.session_id if session else None, turn_id, ) # Only release semantic state after execution # has stopped and provider cleanup succeeded. if cleanup_complete: if turn_id is not None and pipeline is not None: pipeline.interrupt_turn(turn_id) stt = None pending_audio_chunk = None await send_json( websocket, error_event( "STT_START_FAILED", str(startup_error), session_id=session.session_id, turn_id=turn_id, retryable=cleanup_complete, ), ) continue # Startup completed successfully. audio_input_closed = False await send_json( websocket, TurnStarted( session_id=session.session_id, turn_id=turn_id, ), ) continue # --------------------------------------------------------- # audio.chunk # --------------------------------------------------------- if event_type == "audio.chunk": 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}; " f"active_turn_id={session.active_turn_id}" ) if pending_audio_chunk is not None: raise ValueError( "Previous audio chunk is still pending PCM data." ) pending_audio_chunk = request except (ValidationError, ValueError) as exc: await send_json( websocket, error_event( "AUDIO_SEQUENCE_ERROR", str(exc), session_id=session.session_id, turn_id=payload.get("turn_id"), details={ "expected_sequence_number": session.expected_audio_sequence, }, ), ) continue continue # --------------------------------------------------------- # audio.stop # --------------------------------------------------------- if event_type == "audio.stop": try: request = AudioStop.model_validate(payload) if stt is None: raise RuntimeError( "STT session has not been started." ) if pipeline is None: raise RuntimeError( "Voice pipeline is not initialized." ) if agent is None or tts is None: raise RuntimeError( "Agent or TTS adapter is not initialized." ) if not pipeline.is_active_turn(request.turn_id): raise ValueError( f"Unexpected turn_id={request.turn_id}." ) if pending_audio_chunk is not None: raise ValueError( "audio.stop received while an audio chunk is pending." ) if ( finalization_task is not None and not finalization_task.done() ): raise RuntimeError( "Turn finalization is already running." ) audio_input_closed = True finalization_task = asyncio.create_task( finalize_turn(request.turn_id, stt) ) except (ValidationError, RuntimeError, ValueError) as exc: await send_json( websocket, error_event( "INVALID_TURN", str(exc), session_id=session.session_id, turn_id=payload.get("turn_id"), ), ) continue # --------------------------------------------------------- # interrupt # --------------------------------------------------------- if event_type == "interrupt": try: request = Interrupt.model_validate(payload) except ValidationError as exc: await send_json( websocket, error_event( "INVALID_MESSAGE", "Invalid interrupt payload.", session_id=session.session_id, details={"errors": exc.errors()}, ), ) continue if pipeline is None: await send_json( websocket, error_event( "INVALID_TURN", "Voice pipeline is not initialized.", session_id=session.session_id, turn_id=request.turn_id, ), ) continue # Only cancel resources belonging to the requested active turn. if pipeline.is_active_turn(request.turn_id): try: async with turn_transition_lock: interrupting_turns.add(request.turn_id) await cancel_finalization_task() await pipeline.cancel_turn_task() # Abort transcription without requesting a final result. if stt is not None: await stt.abort_session() # Release semantic turn state after execution stops. pipeline.interrupt_if_active(request.turn_id) # Clear references to the interrupted execution. stt_events_task = None pending_audio_chunk = None except Exception as exc: await send_json( websocket, error_event( "TURN_CANCEL_FAILED", str(exc), session_id=session.session_id, turn_id=request.turn_id, ), ) continue else: await send_json( websocket, error_event( "INVALID_TURN", ( f"Turn {request.turn_id} " "is no longer active." ), session_id=session.session_id, turn_id=request.turn_id, ), ) continue await send_json( websocket, Interrupted( session_id=session.session_id, turn_id=request.turn_id, ), ) continue # --------------------------------------------------------- # session.end # --------------------------------------------------------- if event_type == "session.end": try: SessionEnd.model_validate(payload) if pipeline is None: raise RuntimeError( "Voice pipeline is not initialized." ) # Prevent normal completion during session shutdown. async with turn_transition_lock: if session.active_turn_id is not None: interrupting_turns.add(session.active_turn_id) await cancel_finalization_task() # Cancel and await any running STT/Agent/TTS task. await pipeline.cancel_turn_task() # Abort STT without waiting for a final transcript. if stt is not None: await stt.abort_session() # Close semantic turn and session state. pipeline.close_session() # Release Gateway-owned references. stt_events_task = None stt = None pending_audio_chunk = None except ( ValidationError, RuntimeError, ValueError, TimeoutError, ) as exc: await send_json( websocket, error_event( "SESSION_END_FAILED", str(exc), session_id=session.session_id, turn_id=session.active_turn_id, retryable=True, ), ) continue await send_json( websocket, SessionEnded( session_id=session.session_id, ), ) break # --------------------------------------------------------- # Evento desconhecido # --------------------------------------------------------- await send_json( websocket, error_event( "UNSUPPORTED_EVENT", f"Unsupported event type: {event_type!r}.", session_id=session.session_id, turn_id=session.active_turn_id, ), ) except WebSocketDisconnect: # The finally block owns connection cleanup. pass finally: with anyio.CancelScope(shield=True): # Stop Gateway-owned finalization before provider cleanup. try: await cancel_finalization_task() except Exception: logger.exception( "Failed to cancel finalization task. session_id=%s", session.session_id if session else None, ) # Cancel and await the task owned by VoicePipeline. if pipeline is not None: try: await pipeline.cancel_turn_task() except Exception: logger.exception( "Failed to cancel voice turn task. session_id=%s", session.session_id if session else None, ) # Handle tasks created before pipeline registration. if stt_events_task is not None: task = stt_events_task if ( pipeline is None or pipeline.turn_task is not task ): try: if not task.done(): task.cancel() try: await asyncio.wait_for( asyncio.shield(task), timeout=5.0, ) except asyncio.CancelledError: if not task.cancelled(): raise except TimeoutError: logger.error( "Unregistered STT task cleanup timed out. " "session_id=%s", session.session_id if session else None, ) except Exception: logger.exception( "Failed to clean up unregistered STT task. " "session_id=%s", session.session_id if session else None, ) if task.done(): stt_events_task = None # Abort STT without waiting for a final transcript. if stt is not None: try: await stt.abort_session() except Exception: logger.exception( "Failed to abort STT session. session_id=%s", session.session_id if session else None, ) # Release session and semantic turn state. try: if pipeline is not None: pipeline.close_session() elif session is not None: session.close() except Exception: logger.exception( "Failed to close voice session. session_id=%s", session.session_id if session else None, ) # Release provider-owned HTTP connections. if agent is not None: try: await close_optional_provider(agent) except Exception: logger.exception( "Failed to close Agent provider. session_id=%s", session.session_id if session else None, ) finally: agent = None