Detect quick connection failures in WebsocketService to prevent infinite reconnection loops

When a WebSocket server accepts the handshake but immediately closes the
connection (e.g. invalid API key returning close code 1008), the existing
exponential backoff does not help because the handshake keeps succeeding.
This tracks how long each connection survives and emits a non-fatal
ErrorFrame after 3 consecutive sub-5s failures, allowing ServiceSwitcher
failover instead of killing the pipeline.

Fixes #3711
This commit is contained in:
Mark Backman
2026-03-30 12:16:29 -04:00
parent 87b8f38a48
commit 86a16d53bc
2 changed files with 221 additions and 0 deletions

View File

@@ -7,6 +7,7 @@
"""Base websocket service with automatic reconnection and error handling."""
import asyncio
import time
from abc import ABC, abstractmethod
from typing import Awaitable, Callable, Optional
@@ -27,6 +28,13 @@ class WebsocketService(ABC):
Subclasses implement service-specific connection and message handling logic.
"""
# Rapid failure detection: when a server accepts the WebSocket handshake but
# immediately closes the connection (e.g. invalid API key, policy rejection),
# exponential backoff won't help because the handshake keeps succeeding. We
# detect this by tracking how long the connection survives after being established.
_MIN_STABLE_CONNECTION_DURATION = 5.0 # seconds
_MAX_CONSECUTIVE_QUICK_FAILURES = 3
def __init__(self, *, reconnect_on_error: bool = True, **kwargs):
"""Initialize the websocket service.
@@ -38,6 +46,8 @@ class WebsocketService(ABC):
self._reconnect_on_error = reconnect_on_error
self._reconnect_in_progress: bool = False
self._disconnecting: bool = False
self._quick_failure_count: int = 0
self._last_connect_time: float = 0.0
async def _verify_connection(self) -> bool:
"""Verify the websocket connection is active and responsive.
@@ -86,6 +96,7 @@ class WebsocketService(ABC):
logger.warning(f"{self} reconnecting, attempt {attempt}")
if await self._reconnect_websocket(attempt):
logger.info(f"{self} reconnected successfully on attempt {attempt}")
self._last_connect_time = time.monotonic()
return True
except Exception as e:
last_exception = e
@@ -145,6 +156,31 @@ class WebsocketService(ABC):
logger.debug(f"{self} receive loop ended during disconnect")
return False
# Check if the connection died too quickly after being established. This
# catches cases where the handshake succeeds but the server immediately
# closes (e.g. invalid API key). Exponential backoff won't help here
# because the handshake keeps succeeding — we need to stop the loop.
if self._last_connect_time > 0:
connection_duration = time.monotonic() - self._last_connect_time
if connection_duration < self._MIN_STABLE_CONNECTION_DURATION:
self._quick_failure_count += 1
logger.warning(
f"{self} connection lasted only {connection_duration:.1f}s "
f"({self._quick_failure_count}/{self._MAX_CONSECUTIVE_QUICK_FAILURES} "
f"consecutive quick failures)"
)
if self._quick_failure_count >= self._MAX_CONSECUTIVE_QUICK_FAILURES:
msg = (
f"{self} connection failed {self._MAX_CONSECUTIVE_QUICK_FAILURES} "
f"times immediately after connecting"
)
logger.error(msg)
await report_error(ErrorFrame(msg))
return False
else:
# Connection was stable — reset the counter.
self._quick_failure_count = 0
# Log the message
logger.warning(error_message)
@@ -168,6 +204,7 @@ class WebsocketService(ABC):
report_error: Callback function to report connection errors.
"""
while True:
self._last_connect_time = time.monotonic()
try:
await self._receive_messages()
# _receive_messages() returned normally. This happens when the websocket
@@ -205,6 +242,7 @@ class WebsocketService(ABC):
additional setup required.
"""
self._disconnecting = False
self._quick_failure_count = 0
async def _disconnect(self):
"""Disconnect from the service and set disconnecting flag.

View File

@@ -0,0 +1,183 @@
#
# Copyright (c) 2024-2026, Daily
#
# SPDX-License-Identifier: BSD 2-Clause License
#
"""Tests for WebsocketService quick failure detection."""
import time
from unittest.mock import AsyncMock, patch
import pytest
from websockets.exceptions import ConnectionClosedError
from websockets.frames import Close
from pipecat.frames.frames import ErrorFrame
from pipecat.services.websocket_service import WebsocketService
class ConcreteWebsocketService(WebsocketService):
"""Minimal concrete implementation for testing."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self._receive_messages_impl: AsyncMock | None = None
async def _connect_websocket(self):
pass
async def _disconnect_websocket(self):
pass
async def _receive_messages(self):
if self._receive_messages_impl:
await self._receive_messages_impl()
@pytest.fixture
def service():
return ConcreteWebsocketService()
@pytest.fixture
def report_error():
return AsyncMock()
@pytest.mark.asyncio
async def test_quick_failures_emit_error(service, report_error):
"""Connections that fail immediately after being established should emit an error
after MAX_CONSECUTIVE_QUICK_FAILURES consecutive quick failures."""
call_count = 0
async def fail_immediately():
nonlocal call_count
call_count += 1
raise ConnectionClosedError(Close(1008, "Invalid API key"), None)
service._receive_messages_impl = AsyncMock(side_effect=fail_immediately)
# Mock _try_reconnect to succeed (handshake passes but connection dies right away)
service._try_reconnect = AsyncMock(return_value=True)
await service._receive_task_handler(report_error)
# Should have called _receive_messages MAX_RAPID_FAILURES times
assert call_count == service._MAX_CONSECUTIVE_QUICK_FAILURES
# Should have emitted a fatal error
report_error.assert_called_once()
error_frame = report_error.call_args[0][0]
assert isinstance(error_frame, ErrorFrame)
assert error_frame.fatal is False
assert "failed 3 times immediately after connecting" in error_frame.error
@pytest.mark.asyncio
async def test_stable_connection_resets_quick_failure_counter(service, report_error):
"""A connection that survives beyond the threshold should reset the quick failure counter."""
call_count = 0
async def fail_then_stable_then_fail():
nonlocal call_count
call_count += 1
if call_count <= 2:
# First two calls: quick failures
raise ConnectionClosedError(Close(1006, "Abnormal closure"), None)
elif call_count == 3:
# Third call: simulate a stable connection by advancing time past threshold
raise ConnectionClosedError(Close(1006, "Abnormal closure"), None)
else:
# Fourth and beyond: quick failures again
raise ConnectionClosedError(Close(1006, "Abnormal closure"), None)
service._receive_messages_impl = AsyncMock(side_effect=fail_then_stable_then_fail)
service._try_reconnect = AsyncMock(return_value=True)
# Patch time.monotonic to control timing
base_time = 1000.0
time_values = iter(
[
# Call 1: set _last_connect_time
base_time,
# Call 1: check in _maybe_try_reconnect (rapid: 0s elapsed)
base_time,
# Call 2: set _last_connect_time
base_time + 1.0,
# Call 2: check in _maybe_try_reconnect (rapid: 0s elapsed)
base_time + 1.0,
# Call 3: set _last_connect_time
base_time + 2.0,
# Call 3: check in _maybe_try_reconnect (stable: 10s elapsed)
base_time + 12.0,
# Call 4: set _last_connect_time
base_time + 13.0,
# Call 4: check in _maybe_try_reconnect (rapid: 0s elapsed)
base_time + 13.0,
# Call 5: set _last_connect_time
base_time + 14.0,
# Call 5: check in _maybe_try_reconnect (rapid: 0s elapsed)
base_time + 14.0,
# Call 6: set _last_connect_time
base_time + 15.0,
# Call 6: check in _maybe_try_reconnect (rapid: 0s elapsed)
base_time + 15.0,
]
)
with patch("pipecat.services.websocket_service.time") as mock_time:
mock_time.monotonic = lambda: next(time_values)
await service._receive_task_handler(report_error)
# After the stable connection (call 3), counter resets to 0.
# Then calls 4, 5, 6 are quick failures (counter: 1, 2, 3) -> error emitted
assert call_count == 6
report_error.assert_called_once()
error_frame = report_error.call_args[0][0]
assert error_frame.fatal is False
@pytest.mark.asyncio
async def test_graceful_close_counts_toward_quick_failures(service, report_error):
"""A _receive_messages that returns normally (graceful close) should also count
toward quick failures if it happens immediately."""
call_count = 0
async def return_immediately():
nonlocal call_count
call_count += 1
service._receive_messages_impl = AsyncMock(side_effect=return_immediately)
service._try_reconnect = AsyncMock(return_value=True)
await service._receive_task_handler(report_error)
assert call_count == service._MAX_CONSECUTIVE_QUICK_FAILURES
report_error.assert_called_once()
error_frame = report_error.call_args[0][0]
assert isinstance(error_frame, ErrorFrame)
assert error_frame.fatal is False
@pytest.mark.asyncio
async def test_connect_resets_quick_failure_counter(service):
"""Calling _connect() should reset the quick failure counter."""
service._quick_failure_count = 5
await service._connect()
assert service._quick_failure_count == 0
@pytest.mark.asyncio
async def test_intentional_disconnect_skips_quick_failure_logic(service, report_error):
"""When _disconnecting is True, quick failure detection should not run."""
service._disconnecting = True
service._quick_failure_count = 0
service._last_connect_time = time.monotonic()
result = await service._maybe_try_reconnect("test error", report_error)
assert result is False
# Counter should not have been incremented
assert service._quick_failure_count == 0
# No error frame should have been emitted
report_error.assert_not_called()