Implement WebSocket real-time audio streaming for Hermes gateway

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)
+

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()