(stt): reset the retry budget on any event, not only on a final transcript

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

@@ -356,6 +356,26 @@ class SpeakerContext(Protocol):
     def to_instructions(self) -> str: ...
 
 
+class _HealthySignallingChan(aio.Chan[SpeechEvent]):
+    """Event channel that clears the retry budget the moment an event is published.
+
+    The reset has to happen in the producer's own turn. `send_nowait` only wakes the
+    metrics consumer, it does not run it, so an attempt that publishes an event and
+    then raises in the same event-loop turn would reach the terminal branch of
+    `RecognizeStream._main_task` with the budget still exhausted. Signalling here
+    also ties the reset to the attempt that produced the event, rather than to
+    whenever a consumer happens to drain it.
+    """
+
+    def __init__(self, stream: RecognizeStream) -> None:
+        super().__init__()
+        self._stream = stream
+
+    def send_nowait(self, value: SpeechEvent) -> None:
+        self._stream._num_retries = 0
+        super().send_nowait(value)
+
+
 class RecognizeStream(ABC):
     class _FlushSentinel:
         """Sentinel to mark when it was flushed"""
@@ -380,7 +400,7 @@ def __init__(
         self._stt = stt
         self._conn_options = conn_options
         self._input_ch = aio.Chan[rtc.AudioFrame | RecognizeStream._FlushSentinel]()
-        self._event_ch = aio.Chan[SpeechEvent]()
+        self._event_ch = _HealthySignallingChan(self)
 
         self._tee = aio.itertools.tee(self._event_ch, 2)
         self._event_aiter, monitor_aiter = self._tee
@@ -538,9 +558,6 @@ 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
 
     def push_frame(self, frame: rtc.AudioFrame) -> None:
         """Push audio to be recognized"""

tests/test_stt_base.py

@@ -3,13 +3,15 @@
 from __future__ import annotations
 
 import asyncio
+import dataclasses
 import time
 
 import pytest
 
 from livekit.agents import APIConnectionError, APIStatusError
 from livekit.agents.stt import (
     STT,
+    RecognitionUsage,
     RecognizeStream,
     SpeechData,
     SpeechEvent,
@@ -180,3 +182,102 @@ async def test_stream_adapter_keeps_vad_speech_end_on_delayed_final(
     assert end_event.speech_end_time is not None
     assert final_event.speech_end_time == end_event.speech_end_time
     assert final_event.created_at - end_event.created_at == pytest.approx(0.5, abs=0.01)
+
+
+class _FlappingStream(RecognizeStream):
+    """Every connection succeeds, then drops. A real reconnect, every time."""
+
+    def __init__(
+        self,
+        *,
+        stt: STT,
+        drops: int,
+        emit: SpeechEvent | None,
+        silent_runs: int = 0,
+        yield_after_emit: bool = True,
+    ) -> None:
+        super().__init__(
+            stt=stt,
+            conn_options=dataclasses.replace(DEFAULT_API_CONNECT_OPTIONS, retry_interval=0.0),
+        )
+        self.runs = 0
+        self._drops = drops
+        self._emit = emit
+        self._silent_runs = silent_runs
+        self._yield_after_emit = yield_after_emit
+
+    async def _run(self) -> None:
+        self.runs += 1
+        if self._emit is not None and self.runs > self._silent_runs:
+            self._event_ch.send_nowait(self._emit)
+            if self._yield_after_emit:
+                await asyncio.sleep(0)
+        if self.runs <= self._drops:
+            raise APIConnectionError("socket dropped")
+        await asyncio.sleep(3600)  # healthy at last
+
+
+async def _survives(
+    emit: SpeechEvent | None,
+    *,
+    drops: int = 10,
+    silent_runs: int = 0,
+    yield_after_emit: bool = True,
+) -> bool:
+    stream = _FlappingStream(
+        stt=_DummySTT(),
+        drops=drops,
+        emit=emit,
+        silent_runs=silent_runs,
+        yield_after_emit=yield_after_emit,
+    )
+    try:
+        for _ in range(500):
+            if stream._task.done() or stream.runs > drops:
+                break
+            await asyncio.sleep(0.01)
+        return not stream._task.done()
+    finally:
+        if not stream._task.done():
+            await stream.aclose()
+
+
+async def test_retry_budget_survives_drops_while_the_caller_is_silent() -> None:
+    """The budget counts consecutive failures, not the lifetime of the stream.
+
+    Every reconnect here succeeds, so the stream is healthy throughout. Resetting
+    only on FINAL_TRANSCRIPT tied the budget to the caller speaking: an agent
+    talking over a silent caller, or a caller on hold, never earned it back, and
+    max_retry drops spread over one long call killed the stream for good.
+    """
+    usage_only = SpeechEvent(
+        type=SpeechEventType.RECOGNITION_USAGE,
+        recognition_usage=RecognitionUsage(audio_duration=5.0),
+    )
+    assert await _survives(usage_only)
+
+
+async def test_retry_budget_still_gives_up_when_nothing_is_ever_delivered() -> None:
+    """A connection that has never delivered anything must not retry forever."""
+    assert not await _survives(None)
+
+
+async def test_retry_budget_resets_even_when_the_producer_never_yields() -> None:
+    """The reset must land in the producer's turn, not the consumer's.
+
+    `send_nowait` wakes the metrics consumer but does not run it. An attempt that
+    publishes an event and then raises in the same event-loop turn would otherwise
+    reach the terminal branch of `_main_task` with the budget still exhausted, and
+    a plugin can emit its last usage event during teardown, just before the socket
+    error propagates. Enters the event-delivering attempt already at max_retry.
+    """
+    usage_only = SpeechEvent(
+        type=SpeechEventType.RECOGNITION_USAGE,
+        recognition_usage=RecognitionUsage(audio_duration=5.0),
+    )
+    assert await _survives(
+        usage_only,
+        drops=12,
+        silent_runs=DEFAULT_API_CONNECT_OPTIONS.max_retry,
+        yield_after_emit=False,
+    )