Plan 3: backend WS client + dispatch + tool round-trip
This commit is contained in:
@@ -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))
|
||||
)
|
||||
@@ -0,0 +1,2 @@
|
||||
[pytest]
|
||||
asyncio_mode = auto
|
||||
@@ -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"]
|
||||
Reference in New Issue
Block a user