From b88cb34196613292e56e8582dcc0a09facc2ec84 Mon Sep 17 00:00:00 2001 From: Adolfo Reyna Date: Fri, 19 Jun 2026 13:58:42 -0400 Subject: [PATCH] Implement WebSocket real-time audio streaming for Hermes gateway --- demo_websocket_voice.py | 268 ++++++++++++++++++++ hermes_voice_gateway/api_server_endpoint.py | 253 ++++++++++++++++++ hermes_voice_gateway/test_websocket.py | 230 +++++++++++++++++ lib/websocket_client.py | 184 ++++++++++++++ 4 files changed, 935 insertions(+) create mode 100644 demo_websocket_voice.py create mode 100644 hermes_voice_gateway/test_websocket.py create mode 100644 lib/websocket_client.py diff --git a/demo_websocket_voice.py b/demo_websocket_voice.py new file mode 100644 index 0000000..941ae3f --- /dev/null +++ b/demo_websocket_voice.py @@ -0,0 +1,268 @@ +# pyright: reportMissingImports=false, reportAttributeAccessIssue=false +"""Client firmware application for real-time WebSocket audio streaming to Hermes.""" + +import time +import struct +import machine +from machine import Pin, SPI, I2C, I2S +import json +import ili9341 +from ft6336u import FT6336U +from audio_util import ES8311 +from websocket_client import WebSocketClient + +# --- CONFIGURATION --- +HERMES_WS_URL = "ws://192.168.68.126:8642/api/esp32/voice/ws" +DEVICE_ID = "kitchen-button" + +# Placeholder for API Key. Since the key is stored on the Hermes server, +# please copy-paste the token string here. +HERMES_API_KEY = "mcT1YA1vOr9wXSiHpCYalweEGGZKX-PIfZv2drp8BSg" + +def main(): + print("--- Starting Hermes WebSocket Voice Assistant Demo ---") + + # 1. Initialize display + spi = SPI(1, baudrate=40000000, polarity=0, phase=0, sck=Pin(12), mosi=Pin(11), miso=Pin(13)) + display = ili9341.ILI9341(spi, cs=Pin(10), dc=Pin(46), bl=Pin(45), rst=None) + display.clear(0) + display.text("Hermes Assistant", 10, 10, 1) + display.show() + + # 2. Initialize touch + i2c = I2C(0, sda=Pin(16), scl=Pin(15)) + touch = FT6336U(i2c, rst_pin=18, int_pin=17, width=320, height=240, swap_xy=True, invert_x=False, invert_y=True) + + # 3. Configure audio codec ES8311 + # We need MCLK pin (GPIO 4) active at 6.144 MHz + mclk_pin = Pin(4, Pin.OUT) + mclk_pwm = machine.PWM(mclk_pin) + mclk_pwm.freq(6144000) + mclk_pwm.duty_u16(32768) + + codec = ES8311(i2c) + if not codec.init(sample_rate=16000): + print("ES8311 init failed!") + return + + codec.set_volume(90) + + # Configure microphone registers on ES8311 + try: + codec._write(0x14, 0x1A) # Enable analog mic input & set PGA gain + codec._write(0x16, 0x01) # Boost MIC digital gain to +6dB + codec._write(0x17, 0xC8) # Set ADC digital volume + print("Microphone registers initialized.") + except Exception as e: + print("Failed to write microphone registers:", e) + + amp_pin = Pin(1, Pin.OUT, value=1) # start with amp disabled (1 = disabled) + buffer = bytearray(3200) # 100ms chunk at 16kHz mono 16-bit PCM (3200 bytes) + + while True: + display.clear(0) + display.text("Hermes Assistant", 10, 10, 1) + display.line(10, 22, 310, 22, 1) + display.text("Hold screen & ask", 50, 70, 1) + display.text("a question...", 50, 90, 1) + display.text("Status: Idle (WS)", 10, 220, 1) + display.show() + + # Wait for touch + while not touch.is_touched(): + time.sleep_ms(30) + + print("Touch detected! Connecting WebSocket...") + display.clear(0) + display.text("Connecting...", 80, 80, 1) + display.show() + + headers = { + "Authorization": f"Bearer {HERMES_API_KEY}", + "X-Device-ID": DEVICE_ID + } + + ws = WebSocketClient(HERMES_WS_URL, headers=headers) + try: + ws.connect() + + # Send start event + ws.send_text(json.dumps({ + "event": "start", + "device_id": DEVICE_ID, + "sample_rate": 16000, + "channels": 1, + "sample_width": 2, + "format": "pcm_s16le" + })) + + # Read ready and listening events from server + evt1_opcode, evt1_payload = ws.recv_frame() + evt2_opcode, evt2_payload = ws.recv_frame() + + print("WebSocket connected and streaming started.") + display.clear(0) + display.text("Hermes Assistant", 10, 10, 1) + display.line(10, 22, 310, 22, 1) + display.text("Listening...", 80, 80, 1) + display.fill_rect(130, 110, 30, 30, 1) + display.text("Status: Streaming", 10, 220, 1) + display.show() + + # 4. Open I2S RX for recording (Mono 16kHz) + i2s_rx = I2S(1, + sck=Pin(5), + ws=Pin(7), + sd=Pin(6), + mode=I2S.RX, + ibuf=8000, + rate=16000, + bits=16, + format=I2S.MONO) + + total_data_bytes = 0 + + try: + # Record loop - as long as touch is held + while touch.is_touched(): + bytes_read = i2s_rx.readinto(buffer) + if bytes_read > 0: + ws.send_binary(buffer[:bytes_read]) + total_data_bytes += bytes_read + + print(f"Touch released! Sent {total_data_bytes} bytes.") + except Exception as e: + print("Error recording/streaming:", e) + finally: + i2s_rx.deinit() + + if total_data_bytes < 3200: + print("Recording too short, cancelling session.") + ws.send_text(json.dumps({"event": "cancel"})) + ws.close() + continue + + # Send stop event + ws.send_text(json.dumps({"event": "stop"})) + + display.clear(0) + display.text("Hermes Assistant", 10, 10, 1) + display.line(10, 22, 310, 22, 1) + display.text("Processing...", 50, 80, 1) + display.text("Status: Thinking", 10, 220, 1) + display.show() + + # Playback/Events loop + i2s_tx = None + while True: + opcode, payload = ws.recv_frame() + if opcode is None: + break + + if opcode == 0x1: # Text frame (JSON event) + try: + event_data = json.loads(payload.decode('utf-8')) + evt = event_data.get("event") + + if evt == "transcript": + txt = event_data.get("text", "") + print(f"Heard: {txt}") + display.clear(0) + display.text("Hermes Assistant", 10, 10, 1) + display.line(10, 22, 310, 22, 1) + display.text("Heard:", 10, 40, 1) + # Wrap text onto lines + lines = [txt[i:i+30] for i in range(0, min(len(txt), 120), 30)] + y_offset = 60 + for line in lines: + display.text(line, 10, y_offset, 1) + y_offset += 20 + display.show() + + elif evt == "thinking": + pass + + elif evt == "response_text": + txt = event_data.get("text", "") + print(f"Response: {txt}") + display.clear(0) + display.text("Hermes Assistant", 10, 10, 1) + display.line(10, 22, 310, 22, 1) + display.text("Response:", 10, 40, 1) + lines = [txt[i:i+30] for i in range(0, min(len(txt), 120), 30)] + y_offset = 60 + for line in lines: + display.text(line, 10, y_offset, 1) + y_offset += 20 + display.show() + + elif evt == "audio_start": + print("Audio response started.") + # Enable amp and open I2S TX + amp_pin.value(0) # Active Low Enable + i2s_tx = I2S(1, + sck=Pin(5), + ws=Pin(7), + sd=Pin(8), + mode=I2S.TX, + ibuf=4096, + rate=16000, + bits=16, + format=I2S.MONO) + + elif evt == "audio_end": + print("Audio response ended.") + if i2s_tx: + time.sleep_ms(150) + amp_pin.value(1) # Disable amp + i2s_tx.deinit() + i2s_tx = None + + elif evt == "done": + break + + elif evt == "error": + msg = event_data.get("message", "Unknown error") + print(f"Error from server: {msg}") + display.clear(0) + display.text("Error", 10, 10, 1) + display.line(10, 22, 310, 22, 1) + display.text(msg[:100], 10, 80, 1) + display.show() + time.sleep(3) + break + except Exception as e: + print("Error parsing event text:", e) + + elif opcode == 0x2: # Binary frame (Audio WAV chunk) + if i2s_tx: + chunk = payload + if chunk.startswith(b'RIFF') and len(chunk) > 44: + chunk = chunk[44:] # Skip WAV header for direct I2S mono play + try: + i2s_tx.write(chunk) + except Exception as e: + print("Error writing to speaker:", e) + + ws.close() + except Exception as e: + print("Failed to stream to Hermes server:", e) + display.clear(0) + display.text("Server Error", 10, 10, 1) + display.line(10, 22, 310, 22, 1) + display.text("Could not connect", 50, 80, 1) + display.text("to WebSocket server.", 50, 100, 1) + display.show() + time.sleep(3) + try: + ws.close() + except Exception: + pass + + # Debounce touch release + while touch.is_touched(): + time.sleep_ms(30) + time.sleep_ms(300) + +if __name__ == "__main__": + main() diff --git a/hermes_voice_gateway/api_server_endpoint.py b/hermes_voice_gateway/api_server_endpoint.py index e8dfe84..16fd46c 100644 --- a/hermes_voice_gateway/api_server_endpoint.py +++ b/hermes_voice_gateway/api_server_endpoint.py @@ -652,7 +652,260 @@ class HermesESP32VoiceGatewayMixin: }, ) + async def _handle_esp32_voice_websocket(self, request: Any) -> Any: + """GET /api/esp32/voice/ws — real-time WebSocket audio streaming and status route.""" + if web is None: + raise RuntimeError("aiohttp is required for the ESP32 voice websocket endpoint") + + # 1. Check authentication first (HTTP headers) + auth_err = self._check_auth(request) + if auth_err: + return auth_err + + # Prepare websocket response + ws = web.WebSocketResponse() + await ws.prepare(request) + + device_id = (request.headers.get("X-Device-ID") or request.query.get("device_id") or "default").strip()[:80] or "default" + + # 2. Receive JSON start frame + try: + msg = await asyncio.wait_for(ws.receive(), timeout=10) + if msg.type != web.WSMsgType.TEXT: + await ws.send_json({"event": "error", "message": "Expected text start frame", "code": "protocol_error"}) + await ws.close() + return ws + + data = json.loads(msg.data) + if data.get("event") != "start": + await ws.send_json({"event": "error", "message": "Expected start event", "code": "protocol_error"}) + await ws.close() + return ws + + if "device_id" in data: + device_id = str(data["device_id"]).strip()[:80] or device_id + except asyncio.TimeoutError: + await ws.close() + return ws + except Exception as exc: + try: + await ws.send_json({"event": "error", "message": f"Invalid start payload: {exc}", "code": "protocol_error"}) + except Exception: + pass + await ws.close() + return ws + + # Send ready status + await ws.send_json({"event": "ready"}) + await ws.send_json({"event": "listening"}) + + try: + from hermes_constants import get_hermes_dir + audio_dir = get_hermes_dir("cache/audio", "audio_cache") / "esp32" + except Exception: + audio_dir = Path.home() / ".hermes" / "cache" / "audio" / "esp32" + audio_dir.mkdir(parents=True, exist_ok=True) + + safe_device = re.sub(r"[^A-Za-z0-9_.-]+", "_", device_id)[:80] or "default" + raw_path = audio_dir / f"ws_{safe_device}_{int(time.time() * 1000)}_{uuid.uuid4().hex[:8]}.s16le" + + total_bytes = 0 + start_time = time.time() + last_frame_time = time.time() + cancelled = False + timings = {"ws_accepted": start_time} + + try: + with open(raw_path, "wb") as f: + while True: + now = time.time() + elapsed = now - last_frame_time + if elapsed > 10.0: + await ws.send_json({"event": "error", "message": "Idle timeout exceeded", "code": "timeout"}) + break + + if now - start_time > 60.0: + await ws.send_json({"event": "error", "message": "Max session duration exceeded", "code": "duration_limit"}) + break + + timeout_left = min(10.0 - elapsed, 60.0 - (now - start_time)) + if timeout_left <= 0: + break + + try: + msg = await asyncio.wait_for(ws.receive(), timeout=timeout_left) + except asyncio.TimeoutError: + await ws.send_json({"event": "error", "message": "Idle timeout exceeded", "code": "timeout"}) + break + + last_frame_time = time.time() + + if msg.type == web.WSMsgType.BINARY: + if total_bytes + len(msg.data) > 8_000_000: + await ws.send_json({"event": "error", "message": "Max audio size exceeded", "code": "size_limit"}) + break + f.write(msg.data) + total_bytes += len(msg.data) + if "first_audio" not in timings: + timings["first_audio"] = last_frame_time + + elif msg.type == web.WSMsgType.TEXT: + try: + data = json.loads(msg.data) + event = data.get("event") + if event == "stop": + timings["stop_received"] = last_frame_time + break + elif event == "cancel": + cancelled = True + break + except Exception: + pass + + elif msg.type in (web.WSMsgType.CLOSE, web.WSMsgType.CLOSING): + break + elif msg.type == web.WSMsgType.ERROR: + break + except Exception as exc: + logger.exception("Error in websocket audio loop") + try: + await ws.send_json({"event": "error", "message": f"Internal server error: {exc}", "code": "server_error"}) + except Exception: + pass + raw_path.unlink(missing_ok=True) + await ws.close() + return ws + + if cancelled or ws.closed or total_bytes < 3200 or "stop_received" not in timings: + raw_path.unlink(missing_ok=True) + if not ws.closed: + if total_bytes < 3200: + await ws.send_json({"event": "error", "message": "Audio recording too short", "code": "audio_too_short"}) + else: + await ws.send_json({"event": "error", "message": "Session cancelled", "code": "cancelled"}) + await ws.close() + return ws + + wav_path = raw_path.with_suffix(".wav") + try: + with wave.open(str(wav_path), "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(16000) + wav.writeframes(raw_path.read_bytes()) + finally: + raw_path.unlink(missing_ok=True) + + timings["wav_finalized"] = time.time() + + try: + from tools.transcription_tools import transcribe_audio + stt_result = await asyncio.wait_for( + asyncio.to_thread(transcribe_audio, str(wav_path)), + timeout=45 + ) + timings["stt_done"] = time.time() + + if not stt_result.get("success"): + await ws.send_json({"event": "error", "message": stt_result.get("error") or "Transcription failed", "code": "stt_failed"}) + await ws.close() + return ws + + transcript = str(stt_result.get("transcript") or "").strip() + if not transcript: + await ws.send_json({"event": "error", "message": "Transcription empty", "code": "empty_transcript"}) + await ws.close() + return ws + + await ws.send_json({"event": "transcript", "text": transcript}) + await ws.send_json({"event": "thinking"}) + + response_data, final_text = await self._run_esp32_voice_turn( + transcript=transcript, + device_id=device_id, + instructions=request.headers.get("X-Hermes-Instructions") or None, + ) + timings["agent_done"] = time.time() + + await ws.send_json({"event": "response_text", "text": final_text}) + + from tools.tts_tool import text_to_speech_tool + tts_raw = await asyncio.to_thread(text_to_speech_tool, final_text) + tts_result = json.loads(tts_raw) + timings["tts_done"] = time.time() + + if not tts_result.get("success"): + await ws.send_json({"event": "error", "message": tts_result.get("error") or "TTS generation failed", "code": "tts_failed"}) + await ws.close() + return ws + + tts_path = str(tts_result.get("file_path") or "") + if not tts_path or not Path(tts_path).exists(): + await ws.send_json({"event": "error", "message": "TTS file missing", "code": "tts_file_missing"}) + await ws.close() + return ws + + response_audio_path = self._convert_audio_to_esp32_wav( + tts_path, + output_dir=audio_dir, + device_id=device_id + ) + + wav_bytes = response_audio_path.read_bytes() + await ws.send_json({ + "event": "audio_start", + "content_type": "audio/wav", + "bytes": len(wav_bytes) + }) + + pos = 0 + chunk_size = 4096 + while pos < len(wav_bytes): + chunk = wav_bytes[pos : pos + chunk_size] + await ws.send_bytes(chunk) + pos += chunk_size + + await ws.send_json({"event": "audio_end"}) + await ws.send_json({"event": "done"}) + timings["audio_sent"] = time.time() + + except Exception as exc: + logger.exception("Error processing voice session") + try: + await ws.send_json({"event": "error", "message": f"Processing error: {exc}", "code": "processing_failed"}) + except Exception: + pass + finally: + wav_path.unlink(missing_ok=True) + try: + await ws.close() + except Exception: + pass + + if "audio_sent" in timings: + duration = timings["audio_sent"] - timings["ws_accepted"] + status = "completed" + else: + duration = time.time() - timings["ws_accepted"] + status = "failed" + + logger.info("ESP32 Voice WebSocket Session Completed. Timings: %s", timings) + audit_event = { + "timestamp": time.time(), + "device_id": device_id, + "source_ip": getattr(request, "remote", "") or "", + "duration": duration, + "byte_count": total_bytes, + "timings": timings, + "status": status + } + self._append_esp32_voice_audit(audio_dir, audit_event) + + return ws + def register_esp32_voice_route(app: Any, api_server: HermesESP32VoiceGatewayMixin) -> None: """Register the endpoint on an aiohttp app.""" app.router.add_post("/api/esp32/voice", api_server._handle_esp32_voice) + app.router.add_get("/api/esp32/voice/ws", api_server._handle_esp32_voice_websocket) + diff --git a/hermes_voice_gateway/test_websocket.py b/hermes_voice_gateway/test_websocket.py new file mode 100644 index 0000000..ebae5ac --- /dev/null +++ b/hermes_voice_gateway/test_websocket.py @@ -0,0 +1,230 @@ +# pyright: reportMissingImports=false, reportAttributeAccessIssue=false +"""Integration and unit tests for the ESP32 Voice WebSocket gateway endpoint.""" + +import asyncio +import json +import logging +import sys +import unittest +from pathlib import Path +from unittest.mock import MagicMock + +# 1. Mock external tools modules before importing the API endpoint +mock_transcription = MagicMock() +mock_transcription.transcribe_audio.return_value = {"success": True, "transcript": "Hello world"} +sys.modules["tools.transcription_tools"] = mock_transcription + +mock_tts = MagicMock() +mock_tts.text_to_speech_tool.return_value = json.dumps({"success": True, "file_path": __file__}) +sys.modules["tools.tts_tool"] = mock_tts + +# Mock aiohttp +from aiohttp import web, ClientSession + +# Import gateway mixin and register function +sys.path.append(str(Path(__file__).parent)) +from api_server_endpoint import HermesESP32VoiceGatewayMixin, register_esp32_voice_route + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +class MockHermesServer(HermesESP32VoiceGatewayMixin): + def __init__(self): + self.auth_token = "valid_secret_key" + self._background_tasks = set() + + def _check_auth(self, request): + auth_header = request.headers.get("Authorization") + if not auth_header or auth_header != f"Bearer {self.auth_token}": + return web.json_response({"error": "Unauthorized"}, status=401) + return None + + async def _run_esp32_voice_turn(self, transcript, device_id, instructions=None): + return {"id": "resp_test123"}, f"Echo: {transcript}" + + @staticmethod + def _convert_audio_to_esp32_wav(input_path, *, output_dir, device_id): + out_path = Path(output_dir) / "mock_response.wav" + import wave + with wave.open(str(out_path), "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(16000) + wav.writeframes(b"\x00" * 8000) # Dummy PCM data + return out_path + + +class TestWebSocketGateway(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.server = MockHermesServer() + self.app = web.Application() + register_esp32_voice_route(self.app, self.server) + + self.runner = web.AppRunner(self.app) + await self.runner.setup() + self.site = web.TCPSite(self.runner, "127.0.0.1", 0) + await self.site.start() + + self.port = self.runner.addresses[0][1] + self.base_url = f"http://127.0.0.1:{self.port}" + self.ws_url = f"ws://127.0.0.1:{self.port}/api/esp32/voice/ws" + + async def asyncTearDown(self): + await self.runner.cleanup() + + async def test_auth_failure_missing_token(self): + """Verify authentication failure (missing token) returns 401.""" + async with ClientSession() as session: + try: + async with session.ws_connect(self.ws_url) as ws: + pass + self.fail("Connection should have been rejected with 401") + except Exception as e: + # HTTP status code should be checked + pass + + # Test HTTP request directly to assert 401 + async with session.get(self.ws_url) as resp: + self.assertEqual(resp.status, 401) + body = await resp.json() + self.assertEqual(body["error"], "Unauthorized") + + async def test_auth_failure_bad_token(self): + """Verify authentication failure (bad token) returns 401.""" + headers = {"Authorization": "Bearer bad_token"} + async with ClientSession() as session: + async with session.get(self.ws_url, headers=headers) as resp: + self.assertEqual(resp.status, 401) + + async def test_successful_streaming_flow(self): + """Verify a normal WebSocket start, stream, stop, and response flow.""" + headers = { + "Authorization": "Bearer valid_secret_key", + "X-Device-ID": "test-device" + } + + async with ClientSession() as session: + async with session.ws_connect(self.ws_url, headers=headers) as ws: + # 1. Send start event + await ws.send_str(json.dumps({ + "event": "start", + "device_id": "test-device", + "sample_rate": 16000 + })) + + # Read events + ready_msg = await ws.receive_json() + self.assertEqual(ready_msg["event"], "ready") + + listening_msg = await ws.receive_json() + self.assertEqual(listening_msg["event"], "listening") + + # 2. Send simulated binary PCM audio chunks (Total: 4000 bytes, > 3200 byte limit) + await ws.send_bytes(b"\x00" * 2000) + await ws.send_bytes(b"\x00" * 2000) + + # 3. Send stop event + await ws.send_str(json.dumps({"event": "stop"})) + + # 4. Check transcription event + transcript_msg = await ws.receive_json() + self.assertEqual(transcript_msg["event"], "transcript") + self.assertEqual(transcript_msg["text"], "Hello world") + + # 5. Check thinking event + thinking_msg = await ws.receive_json() + self.assertEqual(thinking_msg["event"], "thinking") + + # 6. Check response_text event + resp_text_msg = await ws.receive_json() + self.assertEqual(resp_text_msg["event"], "response_text") + self.assertEqual(resp_text_msg["text"], "Echo: Hello world") + + # 7. Check audio_start event + audio_start_msg = await ws.receive_json() + self.assertEqual(audio_start_msg["event"], "audio_start") + self.assertEqual(audio_start_msg["content_type"], "audio/wav") + audio_bytes_len = audio_start_msg["bytes"] + + # 8. Check binary audio chunks + received_audio = b"" + while True: + msg = await ws.receive() + if msg.type == web.WSMsgType.BINARY: + received_audio += msg.data + elif msg.type == web.WSMsgType.TEXT: + data = json.loads(msg.data) + if data["event"] == "audio_end": + break + else: + self.fail(f"Unexpected msg type: {msg.type}") + + self.assertEqual(len(received_audio), audio_bytes_len) + + # 9. Check done event + done_msg = await ws.receive_json() + self.assertEqual(done_msg["event"], "done") + + async def test_disconnect_cleanup(self): + """Verify that raw files are cleaned up if client disconnects abruptly.""" + headers = { + "Authorization": "Bearer valid_secret_key", + "X-Device-ID": "disconnect-device" + } + + # We need to find the file created during this session. + # Let's verify files in ~/.hermes/cache/audio/esp32 or /tmp/esp32 + audio_dir = Path.home() / ".hermes" / "cache" / "audio" / "esp32" + existing_files_before = set(audio_dir.glob("ws_disconnect_device_*.s16le")) if audio_dir.exists() else set() + + async with ClientSession() as session: + async with session.ws_connect(self.ws_url, headers=headers) as ws: + await ws.send_str(json.dumps({ + "event": "start", + "device_id": "disconnect-device" + })) + + await ws.receive_json() # ready + await ws.receive_json() # listening + + # Send some bytes + await ws.send_bytes(b"\x00" * 1000) + + # Abrupt close/exit without 'stop' + await ws.close() + + # Allow loop to process close + await asyncio.sleep(0.5) + + existing_files_after = set(audio_dir.glob("ws_disconnect_device_*.s16le")) if audio_dir.exists() else set() + new_files = existing_files_after - existing_files_before + + # Verify no orphaned raw files remaining + self.assertEqual(len(new_files), 0, f"Found orphaned temp files: {new_files}") + + async def test_too_short_audio(self): + """Verify session is rejected with audio_too_short error if total PCM is below 3200 bytes.""" + headers = { + "Authorization": "Bearer valid_secret_key", + "X-Device-ID": "short-device" + } + async with ClientSession() as session: + async with session.ws_connect(self.ws_url, headers=headers) as ws: + await ws.send_str(json.dumps({ + "event": "start", + "device_id": "short-device" + })) + await ws.receive_json() # ready + await ws.receive_json() # listening + + # Send only 1000 bytes (less than 3200 bytes limit) + await ws.send_bytes(b"\x00" * 1000) + await ws.send_str(json.dumps({"event": "stop"})) + + error_msg = await ws.receive_json() + self.assertEqual(error_msg["event"], "error") + self.assertEqual(error_msg["code"], "audio_too_short") + + +if __name__ == "__main__": + unittest.main() diff --git a/lib/websocket_client.py b/lib/websocket_client.py new file mode 100644 index 0000000..c221b22 --- /dev/null +++ b/lib/websocket_client.py @@ -0,0 +1,184 @@ +# pyright: reportMissingImports=false, reportAttributeAccessIssue=false +"""Lightweight WebSocket client for MicroPython. + +This module provides a minimal, robust, memory-efficient WebSocket client +supporting HTTP upgrade handshakes and client-to-server frame masking +(RFC 6455). +""" + +import usocket as socket +import ustruct as struct +import urandom as random +import ubinascii as binascii + +def parse_url(url): + if not url.startswith("ws://") and not url.startswith("wss://"): + raise ValueError("URL must start with ws:// or wss://") + is_ssl = url.startswith("wss://") + url_p = url.split("://", 1)[1] + parts = url_p.split("/", 1) + host_port = parts[0] + path = "/" + parts[1] if len(parts) > 1 else "/" + if ":" in host_port: + host, port = host_port.split(":", 1) + port = int(port) + else: + host = host_port + port = 443 if is_ssl else 80 + return host, port, path, is_ssl + +class WebSocketClient: + def __init__(self, url, headers=None, timeout=10): + self.url = url + self.headers = headers or {} + self.timeout = timeout + self.sock = None + self.host, self.port, self.path, self.is_ssl = parse_url(url) + + def connect(self): + # 1. Generate standard Sec-WebSocket-Key + raw_key = bytes([random.getrandbits(8) for _ in range(16)]) + sec_key = binascii.b2a_base64(raw_key).decode('utf-8').strip() + + # 2. Resolve address and connect + addr = socket.getaddrinfo(self.host, self.port)[0][-1] + self.sock = socket.socket() + self.sock.settimeout(self.timeout) + self.sock.connect(addr) + + if self.is_ssl: + import ussl + self.sock = ussl.wrap_socket(self.sock) + + # 3. Construct HTTP GET upgrade request + req = [ + f"GET {self.path} HTTP/1.1", + f"Host: {self.host}:{self.port}", + "Upgrade: websocket", + "Connection: Upgrade", + f"Sec-WebSocket-Key: {sec_key}", + "Sec-WebSocket-Version: 13", + ] + for k, v in self.headers.items(): + req.append(f"{k}: {v}") + req.append("\r\n") + + self.sock.write("\r\n".join(req).encode('utf-8')) + + # 4. Read HTTP response status and headers + status_line = self.sock.readline() + if not status_line or b"101" not in status_line: + self.close() + raise RuntimeError(f"Handshake failed status: {status_line.decode('utf-8', 'ignore').strip()}") + + while True: + line = self.sock.readline() + if not line or line == b"\r\n": + break + + def send_frame(self, opcode, payload, fin=True): + """Send a masked WebSocket frame to the server (client-to-server MUST be masked).""" + if not self.sock: + raise RuntimeError("Not connected") + + # Header byte 0: FIN and Opcode + b0 = 0x80 if fin else 0 + b0 |= (opcode & 0x0F) + + # Header byte 1: Mask bit (always 1 for client) and Payload length + payload_len = len(payload) + if payload_len <= 125: + header = struct.pack("!BB", b0, 0x80 | payload_len) + elif payload_len <= 65535: + header = struct.pack("!BBH", b0, 0x80 | 126, payload_len) + else: + header = struct.pack("!BBQ", b0, 0x80 | 127, payload_len) + + # Generate 4-byte random masking key + mask = bytes([random.getrandbits(8) for _ in range(4)]) + + # Apply masking key (XOR payload) + masked_payload = bytearray(payload_len) + for i in range(payload_len): + masked_payload[i] = payload[i] ^ mask[i % 4] + + # Send header, mask, and masked payload + self.sock.write(header) + self.sock.write(mask) + self.sock.write(masked_payload) + + def send_text(self, text): + self.send_frame(0x1, text.encode('utf-8')) + + def send_binary(self, data): + self.send_frame(0x2, data) + + def recv_frame(self): + """Receive an unmasked WebSocket frame from the server (server-to-client is unmasked).""" + if not self.sock: + raise RuntimeError("Not connected") + + try: + header = self.sock.read(2) + except Exception: + # Handle socket timeout or disconnect + return None, None + + if not header or len(header) < 2: + return None, None + + b0, b1 = header + opcode = b0 & 0x0F + masked = bool(b1 & 0x80) + payload_len = b1 & 0x7F + + if payload_len == 126: + len_bytes = self.sock.read(2) + if not len_bytes or len(len_bytes) < 2: + return None, None + payload_len = struct.unpack("!H", len_bytes)[0] + elif payload_len == 127: + len_bytes = self.sock.read(8) + if not len_bytes or len(len_bytes) < 8: + return None, None + payload_len = struct.unpack("!Q", len_bytes)[0] + + if masked: + mask = self.sock.read(4) + if not mask or len(mask) < 4: + return None, None + + # Read actual payload + payload = b"" + while len(payload) < payload_len: + needed = payload_len - len(payload) + chunk = self.sock.read(needed) + if not chunk: + break + payload += chunk + + if len(payload) < payload_len: + # Socket closed prematurely + return None, None + + if masked: + # Unmask payload if masked + unmasked = bytearray(payload_len) + for i in range(payload_len): + unmasked[i] = payload[i] ^ mask[i % 4] + payload = bytes(unmasked) + + return opcode, payload + + def close(self): + if self.sock: + try: + # Send close frame + self.send_frame(0x8, b"") + except Exception: + pass + try: + self.sock.close() + except Exception: + pass + self.sock = None