From 8417b962a2bb3f0c4f3dbaf706dc822b380a381e Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 7 Jul 2026 09:39:10 +0300 Subject: [PATCH] fix(proxy): recover Prisma DB reconnect loop when client is disconnected Once the active Prisma client is in the disconnected state, every DB call raises ClientNotConnectedError. The reconnect machinery was supposed to recover from this, but _get_engine_pid() inspected the broken client via prisma's _engine property, which re-raises that same error, so recreate_prisma_client failed before it could build a replacement client and the proxy looped on failed reconnects forever (issue #28322 showed 1486+ consecutive failures over 30 days with zero recoveries) Guard both _get_engine_pid implementations with is_connected() so a disconnected client reads as "no engine" (pid 0) and the recreate path proceeds to construct and connect a fresh client --- litellm/proxy/db/prisma_client.py | 11 +++++- litellm/proxy/utils.py | 15 ++++++- tests/test_litellm/proxy/conftest.py | 25 ++++++++++++ .../proxy/db/test_prisma_client.py | 35 +++++++++++++++++ .../db/test_prisma_planned_engine_restart.py | 1 + .../proxy/db/test_prisma_self_heal.py | 39 ++++++++++++++++++- .../test_prisma_client_engine_watcher.py | 10 +++++ 7 files changed, 132 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 4042755f80d..fbccb8a726c 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -137,8 +137,17 @@ class PrismaWrapper: self.on_engine_replaced: Callable[[], None] | None = None def _get_engine_pid(self) -> int: - """Get the PID of the current Prisma engine subprocess, or 0 if unavailable.""" + """Get the PID of the current Prisma engine subprocess, or 0 if unavailable. + + Must never raise: it runs inside the reconnect path, where the client + may be in any broken state. Prisma's ``_engine`` is a property that + raises ``ClientNotConnectedError`` on a disconnected client; if that + escaped here, ``recreate_prisma_client`` would fail before it could + build a replacement client and the reconnect loop could never recover. + """ try: + if self._original_prisma.is_connected() is not True: + return 0 engine = self._original_prisma._engine process = getattr(engine, "process", None) if engine is not None else None if process is not None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2880eef6908..d7649b524aa 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4097,11 +4097,22 @@ class PrismaClient: raise e def _get_engine_pid(self) -> int: + """Get the PID of the writer's engine subprocess, or 0 if unavailable. + + Must never raise: prisma's ``_engine`` property raises + ``ClientNotConnectedError`` on a disconnected client, and an exception + escaping from the reconnect path would leave it unable to recover. + """ try: - engine = self.db._original_prisma._engine # type: ignore[attr-defined] + prisma_obj = self.writer_db._original_prisma + if prisma_obj.is_connected() is not True: + return 0 + engine = prisma_obj._engine process = getattr(engine, "process", None) if engine is not None else None if process is not None: - return process.pid + pid = process.pid + if isinstance(pid, int): + return pid except (AttributeError, TypeError): pass return 0 diff --git a/tests/test_litellm/proxy/conftest.py b/tests/test_litellm/proxy/conftest.py index 607315eb246..1d71035b67f 100644 --- a/tests/test_litellm/proxy/conftest.py +++ b/tests/test_litellm/proxy/conftest.py @@ -20,6 +20,31 @@ _PROXY_MODULE_GLOBALS_TO_ISOLATE = ( ) +class StubClientNotConnectedError(Exception): + pass + + +class DisconnectedPrisma: + """Mimics prisma-client-py after disconnect(): ``is_connected()`` is False + and the ``_engine`` property raises ``ClientNotConnectedError``.""" + + def is_connected(self) -> bool: + return False + + @property + def _engine(self) -> None: + raise StubClientNotConnectedError( + "Client is not connected to the query engine, you must call `connect()` " + "before attempting to query data." + ) + + +@pytest.fixture +def disconnected_prisma() -> DisconnectedPrisma: + """A stand-in for a Prisma client wedged in the disconnected state.""" + return DisconnectedPrisma() + + @pytest.fixture(autouse=True) def _isolate_proxy_module_globals(): """ diff --git a/tests/test_litellm/proxy/db/test_prisma_client.py b/tests/test_litellm/proxy/db/test_prisma_client.py index 397e3f36e41..eeaf726941f 100644 --- a/tests/test_litellm/proxy/db/test_prisma_client.py +++ b/tests/test_litellm/proxy/db/test_prisma_client.py @@ -92,6 +92,7 @@ async def test_recreate_prisma_client_kills_old_engine_on_disconnect_failure( """When disconnect() fails, recreate_prisma_client must SIGTERM/SIGKILL the old engine PID.""" mock_prisma = AsyncMock() mock_prisma.disconnect.side_effect = Exception("engine hung") + mock_prisma.is_connected = MagicMock(return_value=True) # Simulate engine subprocess with a known PID mock_engine = MagicMock() @@ -122,6 +123,7 @@ async def test_recreate_prisma_client_skips_kill_on_successful_disconnect( ): """When disconnect() succeeds, no kill should be attempted.""" mock_prisma = AsyncMock() + mock_prisma.is_connected = MagicMock(return_value=True) mock_prisma.disconnect.return_value = None wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False) @@ -142,6 +144,7 @@ async def test_recreate_prisma_client_handles_missing_engine_pid( ): """When engine PID is unavailable (no _engine attr), kill is skipped gracefully.""" mock_prisma = AsyncMock() + mock_prisma.is_connected = MagicMock(return_value=True) mock_prisma.disconnect.side_effect = Exception("engine hung") mock_prisma._engine = None # No engine subprocess @@ -158,3 +161,35 @@ async def test_recreate_prisma_client_handles_missing_engine_pid( mock_kill.assert_not_called() # PID was 0, kill skipped mock_new_prisma.connect.assert_awaited_once() + + +def test_get_engine_pid_returns_zero_for_disconnected_client(disconnected_prisma): + """A disconnected client must read as "no engine" instead of raising, + otherwise the reconnect path can never recover.""" + wrapper = PrismaWrapper( + original_prisma=disconnected_prisma, iam_token_db_auth=False + ) + + assert wrapper._get_engine_pid() == 0 + + +@pytest.mark.asyncio +async def test_recreate_prisma_client_recovers_from_disconnected_client( + mock_prisma_binary, disconnected_prisma +): + """recreate_prisma_client must still build a replacement client when the + current one is disconnected.""" + wrapper = PrismaWrapper( + original_prisma=disconnected_prisma, iam_token_db_auth=False + ) + + mock_new_prisma = AsyncMock() + mock_prisma_binary.Prisma.return_value = mock_new_prisma + + with patch("os.kill") as mock_kill: + result = await wrapper.recreate_prisma_client("postgresql://new") + + assert result is True + mock_kill.assert_not_called() + assert wrapper._original_prisma is mock_new_prisma + mock_new_prisma.connect.assert_awaited_once() diff --git a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py index 5e74004cc0b..9b382a41964 100644 --- a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py +++ b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py @@ -42,6 +42,7 @@ def mock_prisma_binary(): def _make_wrapper(engine_pid: int = 111, iam: bool = False) -> PrismaWrapper: mock_prisma = MagicMock() mock_prisma.connect = AsyncMock() + mock_prisma.is_connected = MagicMock(return_value=True) mock_prisma._engine = MagicMock() mock_prisma._engine.process.pid = engine_pid return PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=iam) diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index 35ef0a965f3..7f723fa3ae0 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -20,7 +20,7 @@ def mock_prisma_binary(): """Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests.""" mock_module = MagicMock() with patch.dict(sys.modules, {"prisma": mock_module}): - yield + yield mock_module @pytest.fixture @@ -515,6 +515,43 @@ async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect( assert client._engine_confirmed_dead is True +@pytest.mark.asyncio +async def test_heavy_reconnect_recovers_from_disconnected_prisma_client( + mock_proxy_logging, mock_prisma_binary, disconnected_prisma +): + """Once the active Prisma client is in the disconnected state, every DB + call raises ClientNotConnectedError. The heavy reconnect path is the only + way out, so it must not re-raise that same error while inspecting the + broken client; otherwise `recreate_prisma_client` fails before it can + build a replacement and the proxy loops on failed reconnects forever. + + The full real reconnect path (attempt_db_reconnect -> _run_reconnect_cycle + -> recreate_prisma_client) must succeed from that wedged state. + """ + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db._original_prisma = disconnected_prisma + client._engine_confirmed_dead = True + client._start_engine_watcher = AsyncMock() + + replacement = MagicMock() + replacement.connect = AsyncMock() + mock_prisma_binary.Prisma.return_value = replacement + + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + result = await client.attempt_db_reconnect( + reason="unit_test_disconnected_client", + force=True, + ) + + assert result is True + assert client.db._original_prisma is replacement + replacement.connect.assert_awaited_once() + assert client._consecutive_reconnect_failures == 0 + assert client._engine_confirmed_dead is False + + @pytest.mark.asyncio async def test_db_health_watchdog_should_reconnect_degraded_writer( mock_proxy_logging, diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_engine_watcher.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_engine_watcher.py index 2fedd6bb134..25c04caabba 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_engine_watcher.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_engine_watcher.py @@ -42,6 +42,7 @@ def test_get_engine_pid_extracts_process_pid(prisma_client: PrismaClient) -> Non fake_engine.process = MagicMock() fake_engine.process.pid = 4242 prisma_client.db._original_prisma = MagicMock() + prisma_client.db._original_prisma.is_connected = MagicMock(return_value=True) prisma_client.db._original_prisma._engine = fake_engine actual = { "pid": prisma_client._get_engine_pid(), @@ -58,6 +59,15 @@ def test_get_engine_pid_returns_zero_when_engine_attr_missing( assert prisma_client._get_engine_pid() == 0 +def test_get_engine_pid_returns_zero_when_client_disconnected( + prisma_client: PrismaClient, disconnected_prisma +) -> None: + """The reconnect path calls this on an arbitrarily-broken client; it must + report "no engine" instead of re-raising ClientNotConnectedError.""" + prisma_client.db._original_prisma = disconnected_prisma + assert prisma_client._get_engine_pid() == 0 + + def test_is_engine_alive_true_when_pid_zero(prisma_client: PrismaClient) -> None: prisma_client._engine_pid = 0 pinned = {