""" 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 fastapi import FastAPI, WebSocket, WebSocketDisconnect from pydantic import BaseModel, ValidationError import os 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, TurnState, ) 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, AudioStart, AudioStop, Interrupt, SessionEnd, SessionStart, error_event, AgentResponse, Interrupted, SessionEnded, SessionReady, STTFinal, STTPartial, TurnEnded, TurnStarted, ) from voice_assistant.audio_gateway.app.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 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 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, ), ) 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 turn_manager = create_turn_manager() agent: AgentAdapter | None = None tts: TTSAdapter | None = None pipeline: VoicePipeline | None = 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 event_type = payload.get("type") if not isinstance(event_type, str): await send_json( websocket, error_event( "INVALID_MESSAGE", "Field 'type' is required.", ), ) 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.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": 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." ) 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." ) # 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. 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 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 pending_audio_chunk is not None: raise ValueError( "audio.stop received while an audio chunk is pending." ) if agent is None: raise RuntimeError( "Agent adapter is not initialized." ) if tts is None: raise RuntimeError( "TTS adapter is not initialized." ) if pipeline is None: raise RuntimeError( "Voice pipeline is not initialized." ) # Finish the STT session normally. await stt.stop_session() await pipeline.wait_turn_task() stt_events_task = None try: pipeline.complete_turn(request.turn_id) except (RuntimeError, ValueError) as exc: raise RuntimeError( "Turn pipeline did not complete successfully." ) from exc 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 stt = None pending_audio_chunk = None await send_json( websocket, TurnEnded( session_id=session.session_id, turn_id=request.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: # Cancel and await STT/Agent/TTS execution before # confirming the interruption to the client. 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: # Preserve the existing tolerant behavior for unknown turns. pipeline.interrupt_if_active(request.turn_id) 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." ) # 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) 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: # 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, ) # A task may have been created before registration failed. # Handle it independently if the pipeline does not own it. 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, )