refactor(proxy): inject database repositories through per-repository dependencies (#48507)

* refactor(proxy): inject database repositories through per-repository dependencies

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

* test(proxy): model Prisma client cleanup

* test(rust): include NativeDatabase in bridge API surface

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

.agents/skills/database-access/SKILL.md

@@ -1,54 +1,73 @@
 ---
 name: database-access
-description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope
+description: Organize database repositories and proxy dependency wiring, manage connection ownership, or port Prisma operations to Rust and SQLx without expanding migration scope
 metadata:
   internal: true
 ---
 
 # Database access
 
+## Situations
+
+- Adding a query or an operation to a repository, read [adding a query](references/python-repository.md#adding-a-query)
+- Moving existing Prisma calls behind a repository, read [moving existing access](references/python-repository.md#moving-existing-access-into-a-repository)
+- Porting a repository from Prisma to Rust and SQLx, read [porting to SQLx](references/python-repository.md#porting-a-repository-to-sqlx)
+- Injecting a repository into a route, read [Python and FastAPI](references/python-fastapi.md) or [Rust and Axum](references/rust-axum.md)
+
 ## Adding new database access
 
 - Put Python database access behind an operation in `litellm/repositories`
-- Give callers typed inputs and outputs, not Prisma clients, models, or query dictionaries
-- Keep query execution and backend selection inside the repository
+- Give each repository its own FastAPI provider, so a route builds only the repositories it uses
+- A provider borrows existing handles and performs no database I/O during construction
+- Define the operations a route calls as a Protocol in `litellm/models/<feature>.py`, next to its domain types
+  - The Protocol takes and returns domain types, never Prisma models, raw rows or serialized payloads
+  - Parsing and serialization of stored values live in the repository
+  - Add a Protocol when a repository gets a provider, not up front for every repository
+- Keep query execution inside the repository, with one implementation per repository and no backend switch or fallback
+- A repository is either all Prisma or all native, never a mix of both
 - Prisma is fine for the implementation, adding a query does not require a Rust port
 - Native Rust callers use `litellm-rust/crates/db` directly
 
-## Asked to move existing access behind a repository
+## Framework wiring
 
-- Identify the calls being moved and their transaction boundary
-- Move the existing Prisma implementation into the relevant repository
-- Replace those direct database calls with repository operations
-- Preserve return values, errors, and transaction behavior, with focused regression tests
-- Keep this an extraction, not a backend migration or unrelated cleanup
+- Python routes receive a per-repository dependency such as `UserBannerRepositoryDep` and call its operations, such as `banners.get_banner()`
+- `DatabaseDep` lives in `litellm/proxy/deps/database.py`, repository providers in `litellm/proxy/deps/repositories.py`
+  - A provider returns the feature's Protocol, so routes never see which backend implements it
+  - Dependencies used by one feature only, like the banner audit binding, stay next to that feature's routes
+- Repositories never import FastAPI, so jobs and helpers can use them without a request
+- Use named FastAPI providers with `Annotated[..., Depends(...)]`, ordinary helpers and jobs receive explicit arguments
+- Preserve each route's no-database behavior, authorization precedence and background-task scheduling
+- Read [Python and FastAPI](references/python-fastapi.md) for provider examples, lifespan state and route verification
+- Read [Rust and Axum](references/rust-axum.md) when adding database access to Rust handlers
 
-## An equivalent Rust or SQLx query already exists
+## Connection ownership
 
-- Leave it and the current backend selection unchanged unless integration is explicitly requested
-- Its existence alone does not justify switching a Prisma caller to Rust
-- Repository extraction can keep using Prisma even when the Rust query is available
+- Keep connection setup and shutdown separate from repositories, which own no pools
+- Reuse existing handles, never connect per request or disconnect from a repository operation
+- Give each pool one owner for setup and shutdown, borrowed handles never close it
+- Finish requests, jobs and queued database work before closing connections
+- Clean up resources opened before a startup failure, without a second disconnect fallback
+- Add a backend to the existing lifecycle owner and pass its handle only to repositories that need it
+- Preserve cancellation, timeout, error and transaction semantics across the native bridge
 
-## Asked to port a Prisma query to Rust or SQLx
+## Changing ownership or reloading resources
 
-- Identify the existing inputs, outputs, errors, and transaction boundary
-- Check `litellm-rust/crates/db` for an equivalent query and reuse it if its behavior matches
-- Otherwise follow `litellm-rust/crates/db/AGENTS.md`
-  - Put static SQL in `queries/` and compile-check it with SQLx
-  - Keep typed inputs and reusable row projections in `src/`
-  - Refresh SQLx query metadata after SQL changes
-- Keep Python conversion in the Python bridge, not in `litellm-rust/crates/db`
-- When caller integration is requested, expose the Rust operation through the bridge and call it from the repository
-  - A query-only port can land without switching callers
-- Verify behavior against real PostgreSQL before switching callers
-  - Check relevant null handling, conflicts, ordering, errors, and transaction semantics
-- Remove the old implementation only when no callers need it
+- Inventory initialization, reload, tests, background jobs and shutdown together
+- Move every writer when changing authoritative storage, copying a replaceable global into state is insufficient
+- A frozen repository does not make its clients immutable
+- Decide whether each request captures one resource generation or reads it again during processing
+- Build and validate replacements before publishing them, preserving the old generation on failure
+- Serialize reloads and coordinate publication with shutdown
+- Drain streams and background users before retiring old resources, retaining every generation grows resource use
+- Define multi-worker coordination separately from an in-process swap
+- Choose mechanisms such as `ArcSwap` only after defining publication, capture and retirement
 
 ## Migration context and boundaries
 
-- Long-term direction is to centralize Python database access in repositories, then replace their Prisma implementations incrementally with Rust APIs
-- Mixed Prisma and Rust implementations are expected during the transition, unrelated existing calls can stay as they are
-- Both backends may have separate pools through PgBouncer, size them together and verify pooling-mode compatibility
+- Long-term direction is to centralize Python database access in repositories, then replace their Prisma implementations with Rust APIs one repository at a time
+- Prisma and native repositories coexist during the transition, unrelated existing calls can stay as they are
+- Prisma and SQLx keep separate client pools to PgBouncer, which pools connections to PostgreSQL
+  - Size both layers together and preserve reader/writer routing and pooling-mode compatibility
 - PgBouncer does not share transactions between Prisma and SQLx, keep each atomic unit of work on one backend and transaction
 - Never retry a failed Rust write through Prisma when it may already have committed
 - Prisma owns schema migrations today, eventual SQLx migration ownership and Prisma removal are separate planned steps, not part of a query port

.agents/skills/database-access/references/python-fastapi.md

@@ -0,0 +1,101 @@
+# Python and FastAPI
+
+## Existing pattern
+
+- `UserBannerStore` in `litellm/models/user_banner.py` is the Protocol the banner routes call
+- `get_database` in `litellm/proxy/deps/database.py` reads the lifespan-owned handle from request state
+- `get_user_banner_repository` in `litellm/proxy/deps/repositories.py` builds the banner repository from that handle
+- Choose each repository's replica routing once in its provider, never at call sites
+- A handler receives `UserBannerStore | None`, never a Prisma client or a concrete repository
+- `UserBannerRepository` implements `UserBannerStore` structurally, keeping 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
+- A native repository can later implement the same Protocol without changing handler signatures, see [porting a repository to SQLx](python-repository.md#porting-a-repository-to-sqlx)
+
+## Dependency injection
+
+- Add one provider and one `Annotated` alias per repository in `litellm/proxy/deps/repositories.py`
+- Annotate the provider and alias with the feature's Protocol, not the concrete repository
+- A route depends only on the repositories it calls, so a request never builds repositories it does not use
+- Do not build repositories once at lifespan, a reconnect can replace the Prisma client they wrap
+
+```python
+class UserBannerStore(Protocol):
+    async def get_banner(self) -> UserBanner: ...
+
+    async def save_banner(self, banner: UserBanner) -> None: ...
+```
+
+```python
+async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None:
+    return None if database is None else UserBannerRepository(database.prisma)
+
+
+UserBannerRepositoryDep: TypeAlias = Annotated[UserBannerStore | None, Depends(get_user_banner_repository)]
+```
+
+- FastAPI caches each provider within a request, so routes and sub-dependencies share one instance
+- Ordinary helpers and jobs receive explicit arguments, never unresolved `Depends` defaults
+- A route still reading the `prisma_client` global or building a repository inline switches to the repository's `...RepositoryDep`, keeping its no-database behavior and error precedence
+
+```python
+async def get_user_banner(banners: UserBannerRepositoryDep) -> UserBanner:
+    if banners is None:
+        return UserBanner()
+    return await banners.get_banner()
+```
+
+## Lifespan state
+
+- Prisma startup owns connection setup, authentication refresh, health checks and migrations
+- `Database` holds the initialized Prisma client and the `NativeDatabase` handle, without holding fields per feature
+  - The native handle opens no pool yet, its `connect` and `close` are noops
+- The proxy lifespan builds the handle with `Database.open`, registers `Database.aclose` on its exit stack and publishes it in request state
+  - Callers only open and close `Database`, which closes Prisma and the native handle together, and `open` closes Prisma itself when the native connect fails
+- Starlette shallow-copies lifespan state into request state, which is distinct from `app.state`
+- The database module never imports a repository, legacy globals borrow its handle without owning shutdown
+- Drain requests, jobs, logging, spend and audit work before database teardown
+- The exit stack closes the handle on exit, including startup failures after the client is created
+- When adding native resources, initialize them in the serving process, respecting runtime fork restrictions
+- Await native shutdown instead of depending on Python garbage collection
+- Apply the [connection ownership rules](../SKILL.md#connection-ownership) when changing the owner
+
+## Timing and errors
+
+- FastAPI resolves dependencies before entering the handler
+- Handler exception handling cannot catch failures raised during dependency resolution
+- A universal required-database dependency can override authorization and validation behavior
+- 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 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
+
+- Exercise real FastAPI handlers using `dependency_overrides` for external boundaries
+- Keep production parsing and persistence code real
+- Restore overrides after each test
+- Verify complete response bodies, status codes and OpenAPI against the base
+- Include combined failure conditions to pin error precedence
+- Check saved values, fresh revisions, audit contents and failed-write behavior
+- Check lifespan state with a running lifespan, not only a TestClient that skips startup
+- Verify providers do not connect, query, disconnect or audit eagerly
+- Use real PostgreSQL and live proxy requests when changing persistence behavior
+
+## Scope
+
+- Preserve direct Python call sites or migrate them explicitly
+- Inventory imports, state writers, side effects and error ordering before migrating a feature
+- Preserve configuration sync and serialization, and state which legacy reads remain
+- Do not migrate every route in a large module just to establish the pattern
+- Do not invent providers for resources without consumers
+
+## Framework references
+
+- [Lifespan events](https://fastapi.tiangolo.com/advanced/events/) and [Starlette lifespan state](https://starlette.dev/lifespan/)
+- [`SessionDep`-style `Annotated` aliases](https://fastapi.tiangolo.com/tutorial/sql-databases/#create-a-session-dependency), which `DatabaseDep` and `UserBannerRepositoryDep` follow
+- [Dependency composition and caching](https://fastapi.tiangolo.com/tutorial/dependencies/sub-dependencies/)
+- [Background tasks and dependencies with `yield`](https://fastapi.tiangolo.com/advanced/advanced-dependencies/#background-tasks-and-dependencies-with-yield-technical-details)
+- [Dependency overrides](https://fastapi.tiangolo.com/advanced/testing-dependencies/)
+- [Testing lifespan](https://fastapi.tiangolo.com/advanced/testing-events/)

.agents/skills/database-access/references/python-repository.md

@@ -0,0 +1,119 @@
+# Python repositories
+
+## Adding a query
+
+Use this when a feature needs a read or write that no repository offers yet, including one more operation on an existing repository
+
+### Before starting
+
+- Search `litellm/repositories` for a repository that already owns the table, and extend it instead of adding a second one
+- Check `litellm/models/<feature>.py` for an existing Protocol and domain types
+- Decide the route's behavior when no database is configured, and where that check sits relative to authorization and validation
+
+### Steps
+
+1. Domain types and Protocol in `litellm/models/<feature>.py`
+   - The Protocol takes and returns domain types, never Prisma models, raw rows or serialized payloads
+   - Add the Protocol together with the provider, not for repositories nothing injects
+2. Repository in `litellm/repositories/<feature>_repository.py`
+   - Implement the operation with Prisma, a Rust port is not required
+   - Compose the Prisma client and type the tables it touches with `TableActions` from `prisma_protocols.py`, never inherit `PrismaTableRepository`
+   - Parse and serialize stored values here, so callers only see domain types
+   - Never import FastAPI, connect, disconnect or choose replicas per call
+3. Provider and `Annotated` alias in `litellm/proxy/deps/repositories.py`, see [dependency injection](python-fastapi.md#dependency-injection)
+4. The route takes `<Feature>RepositoryDep`, handles `None` with the behavior decided above, and calls the operation
+5. Jobs and helpers that need the operation receive the repository as an explicit argument, built by their caller from the handle it already holds
+
+### Done when
+
+- Handler tests drive the real route with `dependency_overrides` and the real repository code, see [verification](python-fastapi.md#verification)
+- Tests cover the no-database path, error precedence, and what was actually stored
+- Persistence changes are checked against real PostgreSQL through a live proxy
+
+### Out of scope
+
+- Porting the query to Rust, see [porting a repository to SQLx](#porting-a-repository-to-sqlx)
+- Moving unrelated direct Prisma calls in the same module
+
+## Moving existing access into a repository
+
+Use this when code calls Prisma directly, through `prisma_client.db.<table>` or a helper, and the calls should live behind a repository. The backend stays Prisma
+
+### Before starting
+
+- List every call being moved, its callers, and the transaction or batch it runs in
+- List the module's imports, global state writers, side effects such as audit and config sync, and error ordering
+- Record current return values, raised errors and response bodies, these are what the change must preserve
+
+### Steps
+
+1. Move the existing Prisma code into the relevant repository in `litellm/repositories` without rewriting the queries
+2. Add domain types and a Protocol in `litellm/models/<feature>.py` when a route will inject the repository
+3. Replace each direct call with the repository operation, through a provider for routes and explicit arguments for jobs and helpers, see [Python and FastAPI](python-fastapi.md)
+4. Calls sharing a transaction move into the same repository or unit of work in one change
+5. State which legacy reads remain and why, when the change does not move all of them
+
+### Done when
+
+- Focused regression tests pin the preserved return values, errors and transaction behavior
+- Response bodies, status codes and OpenAPI match the base
+- Configuration sync and serialization behave the same, including invalid stored data
+
+### Out of scope
+
+- Switching the backend, even when an equivalent Rust query exists
+- Converting other routes in the module just to establish the pattern
+- Cleanups unrelated to the moved calls
+
+## Porting a repository to SQLx
+
+Status: TBD. No repository has been ported yet. `database.native` is a `NativeDatabase` from `crates/python-bridge/src/db.rs` whose `connect` and `close` are noops that open no pool, and the PyO3 repository surface does not exist. A real pool takes typed writer/reader settings and a token provider, never the Prisma `DATABASE_URL`, which carries Prisma-only params and rotating IAM tokens. The constraints below are settled, the steps are a sketch to revise after the first port lands
+
+Use this only when a port is explicitly requested. Repository extraction does not require one
+
+### When not to port
+
+- An equivalent Rust or SQLx query existing is not a reason to switch a Prisma caller
+- Leave existing backend selection unchanged unless integration is explicitly requested
+- Never port part of a repository, it is either all Prisma or all native
+
+### Before starting
+
+- Identify the repository's inputs, outputs, errors and transaction boundary
+- List repositories that share a transaction with it, they move together
+- Check `litellm-rust/crates/db` for an equivalent query and reuse it when its behavior matches
+
+### Steps
+
+1. Rust queries, following `litellm-rust/crates/db/AGENTS.md`
+   - Static SQL in `queries/`, compile-checked with SQLx, typed inputs and row projections in `src/`
+   - Dynamic filters and sort keys become closed Rust enums that select static SQL, every value is bound
+   - Refresh SQLx query metadata after SQL changes
+2. PyO3 surface exposing typed operations, never SQL text, table names or column names from Python. TBD: where it lives and how the handle reaches it
+3. Native repository implementing the same Protocol in `litellm/repositories`
+   - Converts typed PyO3 results to domain types, preserving defaults, invalid-data handling and storage errors
+   - Python conversion stays in the bridge, not in `crates/db`
+4. Switch the provider in `litellm/proxy/deps/repositories.py`, the Protocol and routes stay unchanged
+
+```diff
+ async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None:
+-    return None if database is None else UserBannerRepository(database.prisma)
++    return None if database is None else NativeUserBannerRepository(database.native)
+```
+
+5. Delete the Prisma repository in the same change and switch every caller, including jobs and helpers
+
+### Done when
+
+- Behavior is verified against real PostgreSQL, covering null handling, conflicts, ordering, errors and transaction semantics
+- `DATABASE_URL=... cargo test -p litellm-db --locked -- --include-ignored` passes against a Prisma-migrated database
+- Live proxy requests show unchanged responses. TBD: a Python-side real-PostgreSQL test harness
+
+### Rollback
+
+- Roll back the release, never retry an uncertain native write through Prisma
+
+### Out of scope
+
+- SQLx owning schema migrations, and removing Prisma
+- Converting unrelated routes to dependency injection

.agents/skills/database-access/references/rust-axum.md

@@ -0,0 +1,4 @@
+# Rust and Axum
+
+- This should use same `crates/db` what is exposed to Python world.
+- Details are TBD.

.github/workflows/test-unit.yml

@@ -577,6 +577,7 @@ jobs:
               tests/unit/proxy/rerank_endpoints
               tests/unit/proxy/realtime_endpoints
               tests/unit/proxy/ui_crud_endpoints
+              tests/unit/proxy/deps
               tests/unit/proxy/config_resolvers
               tests/unit/proxy/utils
             python-version: "3.12"

AGENTS.md

@@ -90,6 +90,8 @@ Do not put names of customers or customer company names in code, PR descriptions
 
 CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`). Download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
 
+For database operations, connection ownership and proxy dependency wiring, follow [database-access](.agents/skills/database-access/SKILL.md)
+
 Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: <what bounds it>`
 
 Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):

litellm-rust/crates/python-bridge/src/db.rs

@@ -0,0 +1,19 @@
+use litellm_host_python::ready_future;
+use pyo3::{prelude::*, types::PyNone};
+
+#[pyclass(frozen)]
+pub(crate) struct NativeDatabase {}
+
+// TODO: take typed writer/reader settings and a token provider, then open a sqlx PgPool, once the first
+// repository moves to native. The Prisma DATABASE_URL carries Prisma-only params and rotating IAM tokens
+#[pymethods]
+impl NativeDatabase {
+    #[staticmethod]
+    fn connect(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
+        ready_future(py, Bound::new(py, NativeDatabase {})?.as_any())
+    }
+
+    fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        ready_future(py, PyNone::get(py).as_any())
+    }
+}

litellm-rust/crates/python-bridge/src/lib.rs

@@ -2,6 +2,7 @@ mod cache;
 mod callable;
 mod coercion;
 mod credentials;
+mod db;
 mod diagnostics;
 mod errors;
 mod execution;
@@ -21,6 +22,8 @@ mod _native {
     use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
     use pyo3::{prelude::*, types::PyModule};
 
+    #[pymodule_export]
+    use crate::db::NativeDatabase;
     #[cfg(feature = "panic-test")]
     #[pymodule_export]
     use crate::diagnostics::_panic_for_test;
@@ -112,6 +115,7 @@ mod tests {
                 "NativeDiagnosticProcessor",
                 "NativeClickHouseSpendConfig",
                 "NativeClickHouseSpendStorage",
+                "NativeDatabase",
                 "trace_encode_error",
                 "bundled_model_cost_map",
                 "TokenCounter",

litellm/models/user_banner.py

@@ -0,0 +1,48 @@
+from typing import Final, Literal, Protocol
+
+from pydantic import ConfigDict, Field, model_validator
+
+from litellm.types.llms.base import LiteLLMBaseModel
+
+USER_BANNER_MAX_MESSAGE_LENGTH: Final = 4000
+UserBannerSeverity = Literal["info", "warning", "error"]
+
+
+class UserBannerUpdate(LiteLLMBaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    enabled: bool = Field(
+        default=False,
+        description="If true, the banner is shown to all authenticated dashboard users.",
+    )
+    message: str = Field(
+        default="",
+        max_length=USER_BANNER_MAX_MESSAGE_LENGTH,
+        description="Banner text shown to dashboard users. Markdown is supported.",
+    )
+    severity: UserBannerSeverity = Field(
+        default="info",
+        description="Visual style of the banner.",
+    )
+
+    @model_validator(mode="after")
+    def _require_message_when_enabled(self) -> "UserBannerUpdate":
+        if self.enabled and not self.message.strip():
+            raise ValueError("message must be non-empty when the banner is enabled")
+        return self
+
+
+class UserBanner(UserBannerUpdate):
+    revision: str = Field(
+        default="",
+        description=(
+            "Server-stamped opaque publish identity; a fresh value is generated on every "
+            "update so clients re-surface dismissed banners on republish."
+        ),
+    )
+
+
+class UserBannerStore(Protocol):
+    async def get_banner(self) -> UserBanner: ...
+
+    async def save_banner(self, banner: UserBanner) -> None: ...

litellm/proxy/deps/database.py

@@ -0,0 +1,22 @@
+from typing import Annotated, Final, TypeAlias
+
+from fastapi import Depends, Request
+
+from litellm.repositories.database import Database
+
+_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:
+    database: Final[object] = getattr(request.state, "database", _MISSING_DATABASE)
+    if database is None or isinstance(database, Database):
+        return database
+    raise DatabaseStateMissingError()
+
+
+DatabaseDep: TypeAlias = Annotated[Database | None, Depends(get_database)]

litellm/proxy/deps/repositories.py

@@ -0,0 +1,14 @@
+from typing import Annotated, TypeAlias
+
+from fastapi import Depends
+
+from litellm.models.user_banner import UserBannerStore
+from litellm.proxy.deps.database import DatabaseDep
+from litellm.repositories.user_banner_repository import UserBannerRepository
+
+
+async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None:
+    return None if database is None else UserBannerRepository(database.prisma)
+
+
+UserBannerRepositoryDep: TypeAlias = Annotated[UserBannerStore | None, Depends(get_user_banner_repository)]

litellm/proxy/proxy_server.py

@@ -1479,21 +1479,21 @@ 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:
-        _own_database(resources, client)
+        resources.push_async_callback(client.aclose)
     await client.connect()
     return client.writer_db
 
 
-def _own_database(resources: AsyncExitStack, client: PrismaClient | None) -> database_runtime.Database | None:
+async 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)
+    database: Final = await database_runtime.Database.open(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
+    global prisma_client  # noqa: PLW0603  # routes still read this global until they move to DatabaseDep
     prisma_client = None  # rebind-ok: a closed client must not survive into the next lifespan entry
 
 
@@ -1525,7 +1525,7 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener
     import json
 
     resources.callback(_forget_prisma_client)
-    existing_database: Final = _own_database(resources, prisma_client)
+    existing_database: Final = await _own_database(resources, prisma_client)
 
     init_verbose_loggers()
 
@@ -1635,7 +1635,7 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener
             )
         )
 
-    database: Final = existing_database or _own_database(resources, prisma_client)
+    database: Final = existing_database or await _own_database(resources, prisma_client)
 
     await migrate_if_requested(
         environ=os.environ,
@@ -12123,7 +12123,7 @@ class ProxyStartupEvent:
             return None
         client: Final = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj)
         async with AsyncExitStack() as setup_resources:
-            setup_resources.push_async_callback(database_runtime.Database(prisma=client).aclose)
+            setup_resources.push_async_callback(client.aclose)
             connected: Final = await cls._initialize_prisma_client(client)
             if connected is not None:
                 setup_resources.pop_all()

litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py

@@ -0,0 +1,31 @@
+from collections.abc import Coroutine
+from typing import Annotated, Protocol, TypeAlias
+
+from fastapi import Depends
+
+from litellm.models.user_banner import UserBanner
+from litellm.proxy._types import LitellmTableNames, UserAPIKeyAuth
+from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID
+
+
+class UserBannerAudit(Protocol):
+    def __call__(self, before: UserBanner, after: UserBanner, user: UserAPIKeyAuth) -> Coroutine[None, None, None]: ...
+
+
+async def get_user_banner_audit() -> UserBannerAudit:
+    from litellm.proxy.proxy_server import create_config_audit_log
+
+    async def audit(before: UserBanner, after: UserBanner, user: UserAPIKeyAuth) -> None:
+        await create_config_audit_log(
+            param_name=USER_BANNER_ROW_ID,
+            action="updated",
+            before_value=before.model_dump(),
+            after_value=after.model_dump(),
+            user_api_key_dict=user,
+            table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME,
+        )
+
+    return audit
+
+
+UserBannerAuditDep: TypeAlias = Annotated[UserBannerAudit, Depends(get_user_banner_audit)]

litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py

@@ -1,88 +1,43 @@
-import asyncio
-import json
-from typing import Annotated, Final, Literal
+from typing import Annotated, Final
 
-from fastapi import APIRouter, Depends, HTTPException
-from pydantic import Field, ValidationError, model_validator
+from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
+from pydantic import ConfigDict
 
 from litellm._uuid import uuid4
-from litellm.proxy._types import LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth
+from litellm.models.user_banner import UserBanner, UserBannerUpdate
+from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
-from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerRepository
+from litellm.proxy.deps.repositories import UserBannerRepositoryDep
+from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import (
+    UserBannerAuditDep,
+)
 from litellm.types.llms.base import LiteLLMBaseModel
 
 router: Final = APIRouter()
 
-USER_BANNER_MAX_MESSAGE_LENGTH: Final = 4000
-
-UserBannerSeverity = Literal["info", "warning", "error"]
-
-
-class UserBannerUpdate(LiteLLMBaseModel):
-    enabled: bool = Field(
-        default=False,
-        description="If true, the banner is shown to all authenticated dashboard users.",
-    )
-    message: str = Field(
-        default="",
-        max_length=USER_BANNER_MAX_MESSAGE_LENGTH,
-        description="Banner text shown to dashboard users. Markdown is supported.",
-    )
-    severity: UserBannerSeverity = Field(
-        default="info",
-        description="Visual style of the banner.",
-    )
-
-    @model_validator(mode="after")
-    def _require_message_when_enabled(self) -> "UserBannerUpdate":
-        if self.enabled and not self.message.strip():
-            raise ValueError("message must be non-empty when the banner is enabled")
-        return self
-
-
-class UserBanner(UserBannerUpdate):
-    revision: str = Field(
-        default="",
-        description=(
-            "Server-stamped opaque publish identity; a fresh value is generated on every "
-            "update so clients re-surface dismissed banners on republish."
-        ),
-    )
-
 
 class UpdateUserBannerResponse(LiteLLMBaseModel):
+    model_config = ConfigDict(frozen=True)
+
     message: str
     banner: UserBanner
 
 
-def parse_user_banner(raw_settings: object) -> UserBanner:
-    if raw_settings is None:
-        return UserBanner()
-    try:
-        parsed: Final = json.loads(raw_settings) if isinstance(raw_settings, str) else raw_settings
-        return UserBanner.model_validate(parsed)
-    except (json.JSONDecodeError, ValidationError):
-        return UserBanner()
-
-
 @router.get(
     "/get/user_banner",
     tags=["UI Settings"],
     dependencies=[Depends(user_api_key_auth)],
     response_model=UserBanner,
 )
-async def get_user_banner() -> UserBanner:
+async def get_user_banner(banners: UserBannerRepositoryDep) -> UserBanner:
     """
     Get the admin-published dashboard banner.
     Readable by any authenticated user; rendered on every dashboard page.
     """
-    from litellm.proxy.proxy_server import prisma_client
-
-    if prisma_client is None:
+    if banners is None:
         return UserBanner()
 
-    raw_settings: Final = await UserBannerRepository(prisma_client).get_raw_settings()
-    return parse_user_banner(raw_settings)
+    return await banners.get_banner()
 
 
 @router.patch(
@@ -93,21 +48,21 @@ async def get_user_banner() -> UserBanner:
 async def update_user_banner(
     banner_update: UserBannerUpdate,
     user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
+    banners: UserBannerRepositoryDep,
+    audit: UserBannerAuditDep,
+    background_tasks: BackgroundTasks,
 ) -> UpdateUserBannerResponse:
     """
     Publish, edit, or unpublish the dashboard banner.
     Only proxy admins are allowed to modify it.
     """
-    from litellm.proxy.proxy_server import create_config_audit_log, prisma_client
-
     if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
         raise HTTPException(status_code=403, detail="Only proxy admins can update the user banner.")
 
-    if prisma_client is None:
+    if banners is None:
         raise HTTPException(status_code=500, detail="Database not connected. Please connect a database.")
 
-    repository: Final = UserBannerRepository(prisma_client)
-    before: Final = parse_user_banner(await repository.get_raw_settings())
+    before: Final = await banners.get_banner()
     banner: Final = UserBanner(
         enabled=banner_update.enabled,
         message=banner_update.message,
@@ -115,17 +70,8 @@ async def update_user_banner(
         revision=uuid4().hex,
     )
 
-    await repository.upsert_settings(json.dumps(banner.model_dump()))
+    await banners.save_banner(banner)
 
-    asyncio.create_task(
-        create_config_audit_log(
-            param_name=USER_BANNER_ROW_ID,
-            action="updated",
-            before_value=before.model_dump(),
-            after_value=banner.model_dump(),
-            user_api_key_dict=user_api_key_dict,
-            table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME,
-        )
-    )
+    background_tasks.add_task(audit, before, banner, user_api_key_dict)
 
     return UpdateUserBannerResponse(message="User banner updated successfully", banner=banner)

litellm/proxy/utils.py

@@ -6692,6 +6692,17 @@ class PrismaClient:
             waiter.cancel()
             verbose_proxy_logger.warning("Shutdown stopped waiting for an in-flight DB reconnect")
 
+    async def aclose(self) -> None:
+        verbose_proxy_logger.debug("Disconnecting from Prisma")
+        try:
+            await self.stop_background_tasks()
+        except Exception:
+            verbose_proxy_logger.exception("Stopping database background tasks failed")
+        try:
+            await self.disconnect()
+        except Exception:
+            verbose_proxy_logger.exception("Database disconnect failed")
+
     async def stop_background_tasks(self) -> None:
         self._closing = True
         self._stop_engine_watcher()

litellm/repositories/database.py

@@ -4,7 +4,7 @@ from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Final
 
-from litellm._logging import verbose_proxy_logger
+from litellm.rust_bridge._native import NativeDatabase
 
 if TYPE_CHECKING:
     from litellm.proxy.utils import PrismaClient
@@ -13,17 +13,22 @@ if TYPE_CHECKING:
 @dataclass(frozen=True, slots=True)
 class Database:
     prisma: "PrismaClient"
+    native: NativeDatabase
+
+    @classmethod
+    async def open(cls, prisma: "PrismaClient") -> "Database":
+        try:
+            native: Final = await NativeDatabase.connect()
+        except BaseException:
+            await prisma.aclose()
+            raise
+        return cls(prisma=prisma, native=native)
 
     async def aclose(self) -> None:
-        verbose_proxy_logger.debug("Disconnecting from Prisma")
         try:
-            await self.prisma.stop_background_tasks()
-        except Exception:
-            verbose_proxy_logger.exception("Stopping database background tasks failed")
-        try:
-            await self.prisma.disconnect()
-        except Exception:
-            verbose_proxy_logger.exception("Database disconnect failed")
+            await self.prisma.aclose()
+        finally:
+            await self.native.close()
 
 
 @asynccontextmanager

litellm/repositories/user_banner_repository.py

@@ -1,5 +1,9 @@
+import json
 from typing import TYPE_CHECKING, Final
 
+from pydantic import ValidationError
+
+from litellm.models.user_banner import UserBanner
 from litellm.repositories.table_repositories import PrismaTableRepository
 
 if TYPE_CHECKING:
@@ -8,9 +12,25 @@ if TYPE_CHECKING:
 USER_BANNER_ROW_ID: Final = "user_banner"
 
 
+def parse_user_banner(raw_settings: object) -> UserBanner:
+    if raw_settings is None:
+        return UserBanner()
+    try:
+        parsed: Final = json.loads(raw_settings) if isinstance(raw_settings, str) else raw_settings
+        return UserBanner.model_validate(parsed)
+    except (json.JSONDecodeError, ValidationError):
+        return UserBanner()
+
+
 class UserBannerRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]):
     table_name = "litellm_uisettings"
 
+    async def get_banner(self) -> UserBanner:
+        return parse_user_banner(await self.get_raw_settings())
+
+    async def save_banner(self, banner: UserBanner) -> None:
+        await self.upsert_settings(json.dumps(banner.model_dump()))
+
     async def get_raw_settings(self) -> object:
         db_record: Final = await self.table.find_unique(where={"id": USER_BANNER_ROW_ID})
         return db_record.ui_settings if db_record is not None else None

litellm/rust_bridge/_native.pyi

@@ -26,6 +26,12 @@ class NativeClickHouseSpendStorage:
     def ensure_schema(self) -> Future[None]: ...
     def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Future[None]: ...
 
+@final
+class NativeDatabase:
+    @staticmethod
+    def connect() -> Future[NativeDatabase]: ...
+    def close(self) -> Future[None]: ...
+
 @final
 class NativeDiagnosticProcessor:
     def __new__(cls, minimum_custom_key_length: int) -> NativeDiagnosticProcessor: ...
@@ -208,6 +214,7 @@ __all__ = [
     "HuggingFaceEncoding",
     "NativeClickHouseSpendConfig",
     "NativeClickHouseSpendStorage",
+    "NativeDatabase",
     "NativeDiagnosticProcessor",
     "ProcessReservedForForking",
     "ResponsesWebSocketConnection",

tests/_support/native_database.py

@@ -0,0 +1,11 @@
+import asyncio
+
+from litellm.rust_bridge._native import NativeDatabase
+
+
+async def _connect() -> NativeDatabase:
+    return await NativeDatabase.connect()
+
+
+def native_database() -> NativeDatabase:
+    return asyncio.run(_connect())

tests/integration/database/test_proxy_database_lifespan.py

@@ -10,6 +10,7 @@ from typing import Final, Literal, TypeAlias
 import psutil
 import pytest
 
+from litellm.rust_bridge._native import NativeDatabase
 from tests.integration._support.client import eventually
 from tests.integration._support.database import read_rows, scratch_database
 
@@ -79,6 +80,7 @@ async def _exercise_lifespans(mode: Mode) -> None:
                 async with proxy_server.proxy_startup_event(proxy_server.app) as state:
                     database: Final = state["database"]
                     assert database is not None
+                    assert isinstance(database.native, NativeDatabase)
                     assert await database.prisma.writer_db.query_raw("SELECT 42 AS answer") == [{"answer": 42}]
                     assert await database.prisma.read_db.query_raw("SELECT 43 AS answer") == [{"answer": 43}]
                     assert psutil.Process().children(recursive=True)

tests/unit/proxy/deps/test_database.py

@@ -0,0 +1,13 @@
+from typing import Final
+
+import pytest
+from starlette.requests import Request
+
+from litellm.proxy.deps.database import DatabaseStateMissingError, get_database
+
+
+@pytest.mark.asyncio
+async def test_missing_lifespan_is_not_treated_as_a_disabled_database():
+    request: Final = Request({"type": "http", "state": {}})
+    with pytest.raises(DatabaseStateMissingError):
+        await get_database(request)

tests/unit/proxy/deps/test_repositories.py

@@ -0,0 +1,63 @@
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from typing import Final
+from unittest.mock import AsyncMock, MagicMock
+
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from litellm.models.user_banner import UserBanner
+from litellm.proxy.deps.database import get_database
+from litellm.proxy.deps.repositories import UserBannerRepositoryDep
+from litellm.repositories import database as database_runtime
+from tests._support.native_database import native_database
+
+NATIVE: Final = native_database()
+
+
+def _app(database: database_runtime.Database | None) -> FastAPI:
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        yield {"database": database}
+
+    app: Final = FastAPI(lifespan=lifespan)
+
+    @app.get("/banner")
+    async def banner(banners: UserBannerRepositoryDep) -> UserBanner:
+        return UserBanner() if banners is None else await banners.get_banner()
+
+    return app
+
+
+def test_lifespan_database_is_shared_until_shutdown():
+    prisma: Final = MagicMock()
+    prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None)
+    app: Final = _app(database_runtime.Database(prisma=prisma, native=NATIVE))
+    with TestClient(app) as client:
+        assert client.get("/banner").json() == UserBanner().model_dump()
+        prisma.db.litellm_uisettings.find_unique.return_value = SimpleNamespace(
+            ui_settings={"enabled": True, "message": "published", "revision": "revision"}
+        )
+        assert (
+            client.get("/banner").json()
+            == UserBanner(enabled=True, message="published", revision="revision").model_dump()
+        )
+
+
+def test_no_database_state_preserves_optional_repository():
+    with TestClient(_app(None)) as client:
+        response: Final = client.get("/banner")
+        assert response.status_code == 200
+        assert response.json() == UserBanner().model_dump()
+
+
+def test_dependency_override_changes_next_request_without_changing_lifespan_owner(monkeypatch):
+    prisma: Final = MagicMock()
+    prisma.db.litellm_uisettings.find_unique = AsyncMock(
+        return_value=SimpleNamespace(ui_settings={"message": "stored"})
+    )
+    app: Final = _app(database_runtime.Database(prisma=prisma, native=NATIVE))
+    with TestClient(app) as client:
+        assert client.get("/banner").json() == UserBanner(message="stored").model_dump()
+        monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None)
+        assert client.get("/banner").json() == UserBanner().model_dump()

tests/unit/proxy/proxy_server/test_lifecycle.py

@@ -23,8 +23,8 @@ import json
 import logging
 import os
 import subprocess
-from collections.abc import Awaitable, Callable
-from contextlib import aclosing
+from collections.abc import AsyncIterator, Awaitable, Callable
+from contextlib import asynccontextmanager
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -51,10 +51,19 @@ from litellm.proxy.proxy_server import (
     run_ollama_serve,
     save_worker_config,
 )
-from litellm.repositories import database as database_runtime
+from litellm.proxy.utils import PrismaClient
 
 from .conftest import normalize
 
+
+@asynccontextmanager
+async def _closing_prisma(prisma: MagicMock) -> AsyncIterator[None]:
+    try:
+        yield
+    finally:
+        await PrismaClient.aclose(prisma)
+
+
 # ---------------------------------------------------------------------------
 # cleanup_router_config_variables
 # ---------------------------------------------------------------------------
@@ -119,7 +128,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 aclosing(database_runtime.Database(prisma=fake_prisma)):
+    async with _closing_prisma(fake_prisma):
         await proxy_shutdown_event()
         fake_prisma.disconnect.assert_not_awaited()
 
@@ -227,7 +236,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 aclosing(database_runtime.Database(prisma=fake_prisma)):
+    async with _closing_prisma(fake_prisma):
         await proxy_shutdown_event()
 
     assert calls == ["flush", "disconnect"]
@@ -269,7 +278,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 aclosing(database_runtime.Database(prisma=fake_prisma)):
+    async with _closing_prisma(fake_prisma):
         await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
 
     assert calls == ["flush", "audit", "disconnect"]
@@ -342,7 +351,7 @@ async def test_database_lifespan_logs_disconnect_error_after_shutdown_instead_of
     monkeypatch.setattr(litellm, "success_callback", [], raising=False)
 
     with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
-        async with aclosing(database_runtime.Database(prisma=fake_prisma)):
+        async with _closing_prisma(fake_prisma):
             await proxy_shutdown_event()
 
     fake_prisma.disconnect.assert_awaited_once()

tests/unit/proxy/test_proxy_server_endpoints_and_startup.py

@@ -1657,6 +1657,12 @@ class MockPrisma:
     async def stop_background_tasks(self):
         pass
 
+    async def aclose(self):
+        try:
+            await self.stop_background_tasks()
+        finally:
+            await self.disconnect()
+
 
 mock_prisma = MockPrisma()
 
@@ -1856,6 +1862,12 @@ class _PrismaClientWhoseDatabaseRejectsQueries:
     async def stop_background_tasks(self) -> None:
         pass
 
+    async def aclose(self) -> None:
+        try:
+            await self.stop_background_tasks()
+        finally:
+            await self.disconnect()
+
     class _Database:
         async def query_raw(self, query, *args):
             raise RuntimeError("permission denied for table LiteLLM_CredentialsTable")
@@ -14271,6 +14283,14 @@ def _mock_startup_prisma_client(health_check_error=None, connect_error=None):
     client.start_view_setup_task = MagicMock()
     client.start_db_health_watchdog_task = AsyncMock()
     client.health_check = AsyncMock(side_effect=health_check_error)
+
+    async def close_client() -> None:
+        try:
+            await client.stop_background_tasks()
+        finally:
+            await client.disconnect()
+
+    client.aclose = AsyncMock(side_effect=close_client)
     return client
 
 

tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py

@@ -1,5 +1,7 @@
+import asyncio
 import json
 from types import SimpleNamespace
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -7,8 +9,14 @@ from fastapi.testclient import TestClient
 
 from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.deps.database import get_database
+from litellm.proxy.deps.repositories import get_user_banner_repository
 from litellm.proxy.proxy_server import app
+from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import get_user_banner_audit
+from litellm.repositories.database import Database
+from litellm.repositories.user_banner_repository import UserBannerRepository
 from tests._master_key import MASTER_KEY
+from tests._support.native_database import native_database
 
 client = TestClient(app)
 
@@ -21,6 +29,8 @@ PUBLISHED_BANNER = {
 PUBLISH_BODY = {k: v for k, v in PUBLISHED_BANNER.items() if k != "revision"}
 DISABLED_BANNER = {"enabled": False, "message": "", "severity": "info", "revision": ""}
 
+NATIVE: Final = native_database()
+
 
 def _auth_override(role: LitellmUserRoles):
     async def override() -> UserAPIKeyAuth:
@@ -46,7 +56,7 @@ def internal_user_auth():
 @pytest.fixture
 def mock_audit_log(monkeypatch):
     audit_mock = AsyncMock(return_value=None)
-    monkeypatch.setattr("litellm.proxy.proxy_server.create_config_audit_log", audit_mock)
+    monkeypatch.setitem(app.dependency_overrides, get_user_banner_audit, lambda: audit_mock)
     return audit_mock
 
 
@@ -54,7 +64,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.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
+    monkeypatch.setitem(app.dependency_overrides, get_database, lambda: Database(prisma=mock_prisma, native=NATIVE))
     return mock_prisma
 
 
@@ -94,21 +104,77 @@ class TestGetUserBanner:
         assert response.json() == DISABLED_BANNER
 
     def test_no_database_returns_disabled_banner(self, admin_auth, monkeypatch):
-        monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+        monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None)
         response = client.get("/get/user_banner")
         assert response.status_code == 200
         assert response.json() == DISABLED_BANNER
 
 
 class TestUpdateUserBanner:
+    def test_storage_finishes_before_audit_runs(self, admin_auth, monkeypatch, mock_audit_log):
+        events = []
+        prisma = _mock_prisma(monkeypatch)
+
+        async def save(**kwargs):
+            await asyncio.sleep(0)
+            events.append("saved")
+
+        async def audit(before, after, user):
+            events.append("audited")
+
+        prisma.db.litellm_uisettings.upsert.side_effect = save
+        mock_audit_log.side_effect = audit
+        response = client.patch("/update/user_banner", json=PUBLISH_BODY)
+        assert response.status_code == 200
+        assert events == ["saved", "audited"]
+
+    @pytest.mark.parametrize(
+        ("role", "status", "detail"),
+        [
+            (LitellmUserRoles.INTERNAL_USER, 403, "Only proxy admins can update the user banner."),
+            (LitellmUserRoles.PROXY_ADMIN, 500, "Database not connected. Please connect a database."),
+        ],
+    )
+    def test_no_database_preserves_authorization_precedence(self, monkeypatch, role, status, detail):
+        monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override(role))
+        monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None)
+        response = client.patch("/update/user_banner", json=PUBLISH_BODY)
+        assert response.status_code == status
+        assert response.json() == {"detail": detail}
+
+    def test_failed_write_does_not_schedule_success_audit(self, admin_auth, monkeypatch, mock_audit_log):
+        prisma = _mock_prisma(monkeypatch)
+        prisma.db.litellm_uisettings.upsert.side_effect = RuntimeError("storage unavailable")
+        with pytest.raises(RuntimeError, match="storage unavailable"):
+            client.patch("/update/user_banner", json=PUBLISH_BODY)
+        mock_audit_log.assert_not_called()
+
+    def test_audit_receives_previous_and_saved_banner(self, admin_auth, monkeypatch, mock_audit_log):
+        prisma = _mock_prisma(monkeypatch, record=SimpleNamespace(ui_settings=json.dumps(PUBLISHED_BANNER)))
+        response = client.patch("/update/user_banner", json=PUBLISH_BODY)
+        assert response.status_code == 200
+        before, after, user = mock_audit_log.call_args.args
+        assert before.model_dump() == PUBLISHED_BANNER
+        assert after.model_dump() == response.json()["banner"]
+        assert user.user_id == "test-user"
+        assert (
+            json.loads(prisma.db.litellm_uisettings.upsert.await_args.kwargs["data"]["create"]["ui_settings"])
+            == after.model_dump()
+        )
+
     def test_rejects_non_admin(self, internal_user_auth, monkeypatch):
         mock_prisma = _mock_prisma(monkeypatch)
         response = client.patch("/update/user_banner", json=PUBLISH_BODY)
         assert response.status_code == 403
         mock_prisma.db.litellm_uisettings.upsert.assert_not_awaited()
 
-    def test_persists_and_round_trips(self, admin_auth, monkeypatch, mock_audit_log):
-        mock_prisma = _mock_prisma(monkeypatch, record=None)
+    @pytest.mark.parametrize("override_repository", [False, True])
+    def test_persists_and_round_trips(self, admin_auth, monkeypatch, mock_audit_log, override_repository: bool):
+        mock_prisma: Final = _mock_prisma(monkeypatch, record=None)
+        if override_repository:
+            banners: Final = UserBannerRepository(mock_prisma)
+            monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None)
+            monkeypatch.setitem(app.dependency_overrides, get_user_banner_repository, lambda: banners)
 
         response = client.patch("/update/user_banner", json=PUBLISH_BODY)
         assert response.status_code == 200

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

@@ -21,7 +21,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import pytest
 
 from litellm.proxy.utils import _SHUTDOWN_RECONNECT_WAIT_SECONDS, PrismaClient
-from litellm.repositories.database import Database
 
 
 @pytest.mark.asyncio
@@ -47,7 +46,7 @@ async def test_shutdown_waits_for_hung_request_reconnect_to_time_out(
             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())
+        closing: Final = asyncio.create_task(prisma_client.aclose())
         try:
             await asyncio.sleep(0)
             assert not closing.done()
@@ -75,7 +74,7 @@ async def test_shutdown_stops_waiting_for_a_held_reconnect_lock_after_the_fixed_
     assert _SHUTDOWN_RECONNECT_WAIT_SECONDS < prisma_client._db_watchdog_reconnect_timeout_seconds
     await prisma_client._db_reconnect_lock.acquire()
     with patch.object(asyncio.get_running_loop(), "time", return_value=1000.0) as clock:
-        closing: Final = asyncio.create_task(Database(prisma=prisma_client).aclose())
+        closing: Final = asyncio.create_task(prisma_client.aclose())
         try:
             await asyncio.sleep(0)
             assert not closing.done()
@@ -598,3 +597,32 @@ async def test_start_view_setup_task_twice_keeps_the_first_task(prisma_client: P
         "second_is_first": second is first,
     }
     assert actual == {"first_is_task": True, "second_is_first": True}
+
+
+@pytest.mark.asyncio
+async def test_aclose_disconnects_and_logs_when_stopping_background_tasks_fails(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    prisma: Final = MagicMock()
+    prisma.stop_background_tasks = AsyncMock(side_effect=RuntimeError("watchdog failed"))
+    prisma.disconnect = AsyncMock()
+
+    with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
+        await PrismaClient.aclose(prisma)
+
+    prisma.disconnect.assert_awaited_once()
+    assert "watchdog failed" in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_aclose_returns_and_logs_when_disconnect_fails(caplog: pytest.LogCaptureFixture) -> None:
+    prisma: Final = MagicMock()
+    prisma.stop_background_tasks = AsyncMock()
+    prisma.disconnect = AsyncMock(side_effect=RuntimeError("engine hung"))
+
+    with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
+        await PrismaClient.aclose(prisma)
+
+    prisma.stop_background_tasks.assert_awaited_once()
+    prisma.disconnect.assert_awaited_once()
+    assert "engine hung" in caplog.text

tests/unit/repositories/test_database.py

@@ -1,4 +1,3 @@
-import logging
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock
 
@@ -8,29 +7,12 @@ from litellm.repositories.database import Database
 
 
 @pytest.mark.asyncio
-async def test_database_disconnects_and_logs_when_stopping_background_tasks_fails(
-    caplog: pytest.LogCaptureFixture,
-) -> None:
+async def test_closing_the_database_closes_prisma() -> None:
     prisma: Final = MagicMock()
-    prisma.stop_background_tasks = AsyncMock(side_effect=RuntimeError("watchdog failed"))
-    prisma.disconnect = AsyncMock()
+    prisma.aclose = AsyncMock()
+    database: Final = await Database.open(prisma)
 
-    with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
-        await Database(prisma=prisma).aclose()
+    prisma.aclose.assert_not_awaited()
+    await database.aclose()
 
-    prisma.disconnect.assert_awaited_once()
-    assert "watchdog failed" in caplog.text
-
-
-@pytest.mark.asyncio
-async def test_database_aclose_returns_and_logs_when_disconnect_fails(caplog: pytest.LogCaptureFixture) -> None:
-    prisma: Final = MagicMock()
-    prisma.stop_background_tasks = AsyncMock()
-    prisma.disconnect = AsyncMock(side_effect=RuntimeError("engine hung"))
-
-    with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
-        await Database(prisma=prisma).aclose()
-
-    prisma.stop_background_tasks.assert_awaited_once()
-    prisma.disconnect.assert_awaited_once()
-    assert "engine hung" in caplog.text
+    prisma.aclose.assert_awaited_once()