import asyncio import json import logging from queue import Empty, Queue from threading import Event from typing import Any import numpy as np from websockets.asyncio.server import ServerConnection from speech_to_speech.pipeline.control import SESSION_END, PipelineControlMessage, is_control_message from speech_to_speech.pipeline.events import PipelineEvent from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, PIPELINE_END from speech_to_speech.pipeline.queue_types import AudioInItem, AudioOutItem, TextEventItem logger = logging.getLogger(__name__) class WebSocketStreamer: """ Handles bidirectional audio streaming over WebSocket. Receives audio from clients and puts it in the input_queue. Sends audio from the output_queue to clients. Sends text messages (transcripts/tools) from text_output_queue to clients. """ def __init__( self, stop_event: Event, input_queue: Queue[AudioInItem], output_queue: Queue[AudioOutItem], should_listen: Event, text_output_queue: Queue[TextEventItem] | None = None, host: str = "0.0.0.0", port: int = 8765, ) -> None: self.stop_event = stop_event self.input_queue = input_queue # clients -> VAD self.output_queue = output_queue # TTS -> clients self.text_output_queue = text_output_queue # Text messages -> clients self.should_listen = should_listen self.host = host self.port = port self.clients: set[ServerConnection] = set() self.loop: asyncio.AbstractEventLoop | None = None self.server: Any = None def run(self) -> None: """Run the WebSocket server (called from a thread).""" self.loop = asyncio.new_event_loop() asyncio.set_event_loop(self.loop) try: self.loop.run_until_complete(self._run_server()) except Exception as e: logger.error(f"WebSocket server error: {e}") finally: self.loop.close() async def _run_server(self) -> None: """Main async server loop.""" import websockets logger.info(f"WebSocket server starting on ws://{self.host}:{self.port}") self.server = await websockets.serve( self._handle_client, self.host, self.port, ) logger.info("WebSocket server ready, waiting for connections...") # Start the sender task sender_task = asyncio.create_task(self._send_loop()) # Wait until stop_event is set while not self.stop_event.is_set(): await asyncio.sleep(0.1) # Cleanup sender_task.cancel() try: await sender_task except asyncio.CancelledError: pass # Close all clients for client in list(self.clients): try: await client.close() except Exception: pass self.server.close() await self.server.wait_closed() logger.info("WebSocket server closed") async def _handle_client(self, websocket: ServerConnection) -> None: """Handle a single WebSocket client connection.""" client_id = id(websocket) logger.info(f"Client {client_id} connected") self.clients.add(websocket) recv_buffer = bytearray() # Enable listening when first client connects if len(self.clients) == 1: # Drain edge queues so stale data from a previous session doesn't # leak into the new one (SESSION_END may not have flushed everything). for q in (self.output_queue, self.text_output_queue): if q is not None: while not q.empty(): try: q.get_nowait() except Empty: break self.should_listen.set() logger.debug("Listening enabled, edge queues drained (should_listen.set())") try: logger.debug(f"Client {client_id}: Starting message receive loop") async for message in websocket: if isinstance(message, bytes): logger.debug(f"Client {client_id}: Received {len(message)} bytes of audio") if self.should_listen.is_set(): # Split into 512-sample (1024 bytes) chunks for VAD. # Keep a per-client remainder buffer so no samples are dropped # when WebSocket frame boundaries are not aligned. chunk_size_bytes = 512 * 2 # 512 samples * 2 bytes per int16 recv_buffer.extend(message) num_chunks = 0 while len(recv_buffer) >= chunk_size_bytes: chunk = bytes(recv_buffer[:chunk_size_bytes]) del recv_buffer[:chunk_size_bytes] self.input_queue.put(chunk) num_chunks += 1 logger.debug(f"Client {client_id}: Queued {num_chunks} chunks for processing") else: logger.debug(f"Client {client_id}: Skipping audio (should_listen not set)") except Exception as e: logger.error(f"Client {client_id} error: {type(e).__name__}: {e}", exc_info=True) finally: self.clients.discard(websocket) logger.info(f"Client {client_id} disconnected (finally block)") if len(self.clients) == 0: logger.debug("Last WebSocket client disconnected, ending session") self.input_queue.put(SESSION_END) async def _send_loop(self) -> None: """Send audio and text from queues to all connected clients.""" # Buffer audio until we have at least 100ms worth (3200 bytes = 1600 samples at 16kHz int16) MIN_AUDIO_BYTES = 3200 audio_buffer = bytearray() while not self.stop_event.is_set(): try: # Check for audio try: audio_chunk = self.output_queue.get_nowait() if isinstance(audio_chunk, bytes) and audio_chunk == PIPELINE_END: if audio_buffer and self.clients: data = bytes(audio_buffer) audio_buffer.clear() await asyncio.gather( *[client.send(data) for client in self.clients], return_exceptions=True, ) break if isinstance(audio_chunk, bytes) and audio_chunk == AUDIO_RESPONSE_DONE: if audio_buffer and self.clients: data = bytes(audio_buffer) audio_buffer.clear() await asyncio.gather( *[client.send(data) for client in self.clients], return_exceptions=True, ) self.should_listen.set() logger.debug("Response complete, listening re-enabled") continue if is_control_message(audio_chunk, SESSION_END.kind): audio_buffer.clear() continue if isinstance(audio_chunk, PipelineControlMessage): continue if self.clients: chunk_bytes: bytes if isinstance(audio_chunk, bytes): chunk_bytes = audio_chunk elif isinstance(audio_chunk, np.ndarray): chunk_bytes = audio_chunk.tobytes() elif hasattr(audio_chunk, "tobytes"): chunk_bytes = audio_chunk.tobytes() else: continue audio_buffer.extend(chunk_bytes) if len(audio_buffer) >= MIN_AUDIO_BYTES: data = bytes(audio_buffer) audio_buffer.clear() logger.debug(f"Sending {len(data)} bytes of audio to {len(self.clients)} client(s)") await asyncio.gather( *[client.send(data) for client in self.clients], return_exceptions=True ) except Empty: # Flush any buffered audio when queue is empty if audio_buffer and self.clients: data = bytes(audio_buffer) audio_buffer.clear() logger.debug(f"Flushing {len(data)} bytes of audio to {len(self.clients)} client(s)") await asyncio.gather(*[client.send(data) for client in self.clients], return_exceptions=True) # Check for text/tool messages if self.text_output_queue: try: text_message = self.text_output_queue.get_nowait() if self.clients: if isinstance(text_message, PipelineEvent): payload = text_message.model_dump() await asyncio.gather( *[client.send(json.dumps(payload)) for client in self.clients], return_exceptions=True, ) elif isinstance(text_message, (PipelineControlMessage, bytes)): continue except Empty: pass await asyncio.sleep(0.01) except asyncio.CancelledError: break except Exception as e: logger.error(f"Send loop error: {e}")