Websockets: reconnect if connection closes unexpectedly

playht · Jan 16, 2025 · ce04bac · ce04bac
1 parent ba4b745
commit ce04bac
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/pyht/async_client.py b/pyht/async_client.py
@@ -13,6 +13,7 @@
 from typing import Any, Dict, AsyncGenerator, AsyncIterable, AsyncIterator, Coroutine, Tuple, Optional, Union
 import uuid
 from websockets.asyncio.client import connect, ClientConnection
+from websockets.exceptions import ConnectionClosed
 
 import aiohttp
 import filelock
@@ -455,7 +456,12 @@ async def _tts_ws(
                 ws_address = self._inference_coordinates[voice_engine]["websocket_url"]
                 if self._ws is None:
                     self._ws = await connect(ws_address)
-                await self._ws.send(json.dumps(json_data))
+                try:
+                    await self._ws.send(json.dumps(json_data))
+                except ConnectionClosed as e:
+                    logging.debug(f"Reconnecting websocket which closed unexpectedly: {e}")
+                    self._ws = await connect(ws_address)
+                    await self._ws.send(json.dumps(json_data))
                 chunk_idx = -1
                 async for chunk in self._ws:
                     chunk_idx += 1

diff --git a/pyht/client.py b/pyht/client.py
@@ -14,6 +14,7 @@
 import time
 import uuid
 from websockets.sync.client import connect, ClientConnection
+from websockets.exceptions import ConnectionClosed
 
 import filelock
 import grpc
@@ -724,7 +725,12 @@ def _tts_ws(
                 ws_address = self._inference_coordinates[voice_engine]["websocket_url"]
                 if self._ws is None:
                     self._ws = connect(ws_address)
-                self._ws.send(json.dumps(json_data))
+                try:
+                    self._ws.send(json.dumps(json_data))
+                except ConnectionClosed as e:
+                    logging.debug(f"Reconnecting websocket which closed unexpectedly: {e}")
+                    self._ws = connect(ws_address)
+                    self._ws.send(json.dumps(json_data))
                 chunk_idx = -1
                 for chunk in self._ws:
                     chunk_idx += 1