diff --git a/src/pipecat/services/websocket_service.py b/src/pipecat/services/websocket_service.py index 85e5b2db7..e5449a539 100644 --- a/src/pipecat/services/websocket_service.py +++ b/src/pipecat/services/websocket_service.py @@ -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. diff --git a/tests/test_websocket_service.py b/tests/test_websocket_service.py new file mode 100644 index 000000000..2f180b987 --- /dev/null +++ b/tests/test_websocket_service.py @@ -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()