scripts/realtime_ws_smoke.py

from __future__ import annotations

import asyncio
import json
import os
import socket
import sys
import threading
import time
import urllib.request
from pathlib import Path
from typing import Any

import uvicorn
import websockets


ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))


def main() -> None:
    os.environ.setdefault("TIME_MACHINE_ADAPTER_PROFILE", "fixture")
    import app as app_module

    launch_events = list(
        app_module.container.encounter_service.start_encounter(
            mode="surprise",
            coordinate_prompt=None,
            session_id="realtime-smoke",
            seed=0,
        )
    )
    encounter_id = str(launch_events[0].metadata["encounter_id"])
    host = "127.0.0.1"
    port = _free_port()

    config = uvicorn.Config(
        app_module.app,
        host=host,
        port=port,
        log_level="warning",
        lifespan="off",
    )
    server = uvicorn.Server(config)
    thread = threading.Thread(target=server.run, daemon=True)
    thread.start()
    _wait_for_port(host, port)

    try:
        _check_http_routes(host, port)
        messages = asyncio.run(_run_ws_turn(host, port, encounter_id))
    finally:
        server.should_exit = True
        thread.join(timeout=10)

    seen = {message.get("type") for message in messages}
    required = {"ready", "listening", "processing", "transcript", "conversation_text", "audio"}
    missing = sorted(required - seen)
    if missing:
        raise SystemExit(f"Realtime WebSocket smoke failed; missing {missing}; saw {seen}")
    print("Realtime WebSocket smoke passed.")


def _check_http_routes(host: str, port: int) -> None:
    root_response = urllib.request.urlopen(f"http://{host}:{port}/", timeout=10)
    if root_response.status not in {200, 307, 308}:
        raise RuntimeError(f"Unexpected / status: {root_response.status}")
    app_response = urllib.request.urlopen(f"http://{host}:{port}/app", timeout=10)
    if app_response.status != 200:
        raise RuntimeError(f"Unexpected /app status: {app_response.status}")


async def _run_ws_turn(host: str, port: int, encounter_id: str) -> list[dict[str, Any]]:
    uri = f"ws://{host}:{port}/realtime/voice"
    messages: list[dict[str, Any]] = []
    async with websockets.connect(uri, max_size=8 * 1024 * 1024) as websocket:
        messages.append(json.loads(await websocket.recv()))
        await websocket.send(
            json.dumps(
                {
                    "type": "start",
                    "encounter_id": encounter_id,
                    "sample_rate": 16000,
                    "channels": 1,
                }
            )
        )
        messages.append(json.loads(await websocket.recv()))
        await websocket.send(b"\x00\x00" * 1600)
        await websocket.send(json.dumps({"type": "commit"}))
        for _ in range(20):
            message = json.loads(await websocket.recv())
            messages.append(message)
            if message.get("type") == "listening" and "audio" in {item.get("type") for item in messages}:
                break
    return messages


def _free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind(("127.0.0.1", 0))
        return int(sock.getsockname()[1])


def _wait_for_port(host: str, port: int) -> None:
    deadline = time.monotonic() + 20
    while time.monotonic() < deadline:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.settimeout(0.2)
            if sock.connect_ex((host, port)) == 0:
                return
    raise RuntimeError(f"Server did not open {host}:{port}")


if __name__ == "__main__":
    main()