refactor(db): keep the database handle feature-agnostic

Database now holds the backend handle and its close operation only, and each feature provider builds its own repository from it, so _db never grows a field per feature. get_database checks the concrete type instead of validating a runtime Protocol on every request, and the banner audit runs as a FastAPI background task so the in-flight drain waits for it before the database closes

The proxy-dependencies skill moves to the repo root .agents/skills and gains bridge rules against passing SQL text or identifiers from Python

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

.agents/skills/proxy-dependencies/SKILL.md

@@ -5,7 +5,7 @@ description: Inject typed repositories and services into Python proxy handlers,
 
 # Proxy dependencies
 
-- Follow the root [database-access skill](../../../../.agents/skills/database-access/SKILL.md) for queries, repository extraction and backend migration
+- Follow the root [database-access skill](../database-access/SKILL.md) for queries, repository extraction and backend migration
 - This skill owns dependency wiring and resource lifetime, not SQL or schema ownership
 
 ## Handler contracts
@@ -19,8 +19,9 @@ description: Inject typed repositories and services into Python proxy handlers,
 
 ## Database ownership
 
-- `litellm._db.Database` exposes typed repositories and an async close operation
-- `from_prisma` adopts an already initialized client without opening a second pool
+- `litellm._db.Database` holds one handle per backend and an async close operation, never a field per feature
+- It wraps the already initialized Prisma client without opening a second pool
+- Feature modules build their repositories from that handle in their own providers, `_db` never imports a repository
 - FastAPI lifespan enters `_db.lifespan` and publishes its handle in request state
 - Drain database consumers before exiting that context, which alone closes the handle
 - Do not keep a parallel Prisma disconnect fallback in proxy shutdown

.agents/skills/proxy-dependencies/references/route-migration.md

@@ -3,7 +3,7 @@
 ## Existing pattern
 
 - `get_database` reads the lifespan-owned handle from request state
-- A feature provider selects the narrow repository, such as `database.user_banners`
+- A feature provider builds its narrow repository from the backend handle, such as `UserBannerRepository(database.prisma)`
 - A handler receives `UserBannerStore | None`, never a Prisma client
 - `UserBannerRepository` keeps the existing Prisma table, serialization and configuration-sync behavior
 - The banner audit provider binds the legacy audit callable at request time, without performing audit I/O
@@ -17,7 +17,8 @@
 - Banner GET without a database returns the default disabled banner
 - Banner PATCH checks the admin role before returning the existing missing-database error
 - Invalid stored banner data falls back to the default, but storage failures still propagate
-- Banner writes finish before the success audit is scheduled
+- Banner writes finish before the success audit is queued as a FastAPI background task
+- Queue post-response work with `BackgroundTasks`, not bare `asyncio.create_task`, so the in-flight drain waits for it before the database closes
 - Do not move background work into an awaited provider merely to remove an import
 
 ## Verification

.agents/skills/proxy-dependencies/references/runtime-ownership.md

@@ -3,8 +3,8 @@
 ## Current foundation
 
 - Prisma startup still owns connection setup, authentication refresh, health checks and migrations
-- Lifespan enters `_db.lifespan` with the initialized client adapted by `from_prisma`
-- The `Database` handle provides typed repositories and `aclose`, without exposing Prisma to handlers
+- Lifespan enters `_db.lifespan` with the initialized client wrapped in `Database`
+- Feature providers build repositories from the `Database` handle, so handlers never see Prisma
 - The existing shutdown sequence drains requests, jobs, logging, spend and audit work before disconnect
 - `_db.lifespan` alone closes the handle, including startup or shutdown failure
 - No parallel disconnect fallback belongs in proxy shutdown
@@ -13,6 +13,9 @@
 ## Adding a backend
 
 - Keep domain operations in repositories and conversion at the native bridge
+- Add the backend as one more handle on `Database`, and let each feature provider pick which repository implementation to build
+- Expose typed operations across the bridge, never a function that accepts SQL text, table names or column names from Python
+- Map dynamic filters and sort keys to closed Rust enums that select static SQL, and pass every value as a bind parameter
 - Define who creates, borrows and closes each pool, including startup-failure cleanup
 - Initialize native resources in the serving process, respecting the existing native runtime's fork restrictions
 - Await native shutdown instead of depending on Python garbage collection

AGENTS.md

@@ -1,4 +1,4 @@
-For proxy dependency wiring and database lifetime changes, follow [proxy-dependencies](litellm-rust/.agents/skills/proxy-dependencies/SKILL.md) alongside [database-access](.agents/skills/database-access/SKILL.md)
+For proxy dependency wiring and database lifetime changes, follow [proxy-dependencies](.agents/skills/proxy-dependencies/SKILL.md) alongside [database-access](.agents/skills/database-access/SKILL.md)
 
 Do not write comments unless they are any of:
 - absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear)

litellm-rust/AGENTS.md

@@ -1,7 +1,5 @@
 # Rust workspace rules
 
-For proxy dependency wiring and database lifetime changes, follow [.agents/skills/proxy-dependencies/SKILL.md](.agents/skills/proxy-dependencies/SKILL.md) alongside the root database-access skill
-
 For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.agents/skills/rust-tracing/SKILL.md)
 
 For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md)

litellm/_db.py

@@ -1,37 +1,18 @@
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Protocol, runtime_checkable
-
-from litellm.repositories.user_banner_repository import UserBannerRepository, UserBannerStore
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from litellm.proxy.utils import PrismaClient
 
 
-@runtime_checkable
-class Database(Protocol):
-    @property
-    def user_banners(self) -> UserBannerStore: ...
-
-    async def aclose(self) -> None: ...
-
-
 @dataclass(frozen=True, slots=True)
-class _PrismaDatabase:
-    client: "PrismaClient"
-
-    @property
-    def user_banners(self) -> UserBannerStore:
-        return UserBannerRepository(self.client)
+class Database:
+    prisma: "PrismaClient"
 
     async def aclose(self) -> None:
-        await self.client.disconnect()
-
-
-def from_prisma(client: "PrismaClient") -> Database:
-    """Adapt an initialized client without opening another pool."""
-    return _PrismaDatabase(client)
+        await self.prisma.disconnect()
 
 
 @asynccontextmanager

litellm/proxy/deps.py

@@ -1,18 +1,22 @@
 from typing import Annotated, Final, TypeAlias
 
 from fastapi import Depends, Request
-from pydantic import ConfigDict, TypeAdapter
 
 from litellm._db import Database
 
-_DATABASE_ADAPTER: Final[TypeAdapter[Database | None]] = TypeAdapter(
-    Database | None, config=ConfigDict(arbitrary_types_allowed=True)
-)
 _MISSING_DATABASE: Final = object()
 
 
+class DatabaseStateMissingError(RuntimeError):
+    def __init__(self) -> None:
+        super().__init__("The proxy lifespan did not publish a Database in request state")
+
+
 async def get_database(request: Request) -> Database | None:
-    return _DATABASE_ADAPTER.validate_python(getattr(request.state, "database", _MISSING_DATABASE))
+    database: Final[object] = getattr(request.state, "database", _MISSING_DATABASE)
+    if database is None or isinstance(database, Database):
+        return database
+    raise DatabaseStateMissingError()
 
 
 DatabaseDependency: TypeAlias = Annotated[Database | None, Depends(get_database)]

litellm/proxy/proxy_server.py

@@ -1552,7 +1552,7 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener
         )
 
     database: Final = await resources.enter_async_context(
-        _db.lifespan(None if prisma_client is None else _db.from_prisma(prisma_client))
+        _db.lifespan(None if prisma_client is None else _db.Database(prisma=prisma_client))
     )
 
     await migrate_if_requested(

litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py

@@ -6,7 +6,7 @@ from fastapi import Depends
 from litellm.models.user_banner import UserBanner
 from litellm.proxy._types import LitellmTableNames, UserAPIKeyAuth
 from litellm.proxy.deps import DatabaseDependency
-from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerStore
+from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerRepository, UserBannerStore
 
 
 class UserBannerAudit(Protocol):
@@ -14,7 +14,7 @@ class UserBannerAudit(Protocol):
 
 
 async def get_user_banner_repository(database: DatabaseDependency) -> UserBannerStore | None:
-    return None if database is None else database.user_banners
+    return None if database is None else UserBannerRepository(database.prisma)
 
 
 async def get_user_banner_audit() -> UserBannerAudit:

litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py

@@ -1,7 +1,6 @@
-import asyncio
 from typing import Annotated, Final
 
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
 from pydantic import ConfigDict
 
 from litellm._uuid import uuid4
@@ -51,6 +50,7 @@ async def update_user_banner(
     user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
     repository: UserBannerRepositoryDependency,
     audit: UserBannerAuditDependency,
+    background_tasks: BackgroundTasks,
 ) -> UpdateUserBannerResponse:
     """
     Publish, edit, or unpublish the dashboard banner.
@@ -72,6 +72,6 @@ async def update_user_banner(
 
     await repository.save_banner(banner)
 
-    asyncio.create_task(audit(before, banner, user_api_key_dict))
+    background_tasks.add_task(audit, before, banner, user_api_key_dict)
 
     return UpdateUserBannerResponse(message="User banner updated successfully", banner=banner)

tests/unit/proxy/proxy_server/test_lifecycle.py

@@ -118,7 +118,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 _db.lifespan(_db.from_prisma(fake_prisma)):
+    async with _db.lifespan(_db.Database(prisma=fake_prisma)):
         await proxy_shutdown_event()
         fake_prisma.disconnect.assert_not_awaited()
 
@@ -225,7 +225,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 _db.lifespan(_db.from_prisma(fake_prisma)):
+    async with _db.lifespan(_db.Database(prisma=fake_prisma)):
         await proxy_shutdown_event()
 
     assert calls == ["flush", "disconnect"]
@@ -266,7 +266,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 _db.lifespan(_db.from_prisma(fake_prisma)):
+    async with _db.lifespan(_db.Database(prisma=fake_prisma)):
         await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
 
     assert calls == ["flush", "audit", "disconnect"]
@@ -338,7 +338,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 _db.lifespan(_db.from_prisma(fake_prisma)):
+        async with _db.lifespan(_db.Database(prisma=fake_prisma)):
             await proxy_shutdown_event()
 
 

tests/unit/proxy/test_deps.py

@@ -6,15 +6,12 @@ from unittest.mock import AsyncMock, MagicMock
 import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
-from pydantic import ValidationError
+from starlette.requests import Request
 
 from litellm import _db
 from litellm.models.user_banner import UserBanner
-from litellm.proxy.deps import get_database
-from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import (
-    UserBannerRepositoryDependency,
-    get_user_banner_repository,
-)
+from litellm.proxy.deps import DatabaseStateMissingError, get_database
+from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import UserBannerRepositoryDependency
 
 
 def _app(database: _db.Database | None) -> FastAPI:
@@ -36,7 +33,7 @@ def test_lifespan_database_is_shared_until_shutdown():
     prisma: Final = MagicMock()
     prisma.disconnect = AsyncMock()
     prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None)
-    app: Final = _app(_db.from_prisma(prisma))
+    app: Final = _app(_db.Database(prisma=prisma))
     with TestClient(app) as client:
         assert client.get("/banner").json() == UserBanner().model_dump()
         prisma.db.litellm_uisettings.find_unique.return_value = SimpleNamespace(
@@ -63,7 +60,7 @@ def test_dependency_override_changes_next_request_without_changing_lifespan_owne
     prisma.db.litellm_uisettings.find_unique = AsyncMock(
         return_value=SimpleNamespace(ui_settings={"message": "stored"})
     )
-    app: Final = _app(_db.from_prisma(prisma))
+    app: Final = _app(_db.Database(prisma=prisma))
     with TestClient(app) as client:
         assert client.get("/banner").json() == UserBanner(message="stored").model_dump()
         monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None)
@@ -74,22 +71,6 @@ def test_dependency_override_changes_next_request_without_changing_lifespan_owne
 
 @pytest.mark.asyncio
 async def test_missing_lifespan_is_not_treated_as_a_disabled_database():
-    from starlette.requests import Request
-
     request: Final = Request({"type": "http", "state": {}})
-    with pytest.raises(ValidationError, match="Database"):
+    with pytest.raises(DatabaseStateMissingError):
         await get_database(request)
-
-
-@pytest.mark.asyncio
-async def test_constructing_repository_does_not_touch_database():
-    prisma: Final = MagicMock()
-    repository: Final = await get_user_banner_repository(_db.from_prisma(prisma))
-    assert repository is not None
-    assert prisma.mock_calls == []
-
-
-def test_dependencies_do_not_become_openapi_parameters():
-    operation: Final = _app(None).openapi()["paths"]["/banner"]["get"]
-    assert "parameters" not in operation
-    assert "requestBody" not in operation

tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py

@@ -58,7 +58,7 @@ def _mock_prisma(monkeypatch, record=None):
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=record)
     mock_prisma.db.litellm_uisettings.upsert = AsyncMock()
-    monkeypatch.setitem(app.dependency_overrides, get_database, lambda: _db.from_prisma(mock_prisma))
+    monkeypatch.setitem(app.dependency_overrides, get_database, lambda: _db.Database(prisma=mock_prisma))
     return mock_prisma
 
 

tests/unit/test__db.py

@@ -12,7 +12,7 @@ from litellm import _db
 async def test_lifespan_closes_database_when_owner_fails(failure):
     prisma: Final = MagicMock()
     prisma.disconnect = AsyncMock()
-    database: Final = _db.from_prisma(prisma)
+    database: Final = _db.Database(prisma=prisma)
 
     async def fail():
         async with _db.lifespan(database) as current: