From b1eeaa8af32fbe5e1a66a76c5b6a7f5e5742341e Mon Sep 17 00:00:00 2001 From: Assistant builder Date: Thu, 11 Jun 2026 22:51:23 +0000 Subject: [PATCH] Plan 3: backend WS client + dispatch + tool round-trip --- client/session.py | 211 +++++++++++++++++++++++++++++++++++ client/tests/pytest.ini | 2 + client/tests/test_session.py | 186 ++++++++++++++++++++++++++++++ 3 files changed, 399 insertions(+) create mode 100644 client/session.py create mode 100644 client/tests/pytest.ini create mode 100644 client/tests/test_session.py diff --git a/client/session.py b/client/session.py new file mode 100644 index 0000000..2f1e9fc --- /dev/null +++ b/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)) + ) diff --git a/client/tests/pytest.ini b/client/tests/pytest.ini new file mode 100644 index 0000000..2f4c80e --- /dev/null +++ b/client/tests/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +asyncio_mode = auto diff --git a/client/tests/test_session.py b/client/tests/test_session.py new file mode 100644 index 0000000..16855e4 --- /dev/null +++ b/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"]