Merge 55a0a25080 into 0a1dbc0fb4

litellm/proxy/_types.py

@@ -124,6 +124,7 @@ from .types_utils.utils import get_instance_fn, validate_custom_validate_return_
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
 
+    from litellm.repositories.database import Database
     from litellm.tracing import TraceReceiver
 
     Span = _Span | Any
@@ -133,6 +134,7 @@ else:
 
 class ProxyLifespanState(TypedDict):
     tracing_receiver: ReadOnly["TraceReceiver | None"]
+    database: ReadOnly["Database | None"]
 
 
 class ReconcileOutcome(NamedTuple):

litellm/proxy/proxy_server.py

@@ -168,6 +168,7 @@ from litellm.proxy.common_utils.token_count_settings import (
 )
 from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog
 from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields
+from litellm.repositories import database as database_runtime
 from litellm.router_utils.access_windows import access_windows_config_error
 from litellm.router_utils.add_retry_fallback_headers import (
     get_fallback_errors_from_headers,
@@ -1288,13 +1289,7 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N
     if worker_heartbeat is not None and prisma_client:
         await worker_heartbeat.deregister()
     if prisma_client:
-        # Drain the SGR fold first: it lives in memory, so an un-drained interval
-        # is lost, and a write attempted after disconnect raises
-        # ClientNotConnectedError rather than persisting anything. Ordering this
-        # inside the same guard is what keeps the two from drifting apart.
         await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client))
-        verbose_proxy_logger.debug("Disconnecting from Prisma")
-        await prisma_client.disconnect()
     else:
         await _drain_audit_tasks_on_shutdown()
 
@@ -1425,8 +1420,19 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries:
     return client.writer_db
 
 
+def _forget_prisma_client() -> None:
+    global prisma_client  # noqa: PLW0603  # routes still read this global until they move to DatabaseDependency
+    prisma_client = None  # rebind-ok: a closed client must not survive into the next lifespan entry
+
+
 @asynccontextmanager
 async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
+    async with AsyncExitStack() as resources, _proxy_lifespan(app, resources) as state:
+        yield state
+
+
+@asynccontextmanager
+async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGenerator[ProxyLifespanState, None]:
     global \
         prisma_client, \
         master_key, \
@@ -1550,6 +1556,11 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState
             )
         )
 
+    resources.callback(_forget_prisma_client)
+    database: Final = await resources.enter_async_context(
+        database_runtime.lifespan(None if prisma_client is None else database_runtime.Database(prisma=prisma_client))
+    )
+
     await migrate_if_requested(
         environ=os.environ,
         master_key=master_key,
@@ -1784,7 +1795,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState
         enabled=tracing_enabled,
         settings=tracing_settings,
     ) as receiver:
-        state: Final[ProxyLifespanState] = {"tracing_receiver": receiver}
+        state: Final[ProxyLifespanState] = {"tracing_receiver": receiver, "database": database}
         from litellm.proxy.admin_mcp import admin_mcp_lifespan
 
         telemetry_settings: Final = load_telemetry_settings()

litellm/repositories/database.py

@@ -0,0 +1,27 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from litellm._logging import verbose_proxy_logger
+
+if TYPE_CHECKING:
+    from litellm.proxy.utils import PrismaClient
+
+
+@dataclass(frozen=True, slots=True)
+class Database:
+    prisma: "PrismaClient"
+
+    async def aclose(self) -> None:
+        verbose_proxy_logger.debug("Disconnecting from Prisma")
+        await self.prisma.disconnect()
+
+
+@asynccontextmanager
+async def lifespan(database: Database | None) -> AsyncGenerator[Database | None, None]:
+    try:
+        yield database
+    finally:
+        if database is not None:
+            await database.aclose()

tests/unit/proxy/proxy_server/test_lifecycle.py

@@ -33,9 +33,9 @@ from pydantic import BaseModel
 from typing_extensions import TypedDict
 
 import litellm.proxy.proxy_server as ps
+from litellm.proxy.management_helpers.audit_logs import track_audit_task
 from litellm.proxy.proxy_server import (
     ProxyStartupEvent,
-    initialize_shared_aiohttp_session,
     _resolve_pydantic_type,
     _resolve_typed_dict_type,
     cleanup_router_config_variables,
@@ -43,14 +43,14 @@ from litellm.proxy.proxy_server import (
     get_litellm_model_info,
     initialize,
     initialize_from_worker_config,
+    initialize_shared_aiohttp_session,
     load_from_azure_key_vault,
     proxy_shutdown_event,
     proxy_startup_event,
     run_ollama_serve,
     save_worker_config,
 )
-
-from litellm.proxy.management_helpers.audit_logs import track_audit_task
+from litellm.repositories import database as database_runtime
 
 from .conftest import normalize
 
@@ -101,7 +101,7 @@ def test_cleanup_router_config_variables_fails_on_unknown_attr_raises():
 
 
 @pytest.mark.asyncio
-async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch):
+async def test_database_lifespan_disconnects_prisma_after_proxy_shutdown(monkeypatch):
     fake_prisma = MagicMock()
     fake_prisma.disconnect = AsyncMock()
     monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
@@ -117,7 +117,9 @@ async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch):
     monkeypatch.setattr(litellm, "cache", None, raising=False)
     monkeypatch.setattr(litellm, "success_callback", [], raising=False)
 
-    await proxy_shutdown_event()
+    async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)):
+        await proxy_shutdown_event()
+        fake_prisma.disconnect.assert_not_awaited()
 
     observed = {
         "disconnect_called": fake_prisma.disconnect.await_count == 1,
@@ -222,7 +224,8 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke
     monkeypatch.setattr(litellm, "cache", None, raising=False)
     monkeypatch.setattr(litellm, "success_callback", [], raising=False)
 
-    await proxy_shutdown_event()
+    async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)):
+        await proxy_shutdown_event()
 
     assert calls == ["flush", "disconnect"]
 
@@ -262,7 +265,8 @@ async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_befor
     monkeypatch.setattr(litellm, "success_callback", [], raising=False)
 
     track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts()))
-    await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
+    async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)):
+        await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
 
     assert calls == ["flush", "audit", "disconnect"]
 
@@ -318,7 +322,7 @@ async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_t
 
 
 @pytest.mark.asyncio
-async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
+async def test_database_lifespan_propagates_disconnect_error_after_shutdown(monkeypatch):
     fake_prisma = MagicMock()
     fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone"))
     monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
@@ -333,7 +337,8 @@ async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
     monkeypatch.setattr(litellm, "success_callback", [], raising=False)
 
     with pytest.raises(RuntimeError, match="db gone"):
-        await proxy_shutdown_event()
+        async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)):
+            await proxy_shutdown_event()
 
 
 # ---------------------------------------------------------------------------
@@ -1060,7 +1065,7 @@ def test_otel_global_provider_published_after_callback_init():
     span was orphaned. The publish (``publish_global_otel_v2_provider``) must
     therefore appear after ``_initialize_startup_logging`` in the lifespan source.
     """
-    wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
+    wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan)
     source = inspect.getsource(wrapped)
     init_pos = source.find("ProxyStartupEvent.initialize_startup_logging(")
     publish_pos = source.find("publish_global_otel_v2_provider(")
@@ -1146,7 +1151,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database():
     The call must follow Prisma setup so DB-backed deployments do not false-positive.
     Direct ``_warn_budget_without_db`` tests cover the warning behavior itself.
     """
-    wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
+    wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan)
     source = inspect.getsource(wrapped)
     budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:")
     warn_pos = source.find("warn_budget_without_db(")

tests/unit/proxy/test_proxy_server.py

@@ -113,7 +113,7 @@ async def test_admin_mcp_failure_still_closes_proxy_resources(
     message: Final = "LITELLM_LICENSE" if failure_phase == "license" else f"{failure_phase} failed"
     async def run_lifespan() -> None:
         async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state:
-            assert state == {"tracing_receiver": None}
+            assert state == {"tracing_receiver": None, "database": None}
             if failure_phase == "serving":
                 raise RuntimeError("serving failed")
             if failure_phase == "cancelled":
@@ -162,7 +162,7 @@ async def test_proxy_shutdown_drains_active_admin_tool_before_closing_connector(
         transport=httpx2.ASGITransport(app=admin_mcp_proxy), base_url="http://localhost:4000"
     ) as client:
         async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state:
-            assert state == {"tracing_receiver": None}
+            assert state == {"tracing_receiver": None, "database": None}
             request: Final = asyncio.create_task(client.post(
                 "/admin/mcp",
                 headers={"Authorization": "Bearer admin", "Accept": "application/json, text/event-stream"},

tests/unit/proxy/test_proxy_server_endpoints_and_startup.py

@@ -1838,6 +1838,12 @@ async def test_proxy_startup_says_a_lingering_migrate_from_variable_can_be_delet
 
 
 class _PrismaClientWhoseDatabaseRejectsQueries:
+    def __init__(self) -> None:
+        self.disconnect_calls = 0
+
+    async def disconnect(self) -> None:
+        self.disconnect_calls += 1
+
     class _Database:
         async def query_raw(self, query, *args):
             raise RuntimeError("permission denied for table LiteLLM_CredentialsTable")
@@ -1852,7 +1858,8 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc
     from litellm.proxy.proxy_server import proxy_startup_event
 
     _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"})
-    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries())
+    prisma: Final = _PrismaClientWhoseDatabaseRejectsQueries()
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
     monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY)
 
     with (
@@ -1865,6 +1872,25 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc
     notices = [record.getMessage() for record in caplog.records if "LITELLM_MIGRATE_FROM_MASTER_KEY" in record.message]
     assert len(notices) == 1
     assert "Could not migrate stored values" in notices[0]
+    assert prisma.disconnect_calls == 1
+
+
+@pytest.mark.asyncio
+async def test_proxy_startup_after_a_failed_migration_does_not_reuse_the_closed_client(monkeypatch, tmp_path):
+    from fastapi import FastAPI
+
+    from litellm.proxy.proxy_server import proxy_startup_event
+
+    _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"})
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries())
+    monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY)
+    with pytest.raises(RuntimeError, match="permission denied"):
+        async with proxy_startup_event(FastAPI()):
+            pass
+
+    monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", "")
+    async with proxy_startup_event(FastAPI()) as state:
+        assert state["database"] is None
 
 
 @pytest.mark.asyncio

tests/unit/repositories/test_database.py

@@ -0,0 +1,71 @@
+import asyncio
+from typing import Final
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm.repositories.database import Database, lifespan
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failure", [RuntimeError, asyncio.CancelledError])
+async def test_lifespan_closes_database_when_owner_fails(failure):
+    prisma: Final = MagicMock()
+    prisma.disconnect = AsyncMock()
+    database: Final = Database(prisma=prisma)
+
+    async def fail():
+        async with lifespan(database) as current:
+            assert current is database
+            prisma.disconnect.assert_not_awaited()
+            raise failure("owner failed")
+
+    with pytest.raises(failure):
+        await fail()
+    prisma.disconnect.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_lifespan_without_database_preserves_owner_error():
+    async def fail():
+        async with lifespan(None) as database:
+            assert database is None
+            raise RuntimeError("owner failed")
+
+    with pytest.raises(RuntimeError, match="owner failed"):
+        await fail()
+
+
+@pytest.mark.asyncio
+async def test_cancelling_owner_waits_for_database_close() -> None:
+    serving: Final = asyncio.Event()
+    closing: Final = asyncio.Event()
+    release_close: Final = asyncio.Event()
+    prisma: Final = MagicMock()
+
+    async def disconnect() -> None:
+        closing.set()
+        await release_close.wait()
+
+    prisma.disconnect = AsyncMock(side_effect=disconnect)
+
+    async def serve() -> None:
+        async with lifespan(Database(prisma=prisma)):
+            serving.set()
+            await asyncio.Event().wait()
+
+    owner: Final = asyncio.create_task(serve())
+    try:
+        await asyncio.wait_for(serving.wait(), timeout=2)
+        owner.cancel()
+        await asyncio.wait_for(closing.wait(), timeout=2)
+        assert not owner.done()
+        release_close.set()
+        with pytest.raises(asyncio.CancelledError):
+            await owner
+    finally:
+        release_close.set()
+        if not owner.done():
+            owner.cancel()
+        await asyncio.gather(owner, return_exceptions=True)
+    prisma.disconnect.assert_awaited_once()