diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 62de50a87c8..6f140266fca 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2732,7 +2732,7 @@ class PrismaClient: self._db_reconnect_circuit_breaker_action = "exit" self._db_reconnect_circuit_breaker_opened: bool = False self._db_reconnect_breaker_attempts: Deque[float] = deque( - maxlen=self._db_reconnect_circuit_breaker_max_attempts + 1 + maxlen=self._db_reconnect_circuit_breaker_max_attempts ) self._db_reconnect_breaker_failures: Deque[float] = deque( maxlen=self._db_reconnect_circuit_breaker_max_failures @@ -4377,7 +4377,7 @@ class PrismaClient: return "engine_process_death_threshold_exceeded" if counts["failures"] >= self._db_reconnect_circuit_breaker_max_failures: return "reconnect_failure_threshold_exceeded" - if counts["attempts"] > self._db_reconnect_circuit_breaker_max_attempts: + if counts["attempts"] >= self._db_reconnect_circuit_breaker_max_attempts: return "reconnect_attempt_threshold_exceeded" return None @@ -4455,6 +4455,8 @@ class PrismaClient: elif event_type == "failure": self._db_reconnect_breaker_failures.append(now) else: + self._db_reconnect_breaker_failures.clear() + self._db_reconnect_breaker_engine_deaths.clear() return False counts = self._get_reconnect_breaker_counts() diff --git a/tests/litellm/proxy/test_prisma_engine_watchdog.py b/tests/litellm/proxy/test_prisma_engine_watchdog.py index 00774dfde9d..11b68e2b7a4 100644 --- a/tests/litellm/proxy/test_prisma_engine_watchdog.py +++ b/tests/litellm/proxy/test_prisma_engine_watchdog.py @@ -15,6 +15,7 @@ Covers: """ import asyncio +from collections import deque import os import threading import time @@ -657,6 +658,54 @@ async def test_reconnect_circuit_breaker_stays_closed_on_transient_success( engine_client._terminate_for_reconnect_breaker.assert_not_called() +@pytest.mark.asyncio +async def test_reconnect_circuit_breaker_opens_on_exact_attempt_threshold( + engine_client, +): + engine_client._db_reconnect_circuit_breaker_enabled = True + engine_client._db_reconnect_circuit_breaker_max_attempts = 2 + engine_client._db_reconnect_circuit_breaker_max_failures = 100 + engine_client._db_reconnect_circuit_breaker_max_engine_deaths = 100 + engine_client._db_reconnect_breaker_attempts = deque(maxlen=2) + engine_client._db_reconnect_breaker_failures = deque(maxlen=100) + engine_client._db_reconnect_breaker_engine_deaths = deque(maxlen=100) + engine_client._db_reconnect_cooldown_seconds = 0 + engine_client._run_reconnect_cycle = AsyncMock(return_value=None) + engine_client._terminate_for_reconnect_breaker = MagicMock() + + first_result = await engine_client._attempt_reconnect_inside_lock( + force=True, + reason="db_health_watchdog_connection_error", + timeout_seconds=5.0, + ) + second_result = await engine_client._attempt_reconnect_inside_lock( + force=True, + reason="db_health_watchdog_connection_error", + timeout_seconds=5.0, + ) + + assert first_result is True + assert second_result is False + assert len(engine_client._db_reconnect_breaker_attempts) == 2 + assert engine_client._db_reconnect_circuit_breaker_opened is True + engine_client._terminate_for_reconnect_breaker.assert_called_once() + + +def test_reconnect_circuit_breaker_success_clears_failure_state(engine_client): + engine_client._db_reconnect_circuit_breaker_enabled = True + engine_client._db_reconnect_breaker_failures.append(100.0) + engine_client._db_reconnect_breaker_engine_deaths.append(100.0) + + result = engine_client._record_reconnect_breaker_event( + event_type="success", + reason="db_health_watchdog_connection_error", + ) + + assert result is False + assert len(engine_client._db_reconnect_breaker_failures) == 0 + assert len(engine_client._db_reconnect_breaker_engine_deaths) == 0 + + @pytest.mark.asyncio async def test_reconnect_circuit_breaker_log_action_does_not_skip_reconnect( engine_client, @@ -702,6 +751,6 @@ def test_reconnect_circuit_breaker_env_vars_are_respected(mock_proxy_logging): assert client._db_reconnect_circuit_breaker_max_failures == 2 assert client._db_reconnect_circuit_breaker_max_engine_deaths == 1 assert client._db_reconnect_circuit_breaker_action == "log" - assert client._db_reconnect_breaker_attempts.maxlen == 5 + assert client._db_reconnect_breaker_attempts.maxlen == 4 assert client._db_reconnect_breaker_failures.maxlen == 2 assert client._db_reconnect_breaker_engine_deaths.maxlen == 1