Implement WebSocket real-time audio streaming for Hermes gateway
This commit is contained in:
Adolfo Reyna
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user