realtime_voice_service/transports/audiosocket.py

from __future__ import annotations

import asyncio
import logging
import struct
import uuid
from typing import Awaitable, Callable

from realtime_voice_service.transports.base import BaseMediaTransport


LOGGER = logging.getLogger("uvicorn.error")

AUDIO_SOCKET_PACKET_HANGUP = 0x00
AUDIO_SOCKET_PACKET_UUID = 0x01
AUDIO_SOCKET_PACKET_DTMF = 0x03
AUDIO_SOCKET_PACKET_PCM16 = 0x10


def normalize_session_id(value: str | bytes) -> str:
    if isinstance(value, bytes):
        try:
            return str(uuid.UUID(bytes=value)).lower()
        except (ValueError, AttributeError):
            return str(uuid.UUID(value.decode("utf-8").strip())).lower()
    return str(uuid.UUID(str(value).strip())).lower()


def encode_packet(packet_type: int, payload: bytes = b"") -> bytes:
    return struct.pack("!BH", packet_type & 0xFF, len(payload)) + payload


def encode_audio_packet(pcm_bytes: bytes) -> bytes:
    return encode_packet(AUDIO_SOCKET_PACKET_PCM16, pcm_bytes)


async def read_packet(
    reader: asyncio.StreamReader,
    *,
    timeout_seconds: float = 5.0,
) -> tuple[int, bytes]:
    header = await asyncio.wait_for(reader.readexactly(3), timeout=timeout_seconds)
    packet_type, payload_length = struct.unpack("!BH", header)
    payload = b""
    if payload_length:
        payload = await asyncio.wait_for(reader.readexactly(payload_length), timeout=timeout_seconds)
    return packet_type, payload


class AudioSocketTransport(BaseMediaTransport):
    def __init__(
        self,
        *,
        transport_id: str,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
        sample_rate_hz: int = 8000,
        frame_duration_ms: int = 20,
        read_timeout_seconds: float = 30.0,
    ) -> None:
        super().__init__(
            transport_id=transport_id,
            sample_rate_hz=sample_rate_hz,
            frame_duration_ms=frame_duration_ms,
        )
        self._reader = reader
        self._writer = writer
        self._read_timeout_seconds = max(read_timeout_seconds, 1.0)
        self._closed = False

    @property
    def protocol(self) -> str:
        return "audiosocket"

    @classmethod
    async def from_streams(
        cls,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
        *,
        handshake_timeout_seconds: float = 5.0,
        sample_rate_hz: int = 8000,
        frame_duration_ms: int = 20,
    ) -> AudioSocketTransport:
        packet_type, payload = await read_packet(reader, timeout_seconds=handshake_timeout_seconds)
        if packet_type != AUDIO_SOCKET_PACKET_UUID:
            raise ValueError(f"expected UUID packet, got {packet_type}")
        session_id = normalize_session_id(payload)
        return cls(
            transport_id=session_id,
            reader=reader,
            writer=writer,
            sample_rate_hz=sample_rate_hz,
            frame_duration_ms=frame_duration_ms,
        )

    async def receive_audio(self) -> bytes | None:
        if self._closed:
            return None
        while not self._closed:
            try:
                packet_type, payload = await read_packet(
                    self._reader,
                    timeout_seconds=self._read_timeout_seconds,
                )
            except (asyncio.IncompleteReadError, asyncio.TimeoutError, ConnectionError):
                return None

            if packet_type == AUDIO_SOCKET_PACKET_HANGUP:
                return None
            if packet_type == AUDIO_SOCKET_PACKET_PCM16:
                return payload
            if packet_type in {AUDIO_SOCKET_PACKET_UUID, AUDIO_SOCKET_PACKET_DTMF}:
                continue
        return None

    async def send_audio(self, audio_chunk: bytes) -> None:
        if self._closed or self._writer.is_closing():
            return
        self._writer.write(encode_audio_packet(audio_chunk))
        await self._writer.drain()

    async def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        if not self._writer.is_closing():
            self._writer.close()
            await self._writer.wait_closed()


class AudioSocketServer:
    def __init__(
        self,
        *,
        host: str,
        port: int,
        session_handler: Callable[[AudioSocketTransport], Awaitable[None]],
        handshake_timeout_seconds: float = 5.0,
        sample_rate_hz: int = 8000,
        frame_duration_ms: int = 20,
    ) -> None:
        self._host = host
        self._port = port
        self._session_handler = session_handler
        self._handshake_timeout_seconds = max(handshake_timeout_seconds, 1.0)
        self._sample_rate_hz = sample_rate_hz
        self._frame_duration_ms = frame_duration_ms
        self._server: asyncio.base_events.Server | None = None
        self._connection_tasks: set[asyncio.Task[None]] = set()

    @property
    def bound_port(self) -> int:
        if self._server is None or not self._server.sockets:
            return self._port
        return int(self._server.sockets[0].getsockname()[1])

    async def start(self) -> None:
        if self._server is not None:
            return
        self._server = await asyncio.start_server(
            self._handle_connection,
            self._host,
            self._port,
        )
        LOGGER.info(
            "AudioSocket server listening on %s:%s",
            self._host,
            self.bound_port,
        )

    async def stop(self) -> None:
        if self._server is None:
            return
        self._server.close()
        await self._server.wait_closed()
        self._server = None

        tasks = list(self._connection_tasks)
        for task in tasks:
            task.cancel()
        for task in tasks:
            try:
                await task
            except asyncio.CancelledError:
                pass

    async def _handle_connection(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        current_task = asyncio.current_task()
        if current_task is not None:
            self._connection_tasks.add(current_task)

        transport: AudioSocketTransport | None = None
        peer = writer.get_extra_info("peername")
        try:
            transport = await AudioSocketTransport.from_streams(
                reader,
                writer,
                handshake_timeout_seconds=self._handshake_timeout_seconds,
                sample_rate_hz=self._sample_rate_hz,
                frame_duration_ms=self._frame_duration_ms,
            )
            LOGGER.info(
                "AudioSocket client accepted: session=%s peer=%s",
                transport.transport_id,
                peer,
            )
            await self._session_handler(transport)
        except asyncio.CancelledError:
            raise
        except Exception:
            LOGGER.exception("AudioSocket connection failed from peer=%s", peer)
        finally:
            if transport is not None:
                await transport.close()
            elif not writer.is_closing():
                writer.close()
                await writer.wait_closed()
            if current_task is not None:
                self._connection_tasks.discard(current_task)