test(proxy): cover reconnect and refresh shutdown barriers

litellm/proxy/utils.py

@@ -6656,6 +6656,9 @@ class PrismaClient:
         for task in tasks:
             task.cancel()
         await asyncio.gather(*tasks, return_exceptions=True)
+        # Request-driven reconnects are not engine tasks. Wait for their bounded
+        # reconnect cycle to release the lock before disconnecting, otherwise a
+        # late reconnect can reopen the client after shutdown.
         async with self._db_reconnect_lock:
             pass
 

tests/integration/database/test_proxy_database_lifespan.py

@@ -14,6 +14,7 @@ from tests.integration._support.database import read_rows, scratch_database
 
 pytestmark: Final = pytest.mark.timeout(180)
 ROOT: Final = Path(__file__).resolve().parents[3]
+PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234"
 Mode: TypeAlias = Literal["late", "unsafe"]
 CONFIG: Final = (
     "model_list: []\n"
@@ -109,7 +110,7 @@ def test_failed_startup_then_retry_releases_database_resources(tmp_path: Path, m
             "DATABASE_URL": database_url,
             "DIRECT_URL": database_url,
             "CONFIG_FILE_PATH": str(config),
-            "LITELLM_MASTER_KEY": "sk-1234" if mode == "unsafe" else f"sk-{uuid.uuid4().hex}",
+            "LITELLM_MASTER_KEY": PUBLICLY_KNOWN_KEY if mode == "unsafe" else f"sk-{uuid.uuid4().hex}",
             "LITELLM_LOCAL_MODEL_COST_MAP": "True",
             "LITELLM_TELEMETRY_ENABLED": "false",
             "PRISMA_HEALTH_WATCHDOG_ENABLED": "true",

tests/unit/proxy/db/test_routing_prisma_wrapper.py

@@ -269,6 +269,40 @@ def test_reader_token_refresh_stops_even_when_the_writer_task_failed() -> None:
     reader.stop_token_refresh_task.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("error", (RuntimeError("refresh failed"), asyncio.CancelledError()))
+async def test_reader_cleanup_finishes_before_writer_stop_failure_propagates(error: BaseException) -> None:
+    from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
+
+    reader_stopping: Final = asyncio.Event()
+    release_reader: Final = asyncio.Event()
+    reader_stopped: Final = asyncio.Event()
+
+    async def stop_reader() -> None:
+        reader_stopping.set()
+        await release_reader.wait()
+        reader_stopped.set()
+
+    writer: Final = MagicMock()
+    reader: Final = MagicMock()
+    writer.stop_token_refresh_task = AsyncMock(side_effect=error)
+    reader.stop_token_refresh_task = AsyncMock(side_effect=stop_reader)
+    routing: Final = RoutingPrismaWrapper(writer=writer, reader=reader)
+    stopping: Final = asyncio.create_task(routing.stop_token_refresh_task())
+    try:
+        await reader_stopping.wait()
+        await asyncio.sleep(0)
+        assert not stopping.done()
+        release_reader.set()
+        with pytest.raises(type(error)):
+            await stopping
+        assert reader_stopped.is_set()
+    finally:
+        release_reader.set()
+        stopping.cancel()
+        await asyncio.gather(stopping, return_exceptions=True)
+
+
 def test_routed_actions_falls_back_to_writer_for_unknown_methods():
     from litellm.proxy.db.routing_prisma_wrapper import _RoutedActions
 

tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py

@@ -15,11 +15,56 @@ from __future__ import annotations
 import asyncio
 import logging
 from typing import Any, Final
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
 from litellm.proxy.utils import PrismaClient
+from litellm.repositories.database import Database
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("engine_dead", (False, True))
+async def test_shutdown_waits_for_hung_request_reconnect_to_time_out(
+    prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch, engine_dead: bool
+) -> None:
+    monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test")
+    reconnecting: Final = asyncio.Event()
+    cancelled: Final = asyncio.Event()
+
+    async def hung_recreate(*args: object, **kwargs: object) -> None:
+        reconnecting.set()
+        try:
+            await asyncio.Event().wait()
+        finally:
+            cancelled.set()
+
+    prisma_client.db.recreate_prisma_client = AsyncMock(side_effect=hung_recreate)
+    prisma_client._engine_confirmed_dead = engine_dead
+    with patch.object(asyncio.get_running_loop(), "time", return_value=1000.0) as clock:
+        reconnect: Final = asyncio.create_task(
+            prisma_client.attempt_db_reconnect(reason="request", force=True, force_recreate=True)
+        )
+        await reconnecting.wait()
+        closing: Final = asyncio.create_task(Database(prisma=prisma_client).aclose())
+        try:
+            await asyncio.sleep(0)
+            assert not closing.done()
+            prisma_client.db.disconnect.assert_not_awaited()
+
+            clock.return_value = 1000.0 + prisma_client._db_watchdog_reconnect_timeout_seconds + 1
+            assert await reconnect is False
+            await closing
+        finally:
+            reconnect.cancel()
+            closing.cancel()
+            await asyncio.gather(reconnect, closing, return_exceptions=True)
+
+    assert cancelled.is_set()
+    assert not prisma_client._db_reconnect_lock.locked()
+    prisma_client.db.disconnect.assert_awaited_once()
+    assert await prisma_client.attempt_db_reconnect(reason="request", force=True) is False
+    prisma_client.db.recreate_prisma_client.assert_awaited_once()
 
 
 @pytest.mark.asyncio