from __future__ import annotations import asyncio import contextlib import logging import time from dataclasses import dataclass, field from typing import Any, Callable from services.ai_voice_runtime_service.audiosocket import ( AUDIO_SOCKET_PACKET_DTMF, AUDIO_SOCKET_PACKET_HANGUP, AUDIO_SOCKET_PACKET_PCM16, AUDIO_SOCKET_PACKET_UUID, EnergyVAD, chunk_audio, encode_audio_packet, normalize_media_uuid, pcm16le_to_wav_bytes, read_packet, resample_pcm16le, ) from services.ai_voice_runtime_service.providers.asr import ASRProvider from services.ai_voice_runtime_service.providers.tts import TTSProvider from services.shared.models import VoiceAITurnDecisionOut logger = logging.getLogger("uvicorn.error") @dataclass(slots=True) class MediaRegistration: voice_session_id: str call_id: str interaction_id: str ai_session_id: str | None = None language: str | None = None media_uuid: str | None = None @dataclass(slots=True) class MediaActor: registration: MediaRegistration reader: asyncio.StreamReader writer: asyncio.StreamWriter vad: EnergyVAD frame_ms: int frame_bytes: int turn_queue: asyncio.Queue[tuple[bytes, bool]] = field(default_factory=asyncio.Queue) worker_task: asyncio.Task | None = None keepalive_task: asyncio.Task | None = None handoff_task: asyncio.Task | None = None playback_interrupt: asyncio.Event = field(default_factory=asyncio.Event) write_lock: asyncio.Lock = field(default_factory=asyncio.Lock) closed: bool = False state: str = "greeting" barge_in_pending: bool = False last_media_touch_monotonic: float = 0.0 last_outbound_audio_monotonic: float = 0.0 first_pcm_logged: bool = False keepalive_loop_logged: bool = False class AudioSocketMediaRuntime: def __init__( self, *, enabled: bool, host: str, port: int, frame_ms: int, idle_timeout_seconds: float, registration_wait_timeout_seconds: float, min_speech_ms: int, trailing_silence_ms: int, max_turn_ms: int, asr_provider: ASRProvider, tts_provider: TTSProvider, load_registration_by_media_uuid: Callable[[str], MediaRegistration | None], mark_media_connected: Callable[[str, str], None], mark_media_ended: Callable[[str, str], None], touch_media_frame: Callable[[str], None], set_state: Callable[[str, str, str | None, dict[str, Any] | None], None], get_pending_greeting: Callable[[str], str | None], mark_reply_delivered: Callable[[str, str, bool], None], process_turn: Callable[[str, str, str | None, bool, dict[str, Any] | None], VoiceAITurnDecisionOut], request_handoff: Callable[[str, str, VoiceAITurnDecisionOut], None], handle_media_error: Callable[[str, str, dict[str, Any] | None], None], ) -> None: self._enabled = enabled self._host = host self._port = port self._frame_ms = max(frame_ms, 20) self._frame_bytes = int((8000 * self._frame_ms / 1000.0) * 2) self._idle_timeout_seconds = max(idle_timeout_seconds, 5.0) self._registration_wait_timeout_seconds = max(registration_wait_timeout_seconds, 0.0) self._outbound_keepalive_interval_seconds = 0.75 self._min_speech_ms = max(min_speech_ms, self._frame_ms) self._trailing_silence_ms = max(trailing_silence_ms, self._frame_ms) self._max_turn_ms = max(max_turn_ms, self._frame_ms) self._asr_provider = asr_provider self._tts_provider = tts_provider self._load_registration_by_media_uuid = load_registration_by_media_uuid self._mark_media_connected = mark_media_connected self._mark_media_ended = mark_media_ended self._touch_media_frame = touch_media_frame self._set_state = set_state self._get_pending_greeting = get_pending_greeting self._mark_reply_delivered = mark_reply_delivered self._process_turn = process_turn self._request_handoff = request_handoff self._handle_media_error = handle_media_error self._server: asyncio.base_events.Server | None = None self._loop: asyncio.AbstractEventLoop | None = None self._actors: dict[str, MediaActor] = {} @property def address(self) -> str: return f"{self._host}:{self._port}" async def start(self) -> None: if not self._enabled or self._server is not None: return self._loop = asyncio.get_running_loop() self._server = await asyncio.start_server(self._handle_connection, self._host, self._port) async def _await_registration( self, media_uuid: str, *, writer: asyncio.StreamWriter | None = None, frame_bytes: int | None = None, ) -> MediaRegistration | None: deadline = time.monotonic() + self._registration_wait_timeout_seconds last_keepalive_at = 0.0 keepalive_payload = b"\x00" * max(frame_bytes or 0, 0) keepalive_logged = False while True: registration = await asyncio.to_thread(self._load_registration_by_media_uuid, media_uuid) if registration is not None: return registration if time.monotonic() >= deadline: logger.warning("audiosocket.registration_timeout media_uuid=%s", media_uuid) return None if writer is not None and keepalive_payload: now = time.monotonic() if now - last_keepalive_at >= self._outbound_keepalive_interval_seconds: writer.write(encode_audio_packet(keepalive_payload)) await writer.drain() if not keepalive_logged: keepalive_logged = True logger.warning( "audiosocket.prereg_keepalive media_uuid=%s frame_bytes=%s", media_uuid, len(keepalive_payload), ) last_keepalive_at = now await asyncio.sleep(0.05) async def stop(self) -> None: actor_ids = list(self._actors.keys()) for session_id in actor_ids: await self._close_session(session_id, reason="runtime_shutdown") if self._server is not None: self._server.close() await self._server.wait_closed() self._server = None def close_session_sync(self, session_id: str, *, reason: str) -> None: if not self._loop: return future = asyncio.run_coroutine_threadsafe( self._close_session(session_id, reason=reason), self._loop, ) with contextlib.suppress(Exception): future.result(timeout=5.0) async def _handle_connection( self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, ) -> None: actor: MediaActor | None = None close_reason = "connection_closed" close_as_error = False peer = writer.get_extra_info("peername") try: logger.warning("audiosocket.accept peer=%s", peer) packet_type, payload = await read_packet(reader, timeout=self._idle_timeout_seconds) logger.warning( "audiosocket.initial_packet peer=%s packet_type=%s payload_len=%s", peer, packet_type, len(payload), ) if packet_type != AUDIO_SOCKET_PACKET_UUID: raise RuntimeError("AudioSocket UUID handshake is required") media_uuid = normalize_media_uuid(payload) logger.warning("audiosocket.handshake peer=%s media_uuid=%s", peer, media_uuid) registration = await self._await_registration( media_uuid, writer=writer, frame_bytes=self._frame_bytes, ) if registration is None: raise RuntimeError(f"Unknown AudioSocket media_uuid: {media_uuid}") logger.warning( "audiosocket.registered peer=%s session_id=%s media_uuid=%s", peer, registration.voice_session_id, media_uuid, ) actor = MediaActor( registration=registration, reader=reader, writer=writer, vad=EnergyVAD( frame_ms=self._frame_ms, min_speech_ms=self._min_speech_ms, trailing_silence_ms=self._trailing_silence_ms, max_turn_ms=self._max_turn_ms, ), frame_ms=self._frame_ms, frame_bytes=self._frame_bytes, ) previous_actor = self._actors.get(registration.voice_session_id) if previous_actor is not None: await self._cleanup_actor(previous_actor, reason="media_replaced", error=False) self._actors[registration.voice_session_id] = actor actor.last_outbound_audio_monotonic = time.monotonic() await asyncio.to_thread(self._mark_media_connected, registration.voice_session_id, media_uuid) actor.keepalive_task = asyncio.create_task(self._keepalive_loop(actor)) actor.worker_task = asyncio.create_task(self._worker(actor)) while not actor.closed: packet_type, payload = await read_packet(reader, timeout=self._idle_timeout_seconds) if packet_type == AUDIO_SOCKET_PACKET_PCM16: await self._handle_pcm(actor, payload) continue if packet_type == AUDIO_SOCKET_PACKET_DTMF: logger.warning( "audiosocket.dtmf session_id=%s payload_len=%s", actor.registration.voice_session_id, len(payload), ) continue if packet_type == AUDIO_SOCKET_PACKET_HANGUP: close_reason = "audiosocket_hangup" break if packet_type == AUDIO_SOCKET_PACKET_UUID: logger.warning( "audiosocket.extra_uuid session_id=%s payload_len=%s", actor.registration.voice_session_id, len(payload), ) continue raise RuntimeError(f"Unsupported AudioSocket packet type: {packet_type}") except asyncio.TimeoutError: close_reason = "media_idle_timeout" close_as_error = True logger.warning("audiosocket.timeout peer=%s", peer) except asyncio.IncompleteReadError: close_reason = "connection_closed" logger.warning("audiosocket.peer_closed peer=%s", peer) except Exception as exc: close_reason = str(exc)[:1000] or "media_runtime_error" close_as_error = True logger.warning("audiosocket.error peer=%s reason=%s", peer, close_reason) finally: if actor is not None: await self._cleanup_actor(actor, reason=close_reason, error=close_as_error) else: writer.close() with contextlib.suppress(Exception): await writer.wait_closed() async def _handle_pcm(self, actor: MediaActor, pcm_frame: bytes) -> None: if actor.closed or not pcm_frame: return if not actor.first_pcm_logged: actor.first_pcm_logged = True logger.warning( "audiosocket.first_pcm session_id=%s state=%s frame_bytes=%s", actor.registration.voice_session_id, actor.state, len(pcm_frame), ) now = time.monotonic() if now - actor.last_media_touch_monotonic >= 1.0: actor.last_media_touch_monotonic = now await asyncio.to_thread(self._touch_media_frame, actor.registration.voice_session_id) if actor.state not in {"listening", "speaking"}: return vad_result = actor.vad.feed(pcm_frame) if actor.state == "speaking" and vad_result.speech_started: actor.playback_interrupt.set() actor.barge_in_pending = True if vad_result.utterance_pcm: await actor.turn_queue.put((vad_result.utterance_pcm, actor.barge_in_pending)) actor.barge_in_pending = False async def _worker(self, actor: MediaActor) -> None: greeting_text = await asyncio.to_thread( self._get_pending_greeting, actor.registration.voice_session_id, ) if greeting_text: logger.warning( "audiosocket.greeting session_id=%s text_len=%s", actor.registration.voice_session_id, len(greeting_text), ) await self._speak_text(actor, greeting_text, is_greeting=True) if not actor.closed: await self._set_actor_state(actor, "listening") while not actor.closed: pcm_bytes, barge_in = await actor.turn_queue.get() if actor.closed: break await self._process_utterance(actor, pcm_bytes, barge_in) async def _process_utterance(self, actor: MediaActor, pcm_bytes: bytes, barge_in: bool) -> None: await self._set_actor_state(actor, "thinking") wav_bytes = pcm16le_to_wav_bytes(pcm_bytes, sample_rate_hz=8000) transcription = await asyncio.to_thread( lambda: self._asr_provider.transcribe( wav_bytes, language_hint=actor.registration.language, ) ) transcript_text = str(transcription.text or "").strip() if not transcript_text: await self._set_actor_state(actor, "listening") return metadata = { "turn_duration_ms": int(len(pcm_bytes) / 16), "media_uuid": actor.registration.media_uuid, } decision = await asyncio.to_thread( self._process_turn, actor.registration.voice_session_id, transcript_text, transcription.language or actor.registration.language, barge_in, metadata, ) handoff_task: asyncio.Task | None = None if decision.needs_handoff: await self._set_actor_state(actor, "handoff_requested", decision.handoff_reason) handoff_task = self._start_handoff_request(actor, transcript_text, decision) if decision.reply_text: await self._speak_text(actor, decision.reply_text, is_greeting=False) if actor.closed: return if decision.needs_handoff: if handoff_task is not None and handoff_task.done(): with contextlib.suppress(asyncio.CancelledError, Exception): await handoff_task return await self._set_actor_state(actor, "listening") def _start_handoff_request( self, actor: MediaActor, customer_request_text: str, decision: VoiceAITurnDecisionOut, ) -> asyncio.Task: existing = actor.handoff_task if existing is not None and not existing.done(): return existing async def _run() -> None: try: await asyncio.to_thread( self._request_handoff, actor.registration.voice_session_id, customer_request_text, decision, ) except Exception as exc: logger.warning( "audiosocket.handoff_request_failed session_id=%s error=%s", actor.registration.voice_session_id, str(exc)[:500], ) actor.handoff_task = asyncio.create_task(_run()) return actor.handoff_task async def _speak_text(self, actor: MediaActor, text: str, *, is_greeting: bool) -> None: if actor.closed or not text: return await self._set_actor_state(actor, "speaking") synth_started_at = time.monotonic() synthesis = await asyncio.to_thread( lambda: self._tts_provider.synthesize( text, language=actor.registration.language, ) ) if not synthesis.audio_bytes: raise RuntimeError("TTS provider returned empty audio") logger.warning( "audiosocket.tts_ready session_id=%s greeting=%s synth_ms=%s audio_bytes=%s", actor.registration.voice_session_id, is_greeting, int((time.monotonic() - synth_started_at) * 1000), len(synthesis.audio_bytes), ) pcm_8k = resample_pcm16le( synthesis.audio_bytes, input_rate_hz=synthesis.sample_rate_hz, output_rate_hz=8000, ) first_frame_sent = False for frame in chunk_audio(pcm_8k, frame_bytes=actor.frame_bytes): if actor.closed or actor.playback_interrupt.is_set(): break await self._write_audio_packet(actor, frame) if not first_frame_sent: first_frame_sent = True logger.warning( "audiosocket.first_frame session_id=%s greeting=%s frame_bytes=%s", actor.registration.voice_session_id, is_greeting, len(frame), ) await asyncio.sleep(actor.frame_ms / 1000.0) interrupted = actor.playback_interrupt.is_set() actor.playback_interrupt.clear() if not interrupted and not actor.closed: await asyncio.to_thread( self._mark_reply_delivered, actor.registration.voice_session_id, text, is_greeting, ) logger.warning( "audiosocket.reply_delivered session_id=%s greeting=%s", actor.registration.voice_session_id, is_greeting, ) async def _write_audio_packet(self, actor: MediaActor, pcm_frame: bytes) -> None: if actor.closed: return async with actor.write_lock: if actor.closed: return actor.writer.write(encode_audio_packet(pcm_frame)) await actor.writer.drain() actor.last_outbound_audio_monotonic = time.monotonic() async def _keepalive_loop(self, actor: MediaActor) -> None: silence_frame = b"\x00" * actor.frame_bytes while not actor.closed: await asyncio.sleep(0.25) if actor.closed: break if (time.monotonic() - actor.last_outbound_audio_monotonic) < self._outbound_keepalive_interval_seconds: continue if not actor.keepalive_loop_logged: actor.keepalive_loop_logged = True logger.warning( "audiosocket.keepalive_loop session_id=%s state=%s frame_bytes=%s", actor.registration.voice_session_id, actor.state, actor.frame_bytes, ) await self._write_audio_packet(actor, silence_frame) async def _set_actor_state( self, actor: MediaActor, state: str, handoff_reason: str | None = None, ) -> None: if actor.closed and state not in {"error", "closed"}: return if actor.state == state and (handoff_reason or None) is None: return logger.warning( "audiosocket.state session_id=%s from_state=%s to_state=%s handoff_reason=%s", actor.registration.voice_session_id, actor.state, state, handoff_reason, ) actor.state = state await asyncio.to_thread( self._set_state, actor.registration.voice_session_id, state, handoff_reason, {"media_uuid": actor.registration.media_uuid}, ) async def _close_session(self, session_id: str, *, reason: str) -> None: actor = self._actors.get(session_id) if actor is None: return await self._cleanup_actor(actor, reason=reason, error=False) async def _cleanup_actor(self, actor: MediaActor, *, reason: str, error: bool) -> None: if actor.closed: return actor.closed = True logger.warning( "audiosocket.cleanup session_id=%s reason=%s error=%s first_pcm=%s", actor.registration.voice_session_id, reason, error, actor.first_pcm_logged, ) self._actors.pop(actor.registration.voice_session_id, None) if actor.keepalive_task is not None: actor.keepalive_task.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): await actor.keepalive_task if actor.handoff_task is not None: actor.handoff_task.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): await actor.handoff_task if actor.worker_task is not None: actor.worker_task.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): await actor.worker_task actor.writer.close() with contextlib.suppress(Exception): await actor.writer.wait_closed() if error: await asyncio.to_thread( self._handle_media_error, actor.registration.voice_session_id, reason, {"media_uuid": actor.registration.media_uuid}, ) return await asyncio.to_thread(self._mark_media_ended, actor.registration.voice_session_id, reason)