Skip to content

fix(stt): raise retryable APIConnectionError on mid-send socket drops #6467

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tinalenguyen merged 1 commit into from
Jul 17, 2026
Merged

fix(stt): raise retryable APIConnectionError on mid-send socket drops #6467

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 34 additions & 27 deletions livekit-agents/livekit/agents/inference/stt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -817,33 +817,40 @@ async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
samples_per_channel=self._opts.sample_rate // 20, # 50ms
)

async for ev in self._input_ch:
frames: list[rtc.AudioFrame] = []
if isinstance(ev, rtc.AudioFrame):
if vad_stream is not None:
vad_stream.push_frame(ev)
frames.extend(audio_bstream.push(ev.data))
elif isinstance(ev, self._FlushSentinel):
frames.extend(audio_bstream.flush())

for frame in frames:
self._speech_duration += frame.duration
audio_bytes = frame.data.tobytes()
base64_audio = base64.b64encode(audio_bytes).decode("utf-8")
audio_msg = {
"type": "input_audio",
"audio": base64_audio,
}
await ws.send_str(json.dumps(audio_msg))

if vad_stream is not None:
vad_stream.end_input()

closing_ws = True
finalize_msg = {
"type": "session.finalize",
}
await ws.send_str(json.dumps(finalize_msg))
try:
async for ev in self._input_ch:
frames: list[rtc.AudioFrame] = []
if isinstance(ev, rtc.AudioFrame):
if vad_stream is not None:
vad_stream.push_frame(ev)
frames.extend(audio_bstream.push(ev.data))
elif isinstance(ev, self._FlushSentinel):
frames.extend(audio_bstream.flush())

for frame in frames:
self._speech_duration += frame.duration
audio_bytes = frame.data.tobytes()
base64_audio = base64.b64encode(audio_bytes).decode("utf-8")
audio_msg = {
"type": "input_audio",
"audio": base64_audio,
}
await ws.send_str(json.dumps(audio_msg))

if vad_stream is not None:
vad_stream.end_input()

closing_ws = True
finalize_msg = {
"type": "session.finalize",
}
await ws.send_str(json.dumps(finalize_msg))
except (aiohttp.ClientError, ConnectionError) as e:
if closing_ws or http_session.closed:
return
raise APIConnectionError(
"LiveKit Inference STT connection closed unexpectedly"
) from e

@utils.log_exceptions(logger=logger)
async def vad_task(ws: aiohttp.ClientWebSocketResponse, stream: vad.VADStream) -> None:
Expand Down
48 changes: 27 additions & 21 deletions livekit-plugins/livekit-plugins-assemblyai/livekit/plugins/assemblyai/stt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@

from livekit.agents import (
DEFAULT_API_CONNECT_OPTIONS,
APIConnectionError,
APIConnectOptions,
APIStatusError,
LanguageCode,
Expand Down Expand Up @@ -685,27 +686,32 @@ async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
# forward inputs to AssemblyAI
# if we receive a close message, signal it to AssemblyAI and break.
# the recv task will then make sure to process the remaining audio and stop
async for data in self._input_ch:
if isinstance(data, self._FlushSentinel):
frames = audio_bstream.flush()
else:
frames = audio_bstream.write(data.data.tobytes())

for frame in frames:
if not anchored:
# Anchor the stream's wall-clock to the moment just
# before the first frame is sent — aligned with the
# server's stream-relative zero used by
# SpeechStarted.timestamp.
self.start_time = time.time()
anchored = True
self._speech_duration += frame.duration
await ws.send_bytes(frame.data.tobytes())
self._last_frame_sent_at = time.time()

closing_ws = True
logger.debug("AssemblyAI sending close message session=%s", self._session_id)
await ws.send_str(SpeechStream._CLOSE_MSG)
try:
async for data in self._input_ch:
if isinstance(data, self._FlushSentinel):
frames = audio_bstream.flush()
else:
frames = audio_bstream.write(data.data.tobytes())

for frame in frames:
if not anchored:
# Anchor the stream's wall-clock to the moment just
# before the first frame is sent — aligned with the
# server's stream-relative zero used by
# SpeechStarted.timestamp.
self.start_time = time.time()
anchored = True
self._speech_duration += frame.duration
await ws.send_bytes(frame.data.tobytes())
self._last_frame_sent_at = time.time()

closing_ws = True
logger.debug("AssemblyAI sending close message session=%s", self._session_id)
await ws.send_str(SpeechStream._CLOSE_MSG)
except (aiohttp.ClientError, ConnectionError) as e:
if closing_ws or self._session.closed:
return
raise APIConnectionError("AssemblyAI connection closed unexpectedly") from e

async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
nonlocal closing_ws
Expand Down
30 changes: 18 additions & 12 deletions livekit-plugins/livekit-plugins-baseten/livekit/plugins/baseten/stt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@

from livekit.agents import (
DEFAULT_API_CONNECT_OPTIONS,
APIConnectionError,
APIConnectOptions,
APIStatusError,
LanguageCode,
Expand Down Expand Up @@ -341,18 +342,23 @@ async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
samples_per_channel=samples_per_buffer,
)

async for data in self._input_ch:
if isinstance(data, self._FlushSentinel):
frames = audio_bstream.flush()
else:
frames = audio_bstream.write(data.data.tobytes())

for frame in frames:
if len(frame.data) % 2 != 0:
logger.warning("Frame data size not aligned to float32 (multiple of 4)")

int16_array = np.frombuffer(frame.data, dtype=np.int16)
await ws.send_bytes(int16_array.tobytes())
try:
async for data in self._input_ch:
if isinstance(data, self._FlushSentinel):
frames = audio_bstream.flush()
else:
frames = audio_bstream.write(data.data.tobytes())

for frame in frames:
if len(frame.data) % 2 != 0:
logger.warning("Frame data size not aligned to float32 (multiple of 4)")

int16_array = np.frombuffer(frame.data, dtype=np.int16)
await ws.send_bytes(int16_array.tobytes())
except (aiohttp.ClientError, ConnectionError) as e:
if closing_ws or self._session.closed:
return
raise APIConnectionError("Baseten connection closed unexpectedly") from e

async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
nonlocal closing_ws
Expand Down
40 changes: 24 additions & 16 deletions ...ns-cartesia/livekit/plugins/cartesia/_recognize_streams/auto_finalize_recognize_stream.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -254,23 +254,31 @@ async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
samples_per_channel=samples_per_chunk,
)

async for data in self._input_ch:
if isinstance(data, rtc.AudioFrame):
for frame in audio_bstream.write(data.data.tobytes()):
self._speech_duration += frame.duration
await ws.send_bytes(frame.data.tobytes())
elif isinstance(data, self._FlushSentinel):
if not self._input_ch.closed:
logger.warning(
"Cartesia STT stream.flush() was ignored. See https://docs.cartesia.ai/use-the-api/compare-stt-endpoints for details."
)

for frame in audio_bstream.flush():
self._speech_duration += frame.duration
await ws.send_bytes(frame.data.tobytes())
try:
async for data in self._input_ch:
if isinstance(data, rtc.AudioFrame):
for frame in audio_bstream.write(data.data.tobytes()):
self._speech_duration += frame.duration
await ws.send_bytes(frame.data.tobytes())
elif isinstance(data, self._FlushSentinel):
if not self._input_ch.closed:
logger.warning(
"Cartesia STT stream.flush() was ignored. See https://docs.cartesia.ai/use-the-api/compare-stt-endpoints for details."
)

self._closing_ws = True
await ws.send_str('{"type":"close"}')
for frame in audio_bstream.flush():
self._speech_duration += frame.duration
await ws.send_bytes(frame.data.tobytes())

self._closing_ws = True
await ws.send_str('{"type":"close"}')
except (aiohttp.ClientError, ConnectionError) as e:
if self._closing_ws or self._session.closed:
return
raise APIConnectionError(
message="Cartesia STT connection closed unexpectedly",
retryable=True,
) from e

@utils.log_exceptions(logger=logger)
async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
Expand Down
50 changes: 29 additions & 21 deletions ...t-plugins-cartesia/livekit/plugins/cartesia/_recognize_streams/legacy_recognize_stream.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,27 +197,35 @@ async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
samples_per_channel=samples_per_chunk,
)

async for data in self._input_ch:
frames: list[rtc.AudioFrame | LegacyRecognizeStream._FlushSentinel] = []
if isinstance(data, rtc.AudioFrame):
frames.extend(audio_bstream.write(data.data.tobytes()))
elif isinstance(data, self._FlushSentinel):
frames.extend(audio_bstream.flush())
frames.append(data)

for frame in frames:
if isinstance(frame, self._FlushSentinel):
await ws.send_str("finalize")
else:
self._speech_duration += frame.duration
await ws.send_bytes(frame.data.tobytes())

for frame in audio_bstream.flush():
self._speech_duration += frame.duration
await ws.send_bytes(frame.data.tobytes())

closing_ws = True
await ws.send_str("close")
try:
async for data in self._input_ch:
frames: list[rtc.AudioFrame | LegacyRecognizeStream._FlushSentinel] = []
if isinstance(data, rtc.AudioFrame):
frames.extend(audio_bstream.write(data.data.tobytes()))
elif isinstance(data, self._FlushSentinel):
frames.extend(audio_bstream.flush())
frames.append(data)

for frame in frames:
if isinstance(frame, self._FlushSentinel):
await ws.send_str("finalize")
else:
self._speech_duration += frame.duration
await ws.send_bytes(frame.data.tobytes())

for frame in audio_bstream.flush():
self._speech_duration += frame.duration
await ws.send_bytes(frame.data.tobytes())

closing_ws = True
await ws.send_str("close")
except (aiohttp.ClientError, ConnectionError) as e:
if closing_ws or self._session.closed:
return
raise APIConnectionError(
message="Cartesia STT connection closed unexpectedly",
retryable=True,
) from e

@utils.log_exceptions(logger=logger)
async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
Expand Down
55 changes: 30 additions & 25 deletions livekit-plugins/livekit-plugins-elevenlabs/livekit/plugins/elevenlabs/stt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -421,34 +421,39 @@ async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
)

has_ended = False
async for data in self._input_ch:
# Write audio bytes to buffer and get 50ms frames
frames: list[rtc.AudioFrame] = []
if isinstance(data, rtc.AudioFrame):
frames.extend(audio_bstream.write(data.data.tobytes()))
elif isinstance(data, self._FlushSentinel):
frames.extend(audio_bstream.flush())
has_ended = True

for frame in frames:
self._audio_duration_collector.push(frame.duration)
audio_b64 = base64.b64encode(frame.data.tobytes()).decode("utf-8")
await ws.send_str(
json.dumps(
{
"message_type": "input_audio_chunk",
"audio_base_64": audio_b64,
"commit": False,
"sample_rate": self._opts.sample_rate,
}
try:
async for data in self._input_ch:
# Write audio bytes to buffer and get 50ms frames
frames: list[rtc.AudioFrame] = []
if isinstance(data, rtc.AudioFrame):
frames.extend(audio_bstream.write(data.data.tobytes()))
elif isinstance(data, self._FlushSentinel):
frames.extend(audio_bstream.flush())
has_ended = True

for frame in frames:
self._audio_duration_collector.push(frame.duration)
audio_b64 = base64.b64encode(frame.data.tobytes()).decode("utf-8")
await ws.send_str(
json.dumps(
{
"message_type": "input_audio_chunk",
"audio_base_64": audio_b64,
"commit": False,
"sample_rate": self._opts.sample_rate,
}
)
)
)

if has_ended:
self._audio_duration_collector.flush()
has_ended = False
if has_ended:
self._audio_duration_collector.flush()
has_ended = False

closing_ws = True
closing_ws = True
except (aiohttp.ClientError, ConnectionError) as e:
if closing_ws or self._session.closed:
return
raise APIConnectionError("ElevenLabs STT connection closed unexpectedly") from e

@utils.log_exceptions(logger=logger)
async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
Expand Down
Loading