Detect stalled STT connections and limit recovery retries

livekit-agents/livekit/agents/__init__.py

@@ -98,6 +98,7 @@
     RecordingOptions,
     RunContext,
     RunOutputOptions,
+    SessionConnectOptions,
     SessionUsageUpdatedEvent,
     SpeechCreatedEvent,
     SpeechSteeringOptions,
@@ -226,6 +227,7 @@ def __getattr__(name: str) -> typing.Any:
     "SimulationRun",
     "SimulationVerdict",
     "AgentSession",
+    "SessionConnectOptions",
     "AudioRecognition",
     "ExpressiveOptions",
     "NonverbalOptions",

livekit-agents/livekit/agents/inference/stt.py

@@ -1186,8 +1186,11 @@ async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
         @utils.log_exceptions(logger=logger)
         async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
             nonlocal session_closed
+            monitor = utils.aio.WebSocketMonitor(
+                ws, liveness_timeout=self._conn_options.liveness_timeout
+            )
             while True:
-                msg = await ws.receive()
+                msg = await monitor.receive()
                 if msg.type in (
                     aiohttp.WSMsgType.CLOSED,
                     aiohttp.WSMsgType.CLOSE,
@@ -1204,7 +1207,7 @@ async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
                 if msg.type == aiohttp.WSMsgType.ERROR:
                     if input_ended.is_set() or http_session.closed:
                         return
-                    # a missed heartbeat pong surfaces here; the reason is only on ws.exception()
+                    # Preserve the underlying transport error.
                     raise APIConnectionError(
                         "LiveKit Inference STT connection lost"
                     ) from ws.exception()
@@ -1256,6 +1259,14 @@ async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
                         body={"code": code},
                         retryable=not input_ended.is_set(),
                     )
+                if msg_type in (
+                    "start_of_speech",
+                    "interim_transcript",
+                    "preflight_transcript",
+                    "final_transcript",
+                    "session.finalized",
+                ):
+                    self._mark_provider_progress()
 
         ws: aiohttp.ClientWebSocketResponse | None = None
         try:
@@ -1351,8 +1362,7 @@ async def _connect_ws(
                 http_session.ws_connect(
                     f"{base_url}/stt?{urlencode({'model': self._opts.model})}",
                     headers=headers,
-                    # without it a half-open socket parks recv_task forever
-                    heartbeat=30.0,
+                    autoping=False,
                 ),
                 self._conn_options.timeout,
             )

livekit-agents/livekit/agents/stt/fallback_adapter.py

@@ -409,6 +409,7 @@ async def _forward_input_task() -> None:
                             retry_interval=self._fallback_adapter._retry_interval,
                         ),
                     )
+                    main_stream._on_provider_progress = self._mark_provider_progress
                     # update main_stream start time offset so transcript timestamps are properly adjusted
                     main_stream.start_time_offset = self.start_time_offset + (
                         time.time() - self._start_time

livekit-agents/livekit/agents/stt/multi_speaker_adapter.py

@@ -142,7 +142,7 @@ async def _metrics_monitor_task(self, event_aiter: AsyncIterable[SpeechEvent]) -
         # accumulate instead of being forgiven by a successful transcript.
         async for ev in event_aiter:
             if ev.type == SpeechEventType.FINAL_TRANSCRIPT:
-                self._num_retries = 0
+                self._mark_provider_progress()
 
     async def _run(self) -> None:
         async def _forward_input(stream: RecognizeStream) -> None:
@@ -171,6 +171,7 @@ async def _forward_output(stream: RecognizeStream) -> None:
                     )
 
         stream = self._wrapped_stt.stream(language=self._language, conn_options=self._conn_options)
+        stream._on_provider_progress = self._mark_provider_progress
         tasks = [
             asyncio.create_task(
                 _forward_input(stream), name="DiarizationAdapterWrapper.forward_input"

livekit-agents/livekit/agents/stt/stream_adapter.py

@@ -144,6 +144,7 @@ async def _recognize() -> None:
                         language=self._language,
                         conn_options=self._wrapped_stt_conn_options,
                     )
+                    self._mark_provider_progress()
 
                     if len(t_event.alternatives) == 0:
                         continue

livekit-agents/livekit/agents/stt/stt.py

@@ -3,7 +3,7 @@
 import asyncio
 import time
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterable, AsyncIterator
+from collections.abc import AsyncIterable, AsyncIterator, Callable
 from dataclasses import dataclass, field
 from enum import Enum, unique
 from types import TracebackType
@@ -399,6 +399,7 @@ def __init__(
         )
 
         self._num_retries = 0
+        self._on_provider_progress: Callable[[], None] | None = None
         self._task = asyncio.create_task(self._main_task())
         self._task.add_done_callback(lambda _: self._event_ch.close())
 
@@ -485,11 +486,6 @@ async def _main_task(self) -> None:
                 last_start_time = time.time()
                 return await self._run()
             except APIError as e:
-                # an attempt that outlived the connect timeout had connected, so this failure
-                # is not consecutive with the previous one and the budget starts over
-                if time.time() - last_start_time > self._conn_options.timeout:
-                    self._num_retries = 0
-
                 if not e.retryable or max_retries == 0:
                     self._emit_error(e, recoverable=False)
                     raise
@@ -518,6 +514,17 @@ async def _main_task(self) -> None:
                 self._emit_error(e, recoverable=False)
                 raise
 
+    def _mark_provider_progress(self) -> None:
+        """Reset recovery retries after a valid provider processing event.
+
+        Call at the receive/parse boundary for speech, turn, transcription, or
+        audio-processing events. Setup, errors, heartbeats, and local events do
+        not indicate provider progress.
+        """
+        self._num_retries = 0
+        if self._on_provider_progress is not None:
+            self._on_provider_progress()
+
     def _emit_error(self, api_error: Exception, recoverable: bool) -> None:
         self._stt.emit(
             "error",
@@ -554,8 +561,8 @@ async def _metrics_monitor_task(self, event_aiter: AsyncIterable[SpeechEvent]) -
 
                 self._stt.emit("metrics_collected", stt_metrics)
             elif ev.type == SpeechEventType.FINAL_TRANSCRIPT:
-                # reset the retry count after a successful recognition
-                self._num_retries = 0
+                # Preserve recovery resets for custom providers that emit only transcripts.
+                self._mark_provider_progress()
 
     def push_frame(self, frame: rtc.AudioFrame) -> None:
         """Push audio to be recognized"""

livekit-agents/livekit/agents/types.py

@@ -1,3 +1,4 @@
+import math
 from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Any, Literal, TypeAlias, TypeVar
@@ -143,6 +144,15 @@ class APIConnectOptions:
     Timeout for connecting to the API in seconds.
     """
 
+    liveness_timeout: float = 6.0
+    """
+    Maximum time, in seconds, without inbound transport activity before the connection fails.
+    Incoming data, PING, and PONG frames reset this timer. An idle connection is probed
+    after one-third of this interval.
+
+    Supported by streaming STT providers that use the framework WebSocket monitor.
+    """
+
     def __post_init__(self) -> None:
         if self.max_retry < 0:
             raise ValueError("max_retry must be greater than or equal to 0")
@@ -153,6 +163,9 @@ def __post_init__(self) -> None:
         if self.timeout < 0:
             raise ValueError("timeout must be greater than or equal to 0")
 
+        if not math.isfinite(self.liveness_timeout) or self.liveness_timeout <= 0:
+            raise ValueError("liveness_timeout must be a finite number greater than 0")
+
     def _interval_for_retry(self, num_retries: int) -> float:
         """
         Return the interval for the given number of retries.

livekit-agents/livekit/agents/utils/aio/__init__.py

@@ -6,6 +6,7 @@
 from .task_set import TaskSet
 from .utils import cancel_and_wait, gracefully_cancel
 from .wait_group import WaitGroup
+from .websocket import WebSocketMonitor
 
 __all__ = [
     "AsyncAtomicCounter",
@@ -20,6 +21,7 @@
     "sleep",
     "TaskSet",
     "WaitGroup",
+    "WebSocketMonitor",
     "debug",
     "cancel_and_wait",
     "duplex_unix",

livekit-agents/livekit/agents/utils/aio/websocket.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+
+import aiohttp
+
+from ..._exceptions import APIConnectionError
+from ...types import DEFAULT_API_CONNECT_OPTIONS
+
+_CLOSE_TIMEOUT = 0.1
+
+
+class WebSocketMonitor:
+    """Receive WebSocket messages with an idle probe and a bounded silence timeout.
+
+    The socket must be connected with ``autoping=False`` and without aiohttp's
+    heartbeat. Incoming control frames are handled here so they count as activity.
+    Use one monitor per active connection, including across cancelled receive calls.
+    """
+
+    def __init__(
+        self,
+        ws: aiohttp.ClientWebSocketResponse,
+        *,
+        liveness_timeout: float = DEFAULT_API_CONNECT_OPTIONS.liveness_timeout,
+    ) -> None:
+        self._ws = ws
+        self._liveness_timeout = liveness_timeout
+        self._ping_interval = liveness_timeout / 3
+        self._loop = asyncio.get_running_loop()
+        self._last_received = self._loop.time()
+        self._ping_sent = False
+
+    async def receive(self) -> aiohttp.WSMessage:
+        receiver = asyncio.create_task(self._ws.receive())
+        ping: asyncio.Task[None] | None = None
+        try:
+            while True:
+                deadline = self._last_received + (
+                    self._liveness_timeout if self._ping_sent else self._ping_interval
+                )
+                tasks: list[asyncio.Task[aiohttp.WSMessage] | asyncio.Task[None]] = [receiver]
+                if ping is not None:
+                    tasks.append(ping)
+                done, _ = await asyncio.wait(
+                    tasks,
+                    timeout=max(0.0, deadline - self._loop.time()),
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                if receiver in done:
+                    msg = receiver.result()
+                    if msg.type in (
+                        aiohttp.WSMsgType.TEXT,
+                        aiohttp.WSMsgType.BINARY,
+                        aiohttp.WSMsgType.PING,
+                        aiohttp.WSMsgType.PONG,
+                    ):
+                        self._last_received = self._loop.time()
+                        self._ping_sent = False
+                        if ping is not None:
+                            ping.cancel()
+                            await asyncio.gather(ping, return_exceptions=True)
+                            ping = None
+
+                    if msg.type == aiohttp.WSMsgType.PING:
+                        await asyncio.wait_for(self._ws.pong(msg.data), self._liveness_timeout)
+                    elif msg.type != aiohttp.WSMsgType.PONG:
+                        return msg
+                    receiver = asyncio.create_task(self._ws.receive())
+                    continue
+
+                if ping is not None and ping in done:
+                    ping.result()
+                    ping = None
+                    continue
+
+                if self._loop.time() - self._last_received >= self._liveness_timeout:
+                    raise aiohttp.ServerTimeoutError(
+                        f"No WebSocket messages received for {self._liveness_timeout} seconds"
+                    )
+                self._ping_sent = True
+                ping = asyncio.create_task(self._ws.ping())
+        except (aiohttp.ClientError, ConnectionError, asyncio.TimeoutError) as exc:
+            # A dead peer cannot complete a close handshake. Bound cleanup so the
+            # original failure reaches retry/fallback without another close timeout.
+            with contextlib.suppress(Exception):
+                await asyncio.wait_for(self._ws.close(), _CLOSE_TIMEOUT)
+            raise APIConnectionError("WebSocket connection lost") from exc
+        finally:
+            receiver.cancel()
+            if ping is not None:
+                ping.cancel()
+            await asyncio.gather(
+                receiver, *([ping] if ping is not None else []), return_exceptions=True
+            )
+
+    def __aiter__(self) -> WebSocketMonitor:
+        return self
+
+    async def __anext__(self) -> aiohttp.WSMessage:
+        msg = await self.receive()
+        if msg.type in (
+            aiohttp.WSMsgType.CLOSE,
+            aiohttp.WSMsgType.CLOSED,
+            aiohttp.WSMsgType.CLOSING,
+        ):
+            raise StopAsyncIteration
+        return msg

livekit-agents/livekit/agents/voice/__init__.py

@@ -5,6 +5,7 @@
     ExpressiveOptions,
     NonverbalOptions,
     RecordingOptions,
+    SessionConnectOptions,
     SpeechSteeringOptions,
     VoiceActivityVideoSampler,
 )
@@ -52,6 +53,7 @@
     "ExpressiveOptions",
     "NonverbalOptions",
     "RecordingOptions",
+    "SessionConnectOptions",
     "SpeechSteeringOptions",
     "RunOutputOptions",
     "VoiceActivityVideoSampler",

livekit-plugins/livekit-plugins-assemblyai/livekit/plugins/assemblyai/stt.py

@@ -717,9 +717,12 @@ async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
         async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
             nonlocal closing_ws
             consecutive_timeouts = 0
+            monitor = utils.aio.WebSocketMonitor(
+                ws, liveness_timeout=self._conn_options.liveness_timeout
+            )
             while True:
                 try:
-                    msg = await asyncio.wait_for(ws.receive(), timeout=5)
+                    msg = await asyncio.wait_for(monitor.receive(), timeout=5)
                     consecutive_timeouts = 0
                 except asyncio.TimeoutError:
                     if closing_ws:
@@ -777,7 +780,10 @@ async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
                     continue
 
                 try:
-                    self._process_stream_event(json.loads(msg.data))
+                    data = json.loads(msg.data)
+                    self._process_stream_event(data)
+                    if data.get("type") in ("SpeechStarted", "Turn"):
+                        self._mark_provider_progress()
                 except Exception:
                     logger.exception(
                         "failed to process AssemblyAI message session=%s",
@@ -900,7 +906,7 @@ async def _connect_ws(self) -> aiohttp.ClientWebSocketResponse:
             self._opts.speech_model,
             self._base_url,
         )
-        ws = await self._session.ws_connect(url, headers=headers)
+        ws = await self._session.ws_connect(url, headers=headers, autoping=False)
         logger.debug(
             "AssemblyAI WebSocket connected status=%s",
             ws._response.status if ws._response is not None else None,