fix(gradium): surface TTS server errors, close codes and transport errors (1/2)

livekit-plugins/livekit-plugins-gradium/livekit/plugins/gradium/tts.py

@@ -26,6 +26,7 @@
 from livekit.agents import (
     APIConnectionError,
     APIConnectOptions,
+    APIError,
     APIStatusError,
     APITimeoutError,
     tokenize,
@@ -41,6 +42,45 @@
 SUPPORTED_SAMPLE_RATE = 48000
 DEFAULT_VOICE_ID = "4SZHfMpw-p46Ywgs"
 
+# Websocket close codes the server uses for requests it would reject again verbatim:
+# protocol error and unsupported data (unknown model, setup validation). Everything
+# else is worth a retry. 1008 (policy violation) covers the org's concurrency cap,
+# which frees up when another session ends, next to credits and API-key failures
+# that will not; the frame does not tell them apart, so the transient reading wins.
+# 1011 covers worker failures and "no worker available" (but also unknown voice ids,
+# which the server does not distinguish for clients either). The server's idle and
+# duration limits close with 1008 after an error frame; a 1000/1001 before
+# end_of_stream means the server shut the socket mid-session, also worth a retry.
+_NON_RETRYABLE_CLOSE_CODES = frozenset({1002, 1003})
+
+
+def _close_error(
+    message: str, *, code: int | None, body: Any, reason: Any = None
+) -> APIStatusError:
+    code = code if isinstance(code, int) else -1
+    if isinstance(reason, str) and reason:
+        # the close frame's reason is the server's own explanation of the failure
+        message = f"{message}: {reason}"
+    return APIStatusError(
+        message,
+        status_code=code,
+        body=body,
+        retryable=code not in _NON_RETRYABLE_CLOSE_CODES,
+    )
+
+
+def _api_error(msg_data: dict[str, Any]) -> APIStatusError:
+    """Build the error for a `{"type": "error"}` frame.
+
+    The server sends it right before closing the socket, so the message is the only
+    place the actual failure reason (invalid voice, no credits, worker failure, ...)
+    shows up. Its `code` is the websocket close code that follows.
+    """
+    message = msg_data.get("message") or "unknown error"
+    return _close_error(
+        f"Gradium returned error: {message}", code=msg_data.get("code"), body=msg_data
+    )
+
 
 @dataclass
 class _TTSOptions:
@@ -241,14 +281,19 @@ async def _run(self, output_emitter: tts.AudioEmitter) -> None:
 
                 while True:
                     msg = await ws.receive()
+                    if msg.type == aiohttp.WSMsgType.ERROR:
+                        raise APIConnectionError("Gradium websocket error") from ws.exception()
                     if msg.type in (
                         aiohttp.WSMsgType.CLOSE,
                         aiohttp.WSMsgType.CLOSED,
                         aiohttp.WSMsgType.CLOSING,
                     ):
-                        # TODO(laurent): once we support returning eos in the api, we should enable this back.
-                        # raise APIStatusError("Gradium websocket connection closed unexpectedly")
-                        break
+                        raise _close_error(
+                            "Gradium connection closed unexpectedly",
+                            code=ws.close_code,
+                            body=f"{msg.data=} {msg.extra=}",
+                            reason=msg.extra,
+                        )
 
                     if msg.type == aiohttp.WSMsgType.TEXT:
                         msg_data = json.loads(msg.data)
@@ -263,12 +308,16 @@ async def _run(self, output_emitter: tts.AudioEmitter) -> None:
                             output_emitter.push(audio)
                         elif type_ == "end_of_stream":
                             break
+                        elif type_ == "error":
+                            raise _api_error(msg_data)
                         else:
                             logger.warning(f"unknown message type: {type_}")
                 output_emitter.flush()
 
         except asyncio.TimeoutError:
             raise APITimeoutError() from None
+        except APIError:
+            raise
         except aiohttp.ClientResponseError as e:
             raise APIStatusError(
                 message=e.message, status_code=e.status, request_id=None, body=None
@@ -324,6 +373,8 @@ async def _run_segments() -> None:
             await asyncio.gather(*tasks)
         except asyncio.TimeoutError:
             raise APITimeoutError() from None
+        except APIError:
+            raise
         except aiohttp.ClientResponseError as e:
             raise APIStatusError(
                 message=e.message, status_code=e.status, request_id=request_id, body=None
@@ -367,15 +418,19 @@ async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
         async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
             while True:
                 msg = await ws.receive()
+                if msg.type == aiohttp.WSMsgType.ERROR:
+                    raise APIConnectionError("Gradium websocket error") from ws.exception()
                 if msg.type in (
                     aiohttp.WSMsgType.CLOSE,
                     aiohttp.WSMsgType.CLOSED,
                     aiohttp.WSMsgType.CLOSING,
                 ):
-                    # TODO(laurent): once we support returning eos in the api, we should enable this back.
-                    # raise APIStatusError("Gradium websocket connection closed unexpectedly")
-                    output_emitter.end_segment()
-                    break
+                    raise _close_error(
+                        "Gradium connection closed unexpectedly",
+                        code=ws.close_code,
+                        body=f"{msg.data=} {msg.extra=}",
+                        reason=msg.extra,
+                    )
 
                 if msg.type == aiohttp.WSMsgType.TEXT:
                     msg_data = json.loads(msg.data)
@@ -391,6 +446,8 @@ async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
                     elif type_ == "end_of_stream":
                         output_emitter.end_segment()
                         break
+                    elif type_ == "error":
+                        raise _api_error(msg_data)
                     else:
                         logger.warning(f"unknown message type: {type_}")