refactor(proxy): own the lifespan database with exit stack callbacks

Drop the hand-rolled lifespan(None | Database) wrapper. The proxy lifespan
already holds an AsyncExitStack, so each owned client registers Database.aclose
on it directly, and tests use contextlib.aclosing.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

litellm/proxy/proxy_server.py

@@ -1429,11 +1429,19 @@ async def _connect_to_count_stored_values(resources: AsyncExitStack) -> Supports
         database_url=str(get_secret("DATABASE_URL")), proxy_logging_obj=proxy_logging_obj
     )
     if prisma_client is None:
-        await resources.enter_async_context(database_runtime.lifespan(database_runtime.Database(prisma=client)))
+        _own_database(resources, client)
     await client.connect()
     return client.writer_db
 
 
+def _own_database(resources: AsyncExitStack, client: PrismaClient | None) -> database_runtime.Database | None:
+    if client is None:
+        return None
+    database: Final = database_runtime.Database(prisma=client)
+    resources.push_async_callback(database.aclose)
+    return database
+
+
 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
@@ -1466,9 +1474,7 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener
     import json
 
     resources.callback(_forget_prisma_client)
-    existing_database: Final = await resources.enter_async_context(
-        database_runtime.lifespan(None if prisma_client is None else database_runtime.Database(prisma=prisma_client))
-    )
+    existing_database: Final = _own_database(resources, prisma_client)
 
     init_verbose_loggers()
 
@@ -1578,9 +1584,7 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener
             )
         )
 
-    database: Final = existing_database or await resources.enter_async_context(
-        database_runtime.lifespan(None if prisma_client is None else database_runtime.Database(prisma=prisma_client))
-    )
+    database: Final = existing_database or _own_database(resources, prisma_client)
 
     await migrate_if_requested(
         environ=os.environ,

litellm/repositories/database.py

@@ -22,15 +22,6 @@ class Database:
             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()
-
-
 @asynccontextmanager
 async def background_task(work: Coroutine[object, object, None]) -> AsyncGenerator[None, None]:
     task: Final = asyncio.create_task(work)

tests/unit/proxy/proxy_server/test_lifecycle.py

@@ -24,6 +24,7 @@ import logging
 import os
 import subprocess
 from collections.abc import Awaitable, Callable
+from contextlib import aclosing
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -118,7 +119,7 @@ async def test_database_lifespan_disconnects_prisma_after_proxy_shutdown(monkeyp
     monkeypatch.setattr(litellm, "cache", None, raising=False)
     monkeypatch.setattr(litellm, "success_callback", [], raising=False)
 
-    async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)):
+    async with aclosing(database_runtime.Database(prisma=fake_prisma)):
         await proxy_shutdown_event()
         fake_prisma.disconnect.assert_not_awaited()
 
@@ -226,7 +227,7 @@ 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)
 
-    async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)):
+    async with aclosing(database_runtime.Database(prisma=fake_prisma)):
         await proxy_shutdown_event()
 
     assert calls == ["flush", "disconnect"]
@@ -268,7 +269,7 @@ 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()))
-    async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)):
+    async with aclosing(database_runtime.Database(prisma=fake_prisma)):
         await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
 
     assert calls == ["flush", "audit", "disconnect"]
@@ -341,7 +342,7 @@ async def test_database_lifespan_propagates_disconnect_error_after_shutdown(monk
     monkeypatch.setattr(litellm, "success_callback", [], raising=False)
 
     with pytest.raises(RuntimeError, match="db gone"):
-        async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)):
+        async with aclosing(database_runtime.Database(prisma=fake_prisma)):
             await proxy_shutdown_event()
 
 
@@ -1017,7 +1018,9 @@ async def test_proxy_shutdown_commits_spend_before_waiting_on_a_slow_logging_tas
 
     task_may_finish: Final = asyncio.Event()
     task_finished: Final = asyncio.Event()
-    flushes: Final[list[tuple[str, bool]]] = []  # mutable-ok: records what each shutdown flush saw, which is the assertion
+    flushes: Final[
+        list[tuple[str, bool]]
+    ] = []  # mutable-ok: records what each shutdown flush saw, which is the assertion
 
     async def slow_logging_task() -> None:
         await task_may_finish.wait()

tests/unit/repositories/test_database.py

@@ -1,76 +1,9 @@
-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.stop_background_tasks = AsyncMock()
-    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()
-    prisma.stop_background_tasks = AsyncMock()
-
-    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()
+from litellm.repositories.database import Database
 
 
 @pytest.mark.asyncio
@@ -80,7 +13,6 @@ async def test_database_disconnects_even_when_stopping_background_tasks_fails()
     prisma.disconnect = AsyncMock()
 
     with pytest.raises(RuntimeError, match="watchdog failed"):
-        async with lifespan(Database(prisma=prisma)):
-            pass
+        await Database(prisma=prisma).aclose()
 
     prisma.disconnect.assert_awaited_once()