From fca08e8acc4bffc922efd6bd4bee2de5f8151312 Mon Sep 17 00:00:00 2001 From: Darien Kindlund Date: Fri, 27 Feb 2026 22:29:25 -0500 Subject: [PATCH] fix: escalate to heavy Prisma reconnect after consecutive lightweight failures (#22211) When the Prisma query engine process is alive but not accepting connections (e.g., startup race condition in containerized deployments), lightweight reconnects (disconnect + connect) will never succeed. The health watchdog retries indefinitely without escalating to a full Prisma client recreation. Adds a consecutive failure counter that triggers a heavy reconnect (full Prisma client and engine recreation) after 3 consecutive lightweight reconnect failures (configurable via PRISMA_RECONNECT_ESCALATION_THRESHOLD env var). Co-authored-by: Claude Opus 4.6 (1M context) --- litellm/proxy/utils.py | 22 +++++- .../proxy/test_prisma_engine_watchdog.py | 72 +++++++++++++++++++ 2 files changed, 93 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5e0d5336aa..afcdd9d0c5 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2304,6 +2304,10 @@ class PrismaClient: 0.0, float(os.getenv("PRISMA_AUTH_RECONNECT_LOCK_TIMEOUT_SECONDS", "0.1")), ) + self._consecutive_reconnect_failures: int = 0 + self._reconnect_escalation_threshold: int = max( + 1, int(os.getenv("PRISMA_RECONNECT_ESCALATION_THRESHOLD", "3")) + ) self._engine_pidfd: int = -1 self._engine_pid: int = 0 self._watching_engine: bool = False @@ -3917,6 +3921,19 @@ class PrismaClient: ) return False + # Escalate to heavy reconnect after consecutive lightweight failures. + # When the Prisma engine process is alive but not accepting connections + # (e.g., startup race condition), lightweight reconnects (disconnect + + # connect) will never succeed. Force a full Prisma client recreation + # to recover from this state. + if self._consecutive_reconnect_failures >= self._reconnect_escalation_threshold: + verbose_proxy_logger.warning( + "Escalating to heavy reconnect after %d consecutive failures. reason=%s", + self._consecutive_reconnect_failures, + reason, + ) + self._engine_confirmed_dead = True + verbose_proxy_logger.warning( "Attempting Prisma DB reconnect. reason=%s", reason ) @@ -3925,12 +3942,15 @@ class PrismaClient: try: await self._run_reconnect_cycle(timeout_seconds=timeout_seconds) reconnect_succeeded = True + self._consecutive_reconnect_failures = 0 verbose_proxy_logger.info( "Prisma DB reconnect succeeded. reason=%s", reason ) except Exception as reconnect_err: + self._consecutive_reconnect_failures += 1 verbose_proxy_logger.error( - "Prisma DB reconnect failed. reason=%s error=%s", + "Prisma DB reconnect failed (%d consecutive). reason=%s error=%s", + self._consecutive_reconnect_failures, reason, reconnect_err, ) diff --git a/tests/litellm/proxy/test_prisma_engine_watchdog.py b/tests/litellm/proxy/test_prisma_engine_watchdog.py index 011b8002db..fb5ace0596 100644 --- a/tests/litellm/proxy/test_prisma_engine_watchdog.py +++ b/tests/litellm/proxy/test_prisma_engine_watchdog.py @@ -444,3 +444,75 @@ def test_on_engine_death_from_thread_ignores_stale_pid(engine_client): engine_client._on_engine_death_from_thread(1234) mock_create_task.assert_not_called() + + +# --------------------------------------------------------------------------- +# Reconnect escalation: lightweight -> heavy after consecutive failures +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_escalation_after_consecutive_lightweight_failures(engine_client): + """After N consecutive lightweight reconnect failures, _engine_confirmed_dead + is set to True so _run_reconnect_cycle takes the heavy reconnect path.""" + engine_client._reconnect_escalation_threshold = 3 + engine_client._consecutive_reconnect_failures = 0 + engine_client._db_reconnect_cooldown_seconds = 0 # disable cooldown for test + + # Make lightweight reconnect fail every time + engine_client.db.disconnect = AsyncMock(return_value=None) + engine_client.db.connect = AsyncMock(side_effect=Exception("connect failed")) + + # Run 3 failed reconnect attempts + for i in range(3): + result = await engine_client._attempt_reconnect_inside_lock( + force=True, reason="test", timeout_seconds=5.0 + ) + assert result is False + + assert engine_client._consecutive_reconnect_failures == 3 + + # Next attempt should escalate: _engine_confirmed_dead set to True before _run_reconnect_cycle + engine_client.db.recreate_prisma_client = AsyncMock(return_value=None) + engine_client._start_engine_watcher = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + result = await engine_client._attempt_reconnect_inside_lock( + force=True, reason="test_escalation", timeout_seconds=5.0 + ) + + # Heavy reconnect should have been attempted (recreate_prisma_client called) + engine_client.db.recreate_prisma_client.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_successful_reconnect_resets_failure_counter(engine_client): + """A successful reconnect resets _consecutive_reconnect_failures to 0.""" + engine_client._consecutive_reconnect_failures = 2 + engine_client._db_reconnect_cooldown_seconds = 0 + + # Make reconnect succeed + engine_client.db.disconnect = AsyncMock(return_value=None) + engine_client.db.connect = AsyncMock(return_value=None) + engine_client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + + result = await engine_client._attempt_reconnect_inside_lock( + force=True, reason="test", timeout_seconds=5.0 + ) + + assert result is True + assert engine_client._consecutive_reconnect_failures == 0 + + +def test_escalation_threshold_env_var(mock_proxy_logging): + """PRISMA_RECONNECT_ESCALATION_THRESHOLD env var is respected.""" + with patch.dict(os.environ, {"PRISMA_RECONNECT_ESCALATION_THRESHOLD": "5"}): + client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) + assert client._reconnect_escalation_threshold == 5 + + +def test_escalation_threshold_min_guard(mock_proxy_logging): + """Escalation threshold cannot be set below 1.""" + with patch.dict(os.environ, {"PRISMA_RECONNECT_ESCALATION_THRESHOLD": "0"}): + client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) + assert client._reconnect_escalation_threshold == 1