refactor: add lifespan-owned database dependency foundation

AGENTS.md

@@ -1,3 +1,5 @@
+For proxy dependency wiring and database lifetime changes, follow [proxy-dependencies](litellm-rust/.agents/skills/proxy-dependencies/SKILL.md) alongside [database-access](.agents/skills/database-access/SKILL.md)
+
 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/.agents/skills/proxy-dependencies/SKILL.md

@@ -0,0 +1,39 @@
+---
+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](../../../../.agents/skills/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` exposes typed repositories and an async close operation
+- `from_prisma` adopts an already initialized client without opening a second pool
+- 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

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

@@ -0,0 +1,44 @@
+# Route migration
+
+## Existing pattern
+
+- `get_database` reads the lifespan-owned handle from request state
+- A feature provider selects the narrow repository, such as `database.user_banners`
+- A 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 scheduled
+- 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/)

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

@@ -0,0 +1,48 @@
+# Runtime ownership
+
+## Current foundation
+
+- Prisma startup still owns connection setup, authentication refresh, health checks and migrations
+- Lifespan enters `_db.lifespan` with the initialized client adapted by `from_prisma`
+- The `Database` handle provides typed repositories and `aclose`, without exposing Prisma to handlers
+- 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
+- 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)

litellm-rust/AGENTS.md

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

litellm/_db.py

@@ -0,0 +1,43 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+from litellm.repositories.user_banner_repository import UserBannerRepository, UserBannerStore
+
+if TYPE_CHECKING:
+    from litellm.proxy.utils import PrismaClient
+
+
+@runtime_checkable
+class Database(Protocol):
+    @property
+    def user_banners(self) -> UserBannerStore: ...
+
+    async def aclose(self) -> None: ...
+
+
+@dataclass(frozen=True, slots=True)
+class _PrismaDatabase:
+    client: "PrismaClient"
+
+    @property
+    def user_banners(self) -> UserBannerStore:
+        return UserBannerRepository(self.client)
+
+    async def aclose(self) -> None:
+        await self.client.disconnect()
+
+
+def from_prisma(client: "PrismaClient") -> Database:
+    """Adapt an initialized client without opening another pool."""
+    return _PrismaDatabase(client)
+
+
+@asynccontextmanager
+async def lifespan(database: Database | None) -> AsyncGenerator[Database | None, None]:
+    try:
+        yield database
+    finally:
+        if database is not None:
+            await database.aclose()

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,18 @@
+from typing import Annotated, Final, TypeAlias
+
+from fastapi import Depends, Request
+from pydantic import ConfigDict, TypeAdapter
+
+from litellm._db import Database
+
+_DATABASE_ADAPTER: Final[TypeAdapter[Database | None]] = TypeAdapter(
+    Database | None, config=ConfigDict(arbitrary_types_allowed=True)
+)
+_MISSING_DATABASE: Final = object()
+
+
+async def get_database(request: Request) -> Database | None:
+    return _DATABASE_ADAPTER.validate_python(getattr(request.state, "database", _MISSING_DATABASE))
+
+
+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,
@@ -1288,13 +1289,7 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N
     if worker_heartbeat is not None and prisma_client:
         await worker_heartbeat.deregister()
     if prisma_client:
-        # Drain the SGR fold first: it lives in memory, so an un-drained interval
-        # is lost, and a write attempted after disconnect raises
-        # ClientNotConnectedError rather than persisting anything. Ordering this
-        # inside the same guard is what keeps the two from drifting apart.
         await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client))
-        verbose_proxy_logger.debug("Disconnecting from Prisma")
-        await prisma_client.disconnect()
     else:
         await _drain_audit_tasks_on_shutdown()
 
@@ -1427,6 +1422,12 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries:
 
 @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 +1551,10 @@ 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.from_prisma(prisma_client))
+    )
+
     await migrate_if_requested(
         environ=os.environ,
         master_key=master_key,
@@ -1784,7 +1789,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,37 @@
+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, 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:
+    return None if database is None else database.user_banners
+
+
+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,44 @@
 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 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 +49,20 @@ 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,
 ) -> 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,
-        )
-    )
+    asyncio.create_task(audit(before, banner, user_api_key_dict))
 
     return UpdateUserBannerResponse(message="User banner updated successfully", banner=banner)

litellm/repositories/user_banner_repository.py

@@ -1,5 +1,9 @@
-from typing import TYPE_CHECKING, Final
+import json
+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:
@@ -8,9 +12,31 @@ if TYPE_CHECKING:
 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()
+
+
 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

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.from_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.from_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.from_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.from_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,95 @@
+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 pydantic import ValidationError
+
+from litellm import _db
+from litellm.models.user_banner import UserBanner
+from litellm.proxy.deps import get_database
+from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import (
+    UserBannerRepositoryDependency,
+    get_user_banner_repository,
+)
+
+
+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.from_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.from_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():
+    from starlette.requests import Request
+
+    request: Final = Request({"type": "http", "state": {}})
+    with pytest.raises(ValidationError, match="Database"):
+        await get_database(request)
+
+
+@pytest.mark.asyncio
+async def test_constructing_repository_does_not_touch_database():
+    prisma: Final = MagicMock()
+    repository: Final = await get_user_banner_repository(_db.from_prisma(prisma))
+    assert repository is not None
+    assert prisma.mock_calls == []
+
+
+def test_dependencies_do_not_become_openapi_parameters():
+    operation: Final = _app(None).openapi()["paths"]["/banner"]["get"]
+    assert "parameters" not in operation
+    assert "requestBody" not in operation

tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py

@@ -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.from_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)

tests/unit/test__db.py

@@ -0,0 +1,36 @@
+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.from_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()