desklock/gateway/src/desklock_gateway/main.py

"""DeskLock gateway: WebSocket voice bridge between the device and Tatlock.

Protocol (see docs/architecture.md — keep in sync):
  device -> gateway : {"type": "utterance_start"} + binary PCM + {"type": "utterance_end"}
  gateway -> device : state / transcript / reply_text events, then audio_start + PCM + audio_end
"""

import asyncio
import logging
from datetime import UTC, datetime

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

from . import commands, stt, tts
from .config import settings
from .tatlock import TatlockClient

logger = logging.getLogger("desklock.gateway")

app = FastAPI(title="DeskLock Gateway")

PCM_CHUNK_BYTES = 4096

# device registry: every endpoint that ever connected, keyed by client IP
DEVICES: dict[str, dict] = {}

# Butler filler audio (the "let me check…" line) is synthesized once, then replayed.
_filler_pcm: bytes | None = None
_filler_tried = False


async def _ensure_filler() -> bytes | None:
    global _filler_pcm, _filler_tried
    if not _filler_tried:
        _filler_tried = True
        try:
            _filler_pcm = await asyncio.to_thread(tts.synthesize, settings.filler_text)
        except Exception:
            logger.exception("filler synth failed; continuing without it")
    return _filler_pcm


def _now() -> str:
    return datetime.now(UTC).astimezone().isoformat(timespec="seconds")


@app.get("/healthz")
async def healthz() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/devices")
async def devices() -> dict[str, dict]:
    return DEVICES


@app.websocket("/ws/voice")
async def voice(ws: WebSocket) -> None:
    await ws.accept()
    ip = ws.client.host if ws.client else "unknown"
    device = DEVICES.setdefault(ip, {"connections": 0})
    device["connections"] += 1
    device.update(connected=True, connected_at=_now(), last_seen=_now())
    logger.info("device connected from %s (connection #%d)", ip, device["connections"])

    tatlock = TatlockClient()
    pcm = bytearray()
    recording = False
    try:
        while True:
            message = await ws.receive()
            if message["type"] == "websocket.disconnect":
                raise WebSocketDisconnect(message.get("code", 1000))
            device["last_seen"] = _now()
            if "bytes" in message and message["bytes"] is not None:
                if recording:
                    pcm.extend(message["bytes"])
                continue
            if "text" not in message or message["text"] is None:
                continue

            import json

            event = json.loads(message["text"])
            if event["type"] == "hello":
                device["device"] = event.get("device", "unknown")
                device["fw"] = event.get("fw", "unknown")
                logger.info("device identified: %s fw=%s", device["device"], device["fw"])
            elif event["type"] == "utterance_start":
                pcm.clear()
                recording = True
            elif event["type"] == "utterance_end":
                recording = False
                try:
                    await _handle_utterance(ws, tatlock, bytes(pcm))
                except WebSocketDisconnect:
                    raise
                except Exception:
                    logger.exception("utterance failed")
                    try:
                        await ws.send_json({"type": "error"})
                        await ws.send_json({"type": "state", "value": "idle"})
                    except Exception:
                        logger.info("device gone before error could be reported")
                        break
    except WebSocketDisconnect:
        logger.info("device disconnected: %s", ip)
    finally:
        device.update(connected=False, disconnected_at=_now())
        await tatlock.aclose()


async def _handle_utterance(ws: WebSocket, tatlock: TatlockClient, pcm: bytes) -> None:
    await ws.send_json({"type": "state", "value": "thinking"})

    transcript = await asyncio.to_thread(stt.transcribe, pcm, settings.sample_rate)
    await ws.send_json({"type": "transcript", "text": transcript})
    if not transcript:
        await ws.send_json({"type": "state", "value": "idle"})
        return

    # Simple device commands (volume/mute) bypass the LLM entirely.
    command = commands.match(transcript)
    if command is not None:
        await ws.send_json(command)
        await ws.send_json({"type": "state", "value": "idle"})
        return

    # Acknowledge immediately with a canned line so the long Tatlock wait isn't dead
    # air, then re-assert "thinking" to keep the spinner up while it works.
    filler = await _ensure_filler()
    if filler:
        await ws.send_json({"type": "reply_text", "text": settings.filler_text})
        await ws.send_json({"type": "audio_start", "sample_rate": settings.sample_rate})
        for offset in range(0, len(filler), PCM_CHUNK_BYTES):
            await ws.send_bytes(filler[offset : offset + PCM_CHUNK_BYTES])
        await ws.send_json({"type": "audio_end"})
        await ws.send_json({"type": "state", "value": "thinking"})

    reply = await tatlock.ask(transcript)
    await ws.send_json({"type": "reply_text", "text": reply})

    audio = await asyncio.to_thread(tts.synthesize, reply)
    await ws.send_json({"type": "audio_start", "sample_rate": settings.sample_rate})
    for offset in range(0, len(audio), PCM_CHUNK_BYTES):
        await ws.send_bytes(audio[offset : offset + PCM_CHUNK_BYTES])
    await ws.send_json({"type": "audio_end"})
    await ws.send_json({"type": "state", "value": "idle"})