fix(gladia): detect a silently dropped STT socket instead of hanging

livekit-plugins/livekit-plugins-gladia/livekit/plugins/gladia/stt.py

@@ -366,6 +366,11 @@ async def _recognize_impl(
             async with self._ensure_session().ws_connect(
                 session_url,
                 timeout=aiohttp.ClientWSTimeout(ws_receive=receive_timeout, ws_close=10),
+                # Without this a silently dropped socket (half-open TCP, no FIN/RST) is
+                # only noticed once ws_receive expires, which is five times the connect
+                # timeout. The heartbeat closes the socket as soon as a ping goes
+                # unanswered, and that surfaces as WSMsgType.ERROR below.
+                heartbeat=30.0,
             ) as ws:
                 # Combine audio frames to get a single frame with all raw PCM data
                 combined_frame = rtc.combine_audio_frames(buffer)
@@ -389,6 +394,7 @@ async def _recognize_impl(
 
                 # Wait for final transcript
                 utterances = []
+                got_final = False
 
                 # Receive messages until we get the post_final_transcript message
                 try:
@@ -406,6 +412,7 @@ async def _recognize_impl(
                             ):
                                 pass
                             elif data["type"] == "post_final_transcript":
+                                got_final = True
                                 break
                             elif data["type"] == "error":
                                 raise APIConnectionError(
@@ -437,6 +444,19 @@ async def _recognize_impl(
                             f"Timeout waiting for Gladia final transcript ({receive_timeout}s)"
                         ) from None
 
+                if not got_final:
+                    # aiohttp closes the socket itself when a heartbeat ping goes unanswered, so
+                    # `async for` ends on WSMsgType.CLOSED before the ERROR branch above runs.
+                    # Without this the caller gets an empty transcript instead of a retryable
+                    # error, which is what the ws_receive timeout used to raise.
+                    exc = ws.exception()
+                    if exc is not None:
+                        raise APIConnectionError("Gladia connection lost") from exc
+                    if not utterances:
+                        raise APIConnectionError(
+                            "Gladia socket closed before the final transcript arrived"
+                        )
+
                 # Create a speech event from the collected final utterances
                 return self._create_speech_event(
                     utterances, session_id, config.language_config.languages
@@ -744,6 +764,9 @@ def __init__(
         self._request_id = ""
         self._reconnect_event = asyncio.Event()
         self._ws: aiohttp.ClientWebSocketResponse | None = None
+        # set once this side has asked Gladia to stop, so the recv loop can tell an
+        # expected shutdown from a socket that went away underneath it
+        self._closing_ws = False
 
     def update_options(
         self,
@@ -874,8 +897,16 @@ async def _run(self) -> None:
                 backoff_time = 1.0
 
                 # Connect to the WebSocket
-                async with self._session.ws_connect(session_url) as ws:
+                async with self._session.ws_connect(
+                    session_url,
+                    # Without this a silently dropped socket (half-open TCP, no FIN/RST)
+                    # is never noticed: _recv_messages_task parks on receive forever, and
+                    # the retry in _main_task only runs when something raises. Matches the
+                    # Deepgram and Telnyx STT plugins.
+                    heartbeat=30.0,
+                ) as ws:
                     self._ws = ws
+                    self._closing_ws = False
                     logger.info(f"Connected to Gladia session {self._request_id}")
 
                     send_task = asyncio.create_task(self._send_audio_task())
@@ -949,7 +980,6 @@ async def _send_audio_task(self) -> None:
 
         has_ended = False
         last_frame: rtc.AudioFrame | None = None
-        closing_ws = False
 
         try:
             async for data in self._input_ch:
@@ -989,11 +1019,11 @@ async def _send_audio_task(self) -> None:
                         has_ended = False
 
             # Tell Gladia we're done sending audio when the stream ends
-            closing_ws = True
+            self._closing_ws = True
             if self._ws:
                 await self._ws.send_str(json.dumps({"type": "stop_recording"}))
         except (aiohttp.ClientError, ConnectionError) as e:
-            if closing_ws or self._session.closed:
+            if self._closing_ws or self._session.closed:
                 return
             raise APIConnectionError("Gladia connection closed unexpectedly") from e
 
@@ -1009,6 +1039,15 @@ async def _recv_messages_task(self) -> None:
                     self._process_gladia_message(data)
                 except Exception as e:
                     logger.exception(f"Error processing Gladia message: {e}")
+            elif msg.type == aiohttp.WSMsgType.ERROR:
+                if self._closing_ws or self._session.closed:
+                    return
+                # The heartbeat closes the socket when a ping goes unanswered, and that
+                # arrives here rather than as a close frame. Raising a retryable error
+                # (instead of logging it as an unexpected type and waiting for the CLOSED
+                # that follows) lets _main_task reconnect with the reason attached;
+                # ws.exception() is the only place it survives.
+                raise APIConnectionError("Gladia connection lost") from self._ws.exception()
             elif msg.type in (
                 aiohttp.WSMsgType.CLOSED,
                 aiohttp.WSMsgType.CLOSE,

tests/test_plugin_gladia_stt.py

@@ -0,0 +1,236 @@
+"""Tests for the Gladia STT plugin: half-open socket detection.
+
+Runs a real `SpeechStream` (its real `_run` loop, driven by `_main_task`) against
+fake sockets handed out by a fake aiohttp session, so the connect kwargs and the
+reconnect behaviour are assertable without a network. Same shape as the Deepgram
+tests added in #7206 and the Telnyx tests added in #7359.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any
+
+import aiohttp
+import pytest
+
+from livekit.agents import APIConnectOptions
+
+pytestmark = pytest.mark.plugin("gladia")
+
+# Retry immediately so a reconnect shows up within the test timeout.
+_FAST_RETRY = APIConnectOptions(max_retry=3, retry_interval=0.01, timeout=1.0)
+
+
+async def _wait_until(predicate, *, timeout: float = 5.0) -> None:
+    deadline = time.monotonic() + timeout
+    while not predicate():
+        assert time.monotonic() < deadline, "timed out waiting for the stream"
+        await asyncio.sleep(0.01)
+
+
+class _ParkedSocket:
+    """A socket that accepts writes and never delivers anything: the half-open case
+    from the client's point of view. The message iterator only ends once the socket
+    is closed, and the test never lets that happen."""
+
+    def __init__(self) -> None:
+        self.closed = False
+        self._closed = asyncio.Event()
+
+    async def send_str(self, data: str) -> None:
+        pass
+
+    async def send_bytes(self, data: bytes) -> None:
+        pass
+
+    def __aiter__(self) -> _ParkedSocket:
+        return self
+
+    async def __anext__(self) -> aiohttp.WSMessage:
+        await self._closed.wait()
+        raise StopAsyncIteration
+
+    async def close(self) -> None:
+        self.closed = True
+        self._closed.set()
+
+
+class _HeartbeatTimeoutSocket:
+    """A socket in the state aiohttp leaves it in when a ping goes unanswered.
+
+    The heartbeat closes the connection itself, so this arrives as WSMsgType.ERROR
+    rather than as a close frame, and the reason lives only on `exception()`. After a
+    few ERROR frames the iterator ends, which is what aiohttp does on the next pass;
+    a recv loop that steps over the ERROR only recovers because of that.
+    """
+
+    def __init__(self) -> None:
+        self.closed = False
+        self.receives = 0
+
+    async def send_str(self, data: str) -> None:
+        pass
+
+    async def send_bytes(self, data: bytes) -> None:
+        pass
+
+    def exception(self) -> BaseException:
+        return aiohttp.ServerTimeoutError("No PONG received after 15.0 seconds")
+
+    def __aiter__(self) -> _HeartbeatTimeoutSocket:
+        return self
+
+    async def __anext__(self) -> aiohttp.WSMessage:
+        self.receives += 1
+        # yield, so a recv loop that steps over the error instead of ending fails the
+        # receive-count assertion rather than starving the event loop
+        await asyncio.sleep(0)
+        if self.receives > 3:
+            raise StopAsyncIteration
+        return aiohttp.WSMessage(aiohttp.WSMsgType.ERROR, self.exception(), None)
+
+    async def close(self) -> None:
+        self.closed = True
+
+
+class _FakeResponse:
+    """The `POST /v2/live` response that hands out the session URL."""
+
+    def __init__(self, payload: dict[str, Any], status: int = 201) -> None:
+        # the one-shot path checks `status`, the streaming path calls `raise_for_status`
+        self._payload = payload
+        self.status = status
+
+    async def __aenter__(self) -> _FakeResponse:
+        return self
+
+    async def __aexit__(self, *exc_info: Any) -> None:
+        return None
+
+    def raise_for_status(self) -> None:
+        pass
+
+    async def json(self) -> dict[str, Any]:
+        return self._payload
+
+
+class _FakeWSContext:
+    """`ws_connect` is used as an async context manager by the Gladia plugin."""
+
+    def __init__(self, ws: Any) -> None:
+        self._ws = ws
+
+    async def __aenter__(self) -> Any:
+        return self._ws
+
+    async def __aexit__(self, *exc_info: Any) -> None:
+        return None
+
+
+class _FakeSession:
+    """Stands in for the aiohttp session: records connect kwargs, hands out sockets."""
+
+    def __init__(self, make_socket) -> None:
+        self.closed = False
+        self.kwargs: list[dict[str, Any]] = []
+        self.sockets: list[Any] = []
+        self._make_socket = make_socket
+
+    def post(self, **kwargs: Any) -> _FakeResponse:
+        return _FakeResponse({"id": "test-session", "url": "wss://example.invalid/live"})
+
+    def ws_connect(self, url: str, **kwargs: Any) -> _FakeWSContext:
+        self.kwargs.append(kwargs)
+        ws = self._make_socket()
+        self.sockets.append(ws)
+        return _FakeWSContext(ws)
+
+    async def close(self) -> None:
+        self.closed = True
+
+
+def _stt(session: _FakeSession):
+    from livekit.plugins.gladia import STT
+
+    return STT(api_key="test-key", http_session=session)  # type: ignore[arg-type]
+
+
+def _stream(session: _FakeSession):
+    from livekit.plugins.gladia import STT
+
+    instance = STT(api_key="test-key", http_session=session)  # type: ignore[arg-type]
+    return instance.stream(conn_options=_FAST_RETRY)
+
+
+async def test_socket_is_opened_with_a_heartbeat():
+    """aiohttp defaults `heartbeat` to None, so without it the read side of a
+    half-open socket parks forever and the retry in `_main_task`, which only runs
+    when something raises, never gets a turn."""
+    session = _FakeSession(_ParkedSocket)
+    stream = _stream(session)
+    try:
+        await _wait_until(lambda: len(session.kwargs) == 1)
+        assert session.kwargs[0].get("heartbeat") == 30.0
+    finally:
+        await stream.aclose()
+
+
+async def test_heartbeat_timeout_reconnects_without_spinning():
+    """The ERROR frame has to end the recv loop and surface as a retryable error.
+    Logging it as an unexpected message type and continuing only reconnected because
+    aiohttp happened to end the iterator next, and it threw away the one value that
+    says why the socket went."""
+    session = _FakeSession(_HeartbeatTimeoutSocket)
+    stream = _stream(session)
+    try:
+        await _wait_until(lambda: len(session.sockets) > 1)
+        assert session.sockets[0].receives == 1
+    finally:
+        await stream.aclose()
+
+
+class _SilentlyClosedSocket:
+    """The one-shot path's view of a heartbeat timeout.
+
+    aiohttp sets the socket closed itself when a ping goes unanswered, so the next
+    `receive()` returns WS_CLOSED_MESSAGE and `__anext__` raises StopAsyncIteration.
+    The ERROR branch never runs, and the reason lives only on `exception()`.
+    """
+
+    def __init__(self) -> None:
+        self.closed = False
+
+    async def send_str(self, data: str) -> None:
+        pass
+
+    async def send_bytes(self, data: bytes) -> None:
+        pass
+
+    def exception(self) -> BaseException:
+        return aiohttp.ServerTimeoutError("No PONG received after 30.0 seconds")
+
+    def __aiter__(self) -> _SilentlyClosedSocket:
+        return self
+
+    async def __anext__(self) -> aiohttp.WSMessage:
+        raise StopAsyncIteration
+
+    async def close(self) -> None:
+        self.closed = True
+
+
+async def test_recognize_raises_when_the_socket_dies_before_the_final_transcript():
+    """Adding a heartbeat means a dead socket now ends the `async for` on CLOSED
+    instead of expiring ws_receive. Without an explicit check the one-shot path
+    returns an empty transcript, which is quieter than the timeout it replaced."""
+    from livekit.agents import APIConnectionError
+    from livekit.rtc import AudioFrame
+
+    session = _FakeSession(_SilentlyClosedSocket)
+    stt_instance = _stt(session)
+    frame = AudioFrame.create(sample_rate=16000, num_channels=1, samples_per_channel=1600)
+
+    with pytest.raises(APIConnectionError):
+        await stt_instance.recognize(buffer=frame, conn_options=_FAST_RETRY)