Plan 3: backend WS client + dispatch + tool round-trip

client/session.py

@@ -0,0 +1,211 @@
+"""Backend WebSocket client (asyncio)."""
+import asyncio
+import json
+import random
+from typing import Callable, Protocol
+
+from client.log import get_logger
+
+_log = get_logger("session")
+
+CLIENT_VERSION = "0.1.0"
+PING_INTERVAL_S = 15.0
+RECONNECT_MIN_S = 1.0
+RECONNECT_MAX_S = 30.0
+
+
+class Callbacks(Protocol):
+    def on_session_started(self, conversation_id: str) -> None: ...
+    def on_assistant_audio(self, pcm16: bytes) -> None: ...
+    def on_assistant_done(self) -> None: ...
+    def on_session_ended(self, reason: str) -> None: ...
+
+
+class Session:
+    def __init__(
+        self,
+        backend_ws: str,
+        device_token: str,
+        device_id: str,
+        callbacks: Callbacks,
+    ) -> None:
+        self._backend_ws = backend_ws
+        self._device_token = device_token
+        self._device_id = device_id
+        self._cb = callbacks
+        self._tools: dict[str, Callable[[dict], dict]] = {}
+        self._ws = None
+        self._loop: asyncio.AbstractEventLoop | None = None
+
+    def register_tool(self, name: str, handler: Callable[[dict], dict]) -> None:
+        self._tools[name] = handler
+
+    async def run_forever(self) -> None:
+        """Connect, pump, reconnect with exponential backoff. Never returns
+        cleanly except on cancellation."""
+        # Lazy import so the module is importable without `websockets`.
+        import websockets
+
+        self._loop = asyncio.get_running_loop()
+        backoff = RECONNECT_MIN_S
+        while True:
+            try:
+                _log.info("connecting to %s", self._backend_ws)
+                async with websockets.connect(
+                    self._backend_ws,
+                    additional_headers={"Authorization": f"Bearer {self._device_token}"},
+                    max_size=4 * 1024 * 1024,
+                ) as ws:
+                    self._ws = ws
+                    backoff = RECONNECT_MIN_S
+                    ping_task = asyncio.create_task(self._ping_loop(ws))
+                    try:
+                        await self._pump(ws)
+                    finally:
+                        ping_task.cancel()
+                        self._ws = None
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:  # noqa: BLE001
+                _log.warning("session error: %s; reconnecting in %.1fs", exc, backoff)
+            jitter = random.uniform(0, backoff * 0.1)
+            await asyncio.sleep(backoff + jitter)
+            backoff = min(backoff * 2.0, RECONNECT_MAX_S)
+
+    async def _pump(self, ws) -> None:
+        # 1. Send hello.
+        await ws.send(json.dumps({
+            "type": "hello",
+            "device_id": self._device_id,
+            "client_version": CLIENT_VERSION,
+        }))
+        # Track ws so the threadsafe shims work when _pump is exercised directly
+        # (i.e. without going through run_forever, as in unit tests).
+        self._ws = ws
+        try:
+            self._loop = asyncio.get_running_loop()
+        except RuntimeError:
+            pass
+        # 2. Pump.
+        while True:
+            frame = await ws.recv()
+            if isinstance(frame, (bytes, bytearray, memoryview)):
+                self._cb.on_assistant_audio(bytes(frame))
+                continue
+            try:
+                env = json.loads(frame)
+            except json.JSONDecodeError:
+                _log.warning("non-JSON text frame; dropping")
+                continue
+            await self._dispatch(env, ws)
+
+    async def _dispatch(self, env: dict, ws) -> None:
+        t = env.get("type")
+        if t == "hello_ack":
+            _log.info("hello_ack config=%s", env.get("config"))
+        elif t == "pong":
+            pass
+        elif t == "session_started":
+            self._cb.on_session_started(str(env.get("conversation_id", "")))
+        elif t == "assistant_done":
+            self._cb.on_assistant_done()
+        elif t == "session_ended":
+            self._cb.on_session_ended(str(env.get("reason", "unknown")))
+        elif t == "tool_call":
+            await self._handle_tool_call(env, ws)
+        elif t == "error":
+            _log.warning(
+                "backend error code=%s message=%s fatal=%s",
+                env.get("code"), env.get("message"), env.get("fatal"),
+            )
+        elif t == "config_updated":
+            _log.info("config_updated: %s", env.get("config"))
+        else:
+            _log.info("unhandled envelope type=%s", t)
+
+    async def _handle_tool_call(self, env: dict, ws) -> None:
+        call_id = env.get("call_id", "")
+        name = env.get("name", "")
+        args = env.get("arguments") or {}
+        if isinstance(args, str):
+            try:
+                args = json.loads(args)
+            except json.JSONDecodeError:
+                args = {}
+        handler = self._tools.get(name)
+        if handler is None:
+            await ws.send(json.dumps({
+                "type": "tool_result",
+                "call_id": call_id,
+                "ok": False,
+                "error": f"unknown tool: {name}",
+            }))
+            return
+        try:
+            result = handler(args)
+        except Exception as exc:  # noqa: BLE001
+            await ws.send(json.dumps({
+                "type": "tool_result",
+                "call_id": call_id,
+                "ok": False,
+                "error": str(exc),
+            }))
+            return
+        await ws.send(json.dumps({
+            "type": "tool_result",
+            "call_id": call_id,
+            "ok": True,
+            "result": result,
+        }))
+
+    async def _ping_loop(self, ws) -> None:
+        try:
+            while True:
+                await asyncio.sleep(PING_INTERVAL_S)
+                await ws.send(json.dumps({"type": "ping"}))
+        except asyncio.CancelledError:
+            pass
+
+    # ---- Public helpers used by main.py from the asyncio loop ----
+
+    async def send_wake(self) -> None:
+        if self._ws is None:
+            return
+        await self._ws.send(json.dumps({"type": "wake"}))
+
+    async def send_binary(self, pcm16: bytes) -> None:
+        if self._ws is None:
+            return
+        await self._ws.send(pcm16)
+
+    async def send_playback_done(self, label: str) -> None:
+        if self._ws is None:
+            return
+        await self._ws.send(json.dumps({"type": "playback_done", "for": label}))
+
+    async def send_cancel(self) -> None:
+        if self._ws is None:
+            return
+        await self._ws.send(json.dumps({"type": "cancel"}))
+
+    def call_soon_threadsafe_send_binary(self, pcm16: bytes) -> None:
+        """Thread-safe shim used by the bridge thread to queue an uplink send."""
+        if self._loop is None or self._ws is None:
+            return
+        self._loop.call_soon_threadsafe(
+            lambda: asyncio.create_task(self.send_binary(pcm16))
+        )
+
+    def call_soon_threadsafe_send_wake(self) -> None:
+        if self._loop is None or self._ws is None:
+            return
+        self._loop.call_soon_threadsafe(
+            lambda: asyncio.create_task(self.send_wake())
+        )
+
+    def call_soon_threadsafe_send_playback_done(self, label: str) -> None:
+        if self._loop is None or self._ws is None:
+            return
+        self._loop.call_soon_threadsafe(
+            lambda: asyncio.create_task(self.send_playback_done(label))
+        )

client/tests/pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+asyncio_mode = auto

client/tests/test_session.py

@@ -0,0 +1,186 @@
+import asyncio
+import json
+
+import pytest
+
+from client.session import Session
+
+
+class FakeWebSocket:
+    """Minimal stand-in for a `websockets` client connection."""
+    def __init__(self, inbox: list):
+        # inbox is a list of either str (text frame) or bytes (binary frame).
+        self._inbox = list(inbox)
+        self.sent: list = []
+
+    async def send(self, frame):
+        self.sent.append(frame)
+
+    async def recv(self):
+        if not self._inbox:
+            await asyncio.sleep(3600)  # block forever; tests cancel
+        return self._inbox.pop(0)
+
+    async def close(self):
+        pass
+
+
+class _FakeCallbacks:
+    """Implements client.session.Callbacks via duck-typing."""
+    def __init__(self):
+        self.started: list = []
+        self.audio: list = []
+        self.assistant_done = 0
+        self.session_ended: list = []
+
+    def on_session_started(self, conversation_id: str) -> None:
+        self.started.append(conversation_id)
+
+    def on_assistant_audio(self, pcm16: bytes) -> None:
+        self.audio.append(pcm16)
+
+    def on_assistant_done(self) -> None:
+        self.assistant_done += 1
+
+    def on_session_ended(self, reason: str) -> None:
+        self.session_ended.append(reason)
+
+
+async def _run_pump_briefly(session, ws, settle: float = 0.05):
+    task = asyncio.create_task(session._pump(ws))
+    await asyncio.sleep(settle)
+    task.cancel()
+    try:
+        await task
+    except (asyncio.CancelledError, BaseException):
+        pass
+
+
+async def test_hello_is_sent_on_connect_and_hello_ack_logged():
+    ws = FakeWebSocket(inbox=[
+        json.dumps({"type": "hello_ack", "config": {"voice": "alloy"}}),
+    ])
+    session = Session(
+        backend_ws="wss://x/device",
+        device_token="T",
+        device_id="00000000-0000-0000-0000-000000000001",
+        callbacks=_FakeCallbacks(),
+    )
+    await _run_pump_briefly(session, ws)
+
+    first = json.loads(ws.sent[0])
+    assert first["type"] == "hello"
+    assert first["device_id"] == "00000000-0000-0000-0000-000000000001"
+    assert first["client_version"]
+
+
+async def test_session_started_calls_callback_with_conversation_id():
+    cb = _FakeCallbacks()
+    ws = FakeWebSocket(inbox=[
+        json.dumps({"type": "hello_ack", "config": {}}),
+        json.dumps({
+            "type": "session_started",
+            "conversation_id": "cccccccc-cccc-cccc-cccc-cccccccccccc",
+        }),
+    ])
+    session = Session("wss://x/device", "T", "did", cb)
+    await _run_pump_briefly(session, ws)
+    assert cb.started == ["cccccccc-cccc-cccc-cccc-cccccccccccc"]
+
+
+async def test_binary_frame_goes_to_audio_callback():
+    cb = _FakeCallbacks()
+    payload = bytes(range(256)) * 15  # 3840 bytes
+    ws = FakeWebSocket(inbox=[
+        json.dumps({"type": "hello_ack", "config": {}}),
+        payload,
+    ])
+    session = Session("wss://x/device", "T", "did", cb)
+    await _run_pump_briefly(session, ws)
+    assert cb.audio == [payload]
+
+
+async def test_assistant_done_and_session_ended_callbacks():
+    cb = _FakeCallbacks()
+    ws = FakeWebSocket(inbox=[
+        json.dumps({"type": "hello_ack", "config": {}}),
+        json.dumps({"type": "assistant_done"}),
+        json.dumps({"type": "session_ended", "reason": "idle"}),
+    ])
+    session = Session("wss://x/device", "T", "did", cb)
+    await _run_pump_briefly(session, ws)
+    assert cb.assistant_done == 1
+    assert cb.session_ended == ["idle"]
+
+
+async def test_tool_call_runs_handler_and_sends_tool_result():
+    cb = _FakeCallbacks()
+    ws = FakeWebSocket(inbox=[
+        json.dumps({"type": "hello_ack", "config": {}}),
+        json.dumps({
+            "type": "tool_call",
+            "call_id": "c1",
+            "name": "set_volume",
+            "arguments": {"level": 42},
+        }),
+    ])
+    session = Session("wss://x/device", "T", "did", cb)
+    captured = {}
+
+    def handler(args):
+        captured["args"] = args
+        return {"ok": True, "level": args["level"]}
+
+    session.register_tool("set_volume", handler)
+    await _run_pump_briefly(session, ws)
+
+    assert captured["args"] == {"level": 42}
+    reply = json.loads(ws.sent[-1])
+    assert reply["type"] == "tool_result"
+    assert reply["call_id"] == "c1"
+    assert reply["ok"] is True
+    assert reply["result"] == {"ok": True, "level": 42}
+
+
+async def test_tool_call_with_exception_sends_ok_false():
+    cb = _FakeCallbacks()
+    ws = FakeWebSocket(inbox=[
+        json.dumps({"type": "hello_ack", "config": {}}),
+        json.dumps({
+            "type": "tool_call",
+            "call_id": "c2",
+            "name": "broken",
+            "arguments": {},
+        }),
+    ])
+    session = Session("wss://x/device", "T", "did", cb)
+
+    def handler(args):
+        raise RuntimeError("kaboom")
+
+    session.register_tool("broken", handler)
+    await _run_pump_briefly(session, ws)
+
+    reply = json.loads(ws.sent[-1])
+    assert reply["type"] == "tool_result"
+    assert reply["call_id"] == "c2"
+    assert reply["ok"] is False
+    assert "kaboom" in reply["error"]
+
+
+async def test_unknown_tool_sends_ok_false():
+    cb = _FakeCallbacks()
+    ws = FakeWebSocket(inbox=[
+        json.dumps({"type": "hello_ack", "config": {}}),
+        json.dumps({
+            "type": "tool_call",
+            "call_id": "c3",
+            "name": "ghost",
+            "arguments": {},
+        }),
+    ])
+    session = Session("wss://x/device", "T", "did", cb)
+    await _run_pump_briefly(session, ws)
+    reply = json.loads(ws.sent[-1])
+    assert reply["ok"] is False
+    assert "ghost" in reply["error"]