Merge cebc7e6dd7 into 14193435b8

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

@@ -0,0 +1,40 @@
+---
+name: proxy-dependencies
+description: Inject typed repositories and services into Python proxy handlers, or change database runtime ownership while preserving behavior for future Rust backends
+---
+
+# Proxy dependencies
+
+- 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
+
+- Inject named FastAPI providers through `Annotated[..., Depends(...)]`
+- Expose narrow repository protocols and typed service callables, not Prisma clients, models, tables or query dictionaries
+- Keep database I/O out of dependency construction
+- Keep no-database behavior and authorization precedence specific to the route
+- Ordinary helpers and background jobs receive explicit arguments, never unresolved `Depends` defaults
+- Override named providers in route tests and keep production parsing and persistence code real
+
+## Database ownership
+
+- `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
+- Legacy globals remain compatibility access for unmigrated code, not a second pool or owner
+- Never close a database from a request dependency or repository operation
+- SQLx integration is a separate explicit change, preserve transaction boundaries and do not replay failed writes through Prisma
+
+## Migration process
+
+- Inventory imports, direct callers, state writers, side effects and error ordering first
+- Migrate one complete feature and its helper calls before expanding
+- Preserve response models, OpenAPI, audit scheduling, configuration sync and serialization
+- Verify default behavior, authorization, invalid input, storage failure and shutdown ordering
+- State which legacy reads and backend leaks remain
+- Read [route migration](references/route-migration.md) for the banner example
+- Read [runtime ownership](references/runtime-ownership.md) before changing resource lifetime or reload behavior

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

@@ -0,0 +1,45 @@
+# Route migration
+
+## Existing pattern
+
+- `get_database` reads the lifespan-owned handle from request state
+- 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
+- SQLx can later implement the same typed operations without changing handler signatures
+
+## 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
+- 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
+- Keep repository extraction, backend switching and error normalization separate
+- Do not migrate every route in a large module just to establish the pattern
+- Do not invent providers for resources without consumers
+
+## Framework references
+
+- [Dependency composition and caching](https://fastapi.tiangolo.com/tutorial/dependencies/sub-dependencies/)
+- [Dependency overrides](https://fastapi.tiangolo.com/advanced/testing-dependencies/)

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

@@ -0,0 +1,51 @@
+# Runtime ownership
+
+## Current foundation
+
+- Prisma startup still owns connection setup, authentication refresh, health checks and migrations
+- 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
+- This does not remove all proxy globals or implement a Rust database backend
+
+## 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
+- Do not create a pool per request or close a shared pool from a borrowed handle
+- Preserve error, cancellation, timeout and transaction semantics across the bridge
+- Keep each atomic operation on one backend and transaction
+- Follow the root database-access skill for backend selection and schema ownership
+
+## Changing ownership
+
+- Inventory initialization, reload, tests, background jobs and shutdown together
+- Migrate every writer of a resource when moving its authoritative storage
+- Starlette lifespan state is shallow-copied into request state
+- `app.state` and yielded lifespan state are distinct mechanisms
+- A frozen container does not make its database client, router or cache immutable
+- Do not copy a replaceable global into state while leaving reload to update only the global
+
+## Runtime swapping
+
+- Backend selection at startup does not require hot swapping
+- Define whether a request captures one runtime generation or reads again during processing
+- Build and validate a replacement before publishing it, leaving the old generation usable on failure
+- Serialize reloads and coordinate publication with shutdown
+- Drain streams and background users before retiring old resources
+- Retaining every old generation until shutdown is simple but can grow resource use without bounds
+- Define multi-worker coordination separately from an in-process swap
+- Axum `State` and `FromRef` can expose the same narrow dependencies to Rust handlers
+- Do not prescribe `ArcSwap` before defining publication, capture and retirement semantics
+
+## Framework references
+
+- [Starlette lifespan](https://www.starlette.io/lifespan/)
+- [Axum state](https://docs.rs/axum/latest/axum/extract/struct.State.html)

AGENTS.md

@@ -1,3 +1,5 @@
+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)
 - used as an input for tools to read and act on. For example:

litellm-rust/Cargo.lock

@@ -4255,6 +4255,7 @@ dependencies = [
  "litellm-cache-response",
  "litellm-callbacks-legacy-python",
  "litellm-core-utils",
+ "litellm-db",
  "litellm-host",
  "litellm-host-python",
  "litellm-http",

litellm-rust/crates/db/.sqlx/query-136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae.json

@@ -0,0 +1,26 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "SELECT ui_settings\nFROM \"LiteLLM_UISettings\"\nWHERE id = 'user_banner'\n",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "ui_settings",
+        "type_info": "Jsonb",
+        "origin": {
+          "Table": {
+            "table": "\"LiteLLM_UISettings\"",
+            "name": "ui_settings"
+          }
+        }
+      }
+    ],
+    "parameters": {
+      "Left": []
+    },
+    "nullable": [
+      false
+    ]
+  },
+  "hash": "136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae"
+}

litellm-rust/crates/db/.sqlx/query-c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc.json

@@ -0,0 +1,14 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "INSERT INTO \"LiteLLM_UISettings\" (id, ui_settings, updated_at)\nVALUES ('user_banner', $1, CURRENT_TIMESTAMP)\nON CONFLICT (id) DO UPDATE\nSET ui_settings = EXCLUDED.ui_settings,\n    updated_at = EXCLUDED.updated_at\n",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Left": [
+        "Jsonb"
+      ]
+    },
+    "nullable": []
+  },
+  "hash": "c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc"
+}

litellm-rust/crates/db/queries/user_banner/find_user_banner.sql

@@ -0,0 +1,3 @@
+SELECT ui_settings
+FROM "LiteLLM_UISettings"
+WHERE id = 'user_banner'

litellm-rust/crates/db/queries/user_banner/save_user_banner.sql

@@ -0,0 +1,5 @@
+INSERT INTO "LiteLLM_UISettings" (id, ui_settings, updated_at)
+VALUES ('user_banner', $1, CURRENT_TIMESTAMP)
+ON CONFLICT (id) DO UPDATE
+SET ui_settings = EXCLUDED.ui_settings,
+    updated_at = EXCLUDED.updated_at

litellm-rust/crates/db/src/lib.rs

@@ -13,6 +13,7 @@ mod failure;
 mod spend;
 #[cfg(test)]
 mod test_support;
+mod user_banner;
 
 pub use daily_spend::{DailyEndUserRow, DailyTeamRow, DailyUserRow};
 pub use end_user::{EndUser, RateLimits, find_end_user};
@@ -24,3 +25,4 @@ pub use spend::{
     begin, insert_spend_logs, lock_team, upsert_daily_end_user_spend, upsert_daily_spend,
     upsert_daily_team_spend, upsert_daily_user_spend,
 };
+pub use user_banner::{find_user_banner, save_user_banner};

litellm-rust/crates/db/src/user_banner.rs

@@ -0,0 +1,114 @@
+use serde_json::Value;
+use sqlx::PgExecutor;
+
+use crate::Error;
+
+pub async fn find_user_banner(executor: impl PgExecutor<'_>) -> Result<Option<Value>, Error> {
+    Ok(
+        sqlx::query_file_scalar!("queries/user_banner/find_user_banner.sql")
+            .fetch_optional(executor)
+            .await?,
+    )
+}
+
+pub async fn save_user_banner(
+    executor: impl PgExecutor<'_>,
+    settings: &Value,
+) -> Result<(), Error> {
+    sqlx::query_file!("queries/user_banner/save_user_banner.sql", settings)
+        .execute(executor)
+        .await?;
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use std::error::Error as StdError;
+
+    use rstest::rstest;
+    use serde_json::{Value, json};
+    use sqlx::PgPool;
+    use sqlx::types::chrono::NaiveDateTime;
+
+    use super::{find_user_banner, save_user_banner};
+    use crate::test_support::database_pool;
+
+    async fn clear(tx: &mut sqlx::PgConnection) -> Result<(), sqlx::Error> {
+        #[expect(
+            clippy::disallowed_methods,
+            reason = "Fixture cleanup is outside the compile-checked query set"
+        )]
+        sqlx::query(r#"DELETE FROM "LiteLLM_UISettings" WHERE id = 'user_banner'"#)
+            .execute(tx)
+            .await?;
+        Ok(())
+    }
+
+    async fn stored_row(
+        tx: &mut sqlx::PgConnection,
+    ) -> Result<(Value, NaiveDateTime, NaiveDateTime), sqlx::Error> {
+        #[expect(
+            clippy::disallowed_methods,
+            reason = "Fixture reads are outside the compile-checked query set"
+        )]
+        sqlx::query_as(
+            r#"SELECT ui_settings, created_at, updated_at FROM "LiteLLM_UISettings" WHERE id = 'user_banner'"#,
+        )
+        .fetch_one(tx)
+        .await
+    }
+
+    #[rstest]
+    #[tokio::test]
+    #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
+    async fn missing_banner_reads_as_none(
+        #[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
+    ) -> Result<(), Box<dyn StdError>> {
+        let mut tx = database_pool?.begin().await?;
+        clear(&mut tx).await?;
+
+        assert_eq!(find_user_banner(&mut *tx).await?, None);
+        Ok(())
+    }
+
+    #[rstest]
+    #[case::first_publish(None)]
+    #[case::republish(Some(json!({"enabled": false, "message": "old", "severity": "info", "revision": "a"})))]
+    #[tokio::test]
+    #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
+    async fn save_replaces_the_stored_banner_and_stamps_updated_at(
+        #[case] previous: Option<Value>,
+        #[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
+    ) -> Result<(), Box<dyn StdError>> {
+        let mut tx = database_pool?.begin().await?;
+        clear(&mut tx).await?;
+        if let Some(previous) = &previous {
+            save_user_banner(&mut *tx, previous).await?;
+        }
+        let banner = json!({"enabled": true, "message": "maintenance at 5pm", "severity": "warning", "revision": "b"});
+
+        save_user_banner(&mut *tx, &banner).await?;
+
+        assert_eq!(find_user_banner(&mut *tx).await?, Some(banner.clone()));
+        let (stored, created_at, updated_at) = stored_row(&mut tx).await?;
+        assert_eq!(stored, banner);
+        assert!(updated_at >= created_at);
+        Ok(())
+    }
+
+    #[rstest]
+    #[tokio::test]
+    #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
+    async fn banner_text_is_stored_as_data_not_sql(
+        #[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
+    ) -> Result<(), Box<dyn StdError>> {
+        let mut tx = database_pool?.begin().await?;
+        clear(&mut tx).await?;
+        let banner = json!({"message": r#"'); DROP TABLE "LiteLLM_UISettings"; --"#});
+
+        save_user_banner(&mut *tx, &banner).await?;
+
+        assert_eq!(find_user_banner(&mut *tx).await?, Some(banner));
+        Ok(())
+    }
+}

litellm-rust/crates/python-bridge/Cargo.toml

@@ -19,18 +19,14 @@ huggingface = ["litellm-token-counter/huggingface"]
 tiktoken = ["litellm-token-counter/tiktoken"]
 
 [dependencies]
-fancy-regex.workspace = true
 litellm-tracing.workspace = true
 litellm-spend-clickhouse.workspace = true
 litellm-storage-clickhouse.workspace = true
 litellm-host.workspace = true
-bytes.workspace = true
-futures-util.workspace = true
 litellm-cache.workspace = true
 litellm-cache-memory.workspace = true
 litellm-cache-redis.workspace = true
 litellm-cache-response.workspace = true
-serde.workspace = true
 litellm-auth.workspace = true
 litellm-auth-aws.workspace = true
 litellm-callbacks-legacy-python.workspace = true
@@ -41,6 +37,7 @@ litellm-inference-messages.workspace = true
 litellm-inference-chat.workspace = true
 litellm-inference-ocr.workspace = true
 litellm-core-utils.workspace = true
+litellm-db.workspace = true
 litellm-http.workspace = true
 litellm-llms.workspace = true
 litellm-secrets = { workspace = true, features = ["aws", "azure", "google", "hashicorp", "cyberark"] }
@@ -49,6 +46,11 @@ litellm-llms-types.workspace = true
 litellm-router-types.workspace = true
 litellm-host-python.workspace = true
 litellm-token-counter = { path = "../token-counter", default-features = false }
+
+serde.workspace = true
+futures-util.workspace = true
+bytes.workspace = true
+fancy-regex.workspace = true
 pyo3.workspace = true
 prost.workspace = true
 pyo3-async-runtimes.workspace = true
@@ -61,9 +63,10 @@ thiserror.workspace = true
 tokio = { workspace = true, features = ["rt", "sync"] }
 
 [dev-dependencies]
-tracing.workspace = true
 litellm-http = { workspace = true, features = ["test-support"] }
 litellm-secrets-aws.workspace = true
+
+tracing.workspace = true
 serde.workspace = true
 serde_with.workspace = true
 criterion.workspace = true

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

@@ -0,0 +1,68 @@
+use std::error::Error as StdError;
+
+use litellm_db::Error;
+use pyo3::{exceptions::PyRuntimeError, prelude::*};
+use serde_json::Value;
+use sqlx::{PgPool, postgres::PgPoolOptions};
+
+fn map_error(error: Error) -> PyErr {
+    let chain: Vec<String> =
+        std::iter::successors(Some(&error as &dyn StdError), |error| error.source())
+            .map(ToString::to_string)
+            .collect();
+    PyRuntimeError::new_err(chain.join(": "))
+}
+
+#[pyclass(frozen)]
+pub struct NativeDatabase {
+    pool: PgPool,
+}
+
+#[pymethods]
+impl NativeDatabase {
+    #[staticmethod]
+    fn connect(py: Python<'_>, url: String, max_connections: u32) -> PyResult<Bound<'_, PyAny>> {
+        crate::execution::run_async_value(py, async move {
+            let pool = PgPoolOptions::new()
+                .max_connections(max_connections)
+                .connect(&url)
+                .await
+                .map_err(|error| map_error(error.into()))?;
+            Ok(Self { pool })
+        })
+    }
+
+    fn find_user_banner<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let pool = self.pool.clone();
+        crate::execution::run_async(
+            py,
+            async move { litellm_db::find_user_banner(&pool).await },
+            map_error,
+        )
+    }
+
+    fn save_user_banner<'py>(
+        &self,
+        py: Python<'py>,
+        #[pyo3(from_py_with = litellm_host_python::from_py_argument)] settings: Value,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let pool = self.pool.clone();
+        crate::execution::run_async(
+            py,
+            async move { litellm_db::save_user_banner(&pool, &settings).await },
+            map_error,
+        )
+    }
+
+    fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let pool = self.pool.clone();
+        crate::execution::run_async(
+            py,
+            async move {
+                pool.close().await;
+                Ok::<(), Error>(())
+            },
+            map_error,
+        )
+    }
+}

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

@@ -2,6 +2,7 @@ mod cache;
 mod callable;
 mod coercion;
 mod credentials;
+mod database;
 mod diagnostics;
 mod errors;
 mod execution;
@@ -20,6 +21,8 @@ mod _native {
     use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
     use pyo3::{prelude::*, types::PyModule};
 
+    #[pymodule_export]
+    use crate::database::NativeDatabase;
     #[cfg(feature = "panic-test")]
     #[pymodule_export]
     use crate::diagnostics::_panic_for_test;
@@ -87,6 +90,7 @@ mod tests {
         Python::attach(|py| {
             let mut expected = vec![
                 "NativeCacheHandle",
+                "NativeDatabase",
                 "RustBridgeDeclined",
                 "RustUpstreamError",
                 "ForkedAfterNativeRuntimeStarted",

litellm/_db.py

@@ -0,0 +1,30 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from litellm.proxy.utils import PrismaClient
+    from litellm.rust_bridge.database import NativeDatabase
+
+
+@dataclass(frozen=True, slots=True)
+class Database:
+    prisma: "PrismaClient"
+    native: "NativeDatabase | None" = None
+
+    async def aclose(self) -> None:
+        try:
+            if self.native is not None:
+                await self.native.close()
+        finally:
+            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()

litellm/constants.py

@@ -2350,3 +2350,4 @@ HARNESS_SNAPSHOT_SKIP_DIRS: Final = frozenset(
 )
 
 DEFAULT_TOOL_LOOP_MAX_ROUNDS: Final = 20
+NATIVE_DATABASE_MAX_CONNECTIONS: Final = 5

litellm/models/user_banner.py

@@ -0,0 +1,42 @@
+from typing import Final, Literal
+
+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."
+        ),
+    )

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._db 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/deps.py

@@ -0,0 +1,22 @@
+from typing import Annotated, Final, TypeAlias
+
+from fastapi import Depends, Request
+
+from litellm._db 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()
+
+
+DatabaseDependency: TypeAlias = Annotated[Database | None, Depends(get_database)]

litellm/proxy/proxy_server.py

@@ -63,6 +63,7 @@ from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, Valida
 from pydantic.fields import FieldInfo, PydanticUndefined
 from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never
 
+from litellm import _db
 from litellm._uuid import uuid
 from litellm.constants import (
     AIOHTTP_CONNECTOR_LIMIT,
@@ -85,6 +86,7 @@ from litellm.constants import (
     LITELLM_SETTINGS_SAFE_DB_OVERRIDES,
     LITELLM_UI_ALLOW_HEADERS,
     LITELLM_UI_SESSION_DURATION,
+    NATIVE_DATABASE_MAX_CONNECTIONS,
     OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS,
     RUNTIME_UPDATABLE_ROUTER_SETTINGS,
 )
@@ -213,6 +215,7 @@ if TYPE_CHECKING:
 
     from litellm.integrations.opentelemetry import OpenTelemetry
     from litellm.proxy.health_check_utils.shared_health_check_manager import SharedHealthCheckManager
+    from litellm.rust_bridge.database import NativeDatabase
 
     Span = _Span | Any
 else:
@@ -1288,13 +1291,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 +1422,23 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries:
     return client.writer_db
 
 
+async def _connect_native_database() -> "NativeDatabase | None":
+    database_url: Final = get_secret_str("DATABASE_URL")
+    if database_url is None or get_secret_bool("LITELLM_NATIVE_DATABASE") is not True:
+        return None
+    from litellm.rust_bridge.database import connect_native_database
+
+    return await connect_native_database(database_url, NATIVE_DATABASE_MAX_CONNECTIONS)
+
+
 @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 +1562,14 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState
             )
         )
 
+    database: Final = await resources.enter_async_context(
+        _db.lifespan(
+            None
+            if prisma_client is None
+            else _db.Database(prisma=prisma_client, native=await _connect_native_database())
+        )
+    )
+
     await migrate_if_requested(
         environ=os.environ,
         master_key=master_key,
@@ -1784,7 +1804,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/proxy/ui_crud_endpoints/user_banner_dependencies.py

@@ -0,0 +1,46 @@
+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.proxy.deps import DatabaseDependency
+from litellm.repositories.user_banner_repository import (
+    USER_BANNER_ROW_ID,
+    NativeUserBannerRepository,
+    UserBannerRepository,
+    UserBannerStore,
+)
+
+
+class UserBannerAudit(Protocol):
+    def __call__(self, before: UserBanner, after: UserBanner, user: UserAPIKeyAuth) -> Coroutine[None, None, None]: ...
+
+
+async def get_user_banner_repository(database: DatabaseDependency) -> UserBannerStore | None:
+    if database is None:
+        return None
+    if database.native is not None:
+        return NativeUserBannerRepository(database.native)
+    return UserBannerRepository(database.prisma)
+
+
+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
+
+
+UserBannerRepositoryDependency: TypeAlias = Annotated[UserBannerStore | None, Depends(get_user_banner_repository)]
+UserBannerAuditDependency: 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.ui_crud_endpoints.user_banner_dependencies import (
+    UserBannerAuditDependency,
+    UserBannerRepositoryDependency,
+)
 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(repository: UserBannerRepositoryDependency) -> 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 repository is None:
         return UserBanner()
 
-    raw_settings: Final = await UserBannerRepository(prisma_client).get_raw_settings()
-    return parse_user_banner(raw_settings)
+    return await repository.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)],
+    repository: UserBannerRepositoryDependency,
+    audit: UserBannerAuditDependency,
+    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 repository 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 repository.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 repository.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/repositories/user_banner_repository.py

@@ -1,16 +1,56 @@
-from typing import TYPE_CHECKING, Final
+import json
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Final, Protocol
 
+from pydantic import ValidationError
+
+from litellm.models.user_banner import UserBanner
 from litellm.repositories.table_repositories import PrismaTableRepository
 
 if TYPE_CHECKING:
     from prisma import models as prisma_models  # noqa: F401  # resolved only from the quoted base-class subscript below
 
+    from litellm.rust_bridge.database import NativeDatabase
+
 USER_BANNER_ROW_ID: Final = "user_banner"
 
 
+class UserBannerStore(Protocol):
+    async def get_banner(self) -> UserBanner: ...
+
+    async def save_banner(self, banner: UserBanner) -> None: ...
+
+
+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()
+
+
+@dataclass(frozen=True, slots=True)
+class NativeUserBannerRepository:
+    native: "NativeDatabase"
+
+    async def get_banner(self) -> UserBanner:
+        return parse_user_banner(await self.native.find_user_banner())
+
+    async def save_banner(self, banner: UserBanner) -> None:
+        await self.native.save_user_banner(banner.model_dump())
+
+
 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

@@ -25,6 +25,14 @@ 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(url: str, max_connections: int) -> Future[NativeDatabase]: ...
+    def close(self) -> Future[None]: ...
+    def find_user_banner(self) -> Future[JsonValue]: ...
+    def save_user_banner(self, settings: Mapping[str, JsonValue]) -> Future[None]: ...
+
 @final
 class NativeDiagnosticProcessor:
     def __new__(cls, minimum_custom_key_length: int) -> NativeDiagnosticProcessor: ...
@@ -207,6 +215,7 @@ __all__ = [
     "HuggingFaceEncoding",
     "NativeClickHouseSpendConfig",
     "NativeClickHouseSpendStorage",
+    "NativeDatabase",
     "NativeDiagnosticProcessor",
     "ProcessReservedForForking",
     "ResponsesWebSocketConnection",

litellm/rust_bridge/database.py

@@ -0,0 +1,33 @@
+from collections.abc import Awaitable, Mapping
+from typing import Final, Protocol, runtime_checkable
+
+from pydantic import ConfigDict, TypeAdapter
+
+from litellm.rust_bridge.loader import get_native_bridge
+
+
+class NativeDatabase(Protocol):
+    def find_user_banner(self) -> Awaitable[object]: ...
+    def save_user_banner(self, settings: Mapping[str, object]) -> Awaitable[None]: ...
+    def close(self) -> Awaitable[None]: ...
+
+
+class NativeDatabaseFactory(Protocol):
+    def connect(self, url: str, max_connections: int) -> Awaitable[NativeDatabase]: ...
+
+
+@runtime_checkable
+class NativeDatabaseBridge(Protocol):
+    NativeDatabase: NativeDatabaseFactory
+
+
+_NATIVE: Final[TypeAdapter[NativeDatabaseBridge]] = TypeAdapter(
+    NativeDatabaseBridge, config=ConfigDict(arbitrary_types_allowed=True)
+)
+
+
+async def connect_native_database(url: str, max_connections: int) -> NativeDatabase:
+    bridge: Final = get_native_bridge()
+    if bridge is None:
+        raise RuntimeError("The native database requires the Rust extension")
+    return await _NATIVE.validate_python(bridge).NativeDatabase.connect(url, max_connections)

tests/unit/proxy/proxy_server/test_lifecycle.py

@@ -33,6 +33,7 @@ from pydantic import BaseModel
 from typing_extensions import TypedDict
 
 import litellm.proxy.proxy_server as ps
+from litellm import _db
 from litellm.proxy.proxy_server import (
     ProxyStartupEvent,
     initialize_shared_aiohttp_session,
@@ -101,7 +102,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 +118,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 _db.lifespan(_db.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 +225,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 _db.lifespan(_db.Database(prisma=fake_prisma)):
+        await proxy_shutdown_event()
 
     assert calls == ["flush", "disconnect"]
 
@@ -262,7 +266,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 _db.lifespan(_db.Database(prisma=fake_prisma)):
+        await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
 
     assert calls == ["flush", "audit", "disconnect"]
 
@@ -318,7 +323,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 +338,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 _db.lifespan(_db.Database(prisma=fake_prisma)):
+            await proxy_shutdown_event()
 
 
 # ---------------------------------------------------------------------------
@@ -1060,7 +1066,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 +1152,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_deps.py

@@ -0,0 +1,76 @@
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from typing import Final
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from starlette.requests import Request
+
+from litellm import _db
+from litellm.models.user_banner import UserBanner
+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:
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        async with _db.lifespan(database) as owned_database:
+            yield {"database": owned_database}
+
+    app: Final = FastAPI(lifespan=lifespan)
+
+    @app.get("/banner")
+    async def banner(repository: UserBannerRepositoryDependency) -> UserBanner:
+        return UserBanner() if repository is None else await repository.get_banner()
+
+    return app
+
+
+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.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(
+            ui_settings={"enabled": True, "message": "published", "revision": "revision"}
+        )
+        assert (
+            client.get("/banner").json()
+            == UserBanner(enabled=True, message="published", revision="revision").model_dump()
+        )
+        prisma.disconnect.assert_not_awaited()
+    prisma.disconnect.assert_awaited_once()
+
+
+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.disconnect = AsyncMock()
+    prisma.db.litellm_uisettings.find_unique = AsyncMock(
+        return_value=SimpleNamespace(ui_settings={"message": "stored"})
+    )
+    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)
+        assert client.get("/banner").json() == UserBanner().model_dump()
+        prisma.disconnect.assert_not_awaited()
+    prisma.disconnect.assert_awaited_once()
+
+
+@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/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,7 @@ 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

tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock
@@ -5,9 +6,12 @@ from unittest.mock import AsyncMock, MagicMock
 import pytest
 from fastapi.testclient import TestClient
 
+from litellm import _db
 from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.deps import get_database
 from litellm.proxy.proxy_server import app
+from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import get_user_banner_audit
 from tests._master_key import MASTER_KEY
 
 client = TestClient(app)
@@ -46,7 +50,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 +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.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
+    monkeypatch.setitem(app.dependency_overrides, get_database, lambda: _db.Database(prisma=mock_prisma))
     return mock_prisma
 
 
@@ -94,13 +98,64 @@ 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)
@@ -174,3 +229,44 @@ class TestUpdateUserBanner:
         response = client.patch("/update/user_banner", json=payload)
         assert response.status_code == 422
         mock_prisma.db.litellm_uisettings.upsert.assert_not_awaited()
+
+
+class _NativeDatabase:
+    def __init__(self, stored=None):
+        self.stored = stored
+
+    async def find_user_banner(self):
+        return self.stored
+
+    async def save_user_banner(self, settings):
+        self.stored = settings
+
+
+class TestNativeUserBannerBackend:
+    @pytest.fixture
+    def backends(self, monkeypatch):
+        prisma = MagicMock()
+        banners = _NativeDatabase()
+        monkeypatch.setitem(app.dependency_overrides, get_database, lambda: _db.Database(prisma=prisma, native=banners))
+        return prisma, banners
+
+    def test_published_banner_round_trips_through_the_native_backend(self, admin_auth, mock_audit_log, backends):
+        prisma, banners = backends
+        published = client.patch("/update/user_banner", json=PUBLISH_BODY).json()["banner"]
+
+        assert banners.stored == published
+        assert client.get("/get/user_banner").json() == published
+        assert prisma.mock_calls == []
+
+    @pytest.mark.parametrize(
+        "stored", [json.dumps(PUBLISHED_BANNER), PUBLISHED_BANNER], ids=["prisma-string", "object"]
+    )
+    def test_native_backend_reads_rows_written_by_either_backend(self, admin_auth, backends, stored):
+        _, banners = backends
+        banners.stored = stored
+        assert client.get("/get/user_banner").json() == PUBLISHED_BANNER
+
+    def test_native_backend_treats_invalid_stored_banner_as_disabled(self, admin_auth, backends):
+        _, banners = backends
+        banners.stored = {"enabled": True, "message": ""}
+        assert client.get("/get/user_banner").json() == DISABLED_BANNER

tests/unit/test__db.py

@@ -0,0 +1,50 @@
+import asyncio
+from typing import Final
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm import _db
+
+
+@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 = _db.Database(prisma=prisma)
+
+    async def fail():
+        async with _db.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 _db.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_closing_disconnects_prisma_even_when_the_native_pool_fails_to_close():
+    prisma: Final = MagicMock()
+    prisma.disconnect = AsyncMock()
+    native: Final = MagicMock()
+    native.close = AsyncMock(side_effect=RuntimeError("native close failed"))
+
+    with pytest.raises(RuntimeError, match="native close failed"):
+        async with _db.lifespan(_db.Database(prisma=prisma, native=native)):
+            pass
+    native.close.assert_awaited_once()
+    prisma.disconnect.assert_awaited_once()