Merge 2465cddf22 into 2825f53ea0

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

@@ -1,6 +1,6 @@
 ---
 name: database-access
-description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope
+description: Organize database repositories and proxy dependency wiring, manage connection ownership, or port Prisma operations to Rust and SQLx without expanding migration scope
 ---
 
 # Database access
@@ -8,11 +8,44 @@ description: Add database queries, extract existing Prisma calls into repositori
 ## Adding new database access
 
 - Put Python database access behind an operation in `litellm/repositories`
+- Group operations by domain in `Repositories`, adding fields as their callers migrate
+- The container borrows existing handles and performs no database I/O during construction
+- Domain protocols are optional substitution boundaries, not required for every repository
 - Give callers typed inputs and outputs, not Prisma clients, models, or query dictionaries
-- Keep query execution and backend selection inside the repository
+- Keep query execution inside the repository, with one implementation per operation and no backend switch or fallback
 - Prisma is fine for the implementation, adding a query does not require a Rust port
 - Native Rust callers use `litellm-rust/crates/db` directly
 
+## Framework wiring
+
+- Python routes receive `RepositoriesDependency` and call domain operations such as `repositories.banners.get_banner()`
+- Use named FastAPI providers with `Annotated[..., Depends(...)]`, ordinary helpers and jobs receive explicit arguments
+- Preserve each route's no-database behavior, authorization precedence and background-task scheduling
+- Read [Python and FastAPI](references/python-fastapi.md) for provider examples, lifespan state and route verification
+- Read [Rust and Axum](references/rust-axum.md) when adding database access to Rust handlers
+
+## Connection ownership
+
+- Keep connection setup and shutdown separate from `Repositories`, which owns no pools
+- Reuse existing handles, never connect per request or disconnect from a repository operation
+- Give each pool one owner for setup and shutdown, borrowed handles never close it
+- Finish requests, jobs and queued database work before closing connections
+- Clean up resources opened before a startup failure, without a second disconnect fallback
+- Add a backend to the existing lifecycle owner and pass its handle only to repositories that need it
+- Preserve cancellation, timeout, error and transaction semantics across the native bridge
+
+## Changing ownership or reloading resources
+
+- Inventory initialization, reload, tests, background jobs and shutdown together
+- Move every writer when changing authoritative storage, copying a replaceable global into state is insufficient
+- A frozen container does not make its clients immutable
+- Decide whether each request captures one resource generation or reads it again during processing
+- Build and validate replacements before publishing them, preserving the old generation on failure
+- Serialize reloads and coordinate publication with shutdown
+- Drain streams and background users before retiring old resources, retaining every generation grows resource use
+- Define multi-worker coordination separately from an in-process swap
+- Choose mechanisms such as `ArcSwap` only after defining publication, capture and retirement
+
 ## Asked to move existing access behind a repository
 
 - Identify the calls being moved and their transaction boundary
@@ -36,17 +69,21 @@ description: Add database queries, extract existing Prisma calls into repositori
   - Keep typed inputs and reusable row projections in `src/`
   - Refresh SQLx query metadata after SQL changes
 - Keep Python conversion in the Python bridge, not in `litellm-rust/crates/db`
-- When caller integration is requested, expose the Rust operation through the bridge and call it from the repository
-  - A query-only port can land without switching callers
+- Expose typed operations through PyO3, never an API accepting SQL text, table names or column names from Python
+- Map dynamic filters and sort keys to closed Rust enums selecting static SQL, and bind every value
+- Switch all callers of the operation and delete its Prisma implementation in the same change
+- Move each transaction as a whole, including reads, writes and locks
 - Verify behavior against real PostgreSQL before switching callers
   - Check relevant null handling, conflicts, ordering, errors, and transaction semantics
-- Remove the old implementation only when no callers need it
+- Repository extraction can happen in the same PR, without converting unrelated routes to dependency injection
+- Read [incremental migration](references/python-fastapi.md#incremental-migration) for the caller, repository and wiring changes
 
 ## Migration context and boundaries
 
 - Long-term direction is to centralize Python database access in repositories, then replace their Prisma implementations incrementally with Rust APIs
 - Mixed Prisma and Rust implementations are expected during the transition, unrelated existing calls can stay as they are
-- Both backends may have separate pools through PgBouncer, size them together and verify pooling-mode compatibility
+- Prisma and SQLx keep separate client pools to PgBouncer, which pools connections to PostgreSQL
+  - Size both layers together and preserve reader/writer routing and pooling-mode compatibility
 - PgBouncer does not share transactions between Prisma and SQLx, keep each atomic unit of work on one backend and transaction
 - Never retry a failed Rust write through Prisma when it may already have committed
 - Prisma owns schema migrations today, eventual SQLx migration ownership and Prisma removal are separate planned steps, not part of a query port

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

@@ -0,0 +1,135 @@
+# Python and FastAPI
+
+## Existing pattern
+
+- `get_database` reads the lifespan-owned handle from request state
+- `get_repositories` builds the shared access point from that handle
+- A handler receives `Repositories | None`, then uses `repositories.banners`, 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
+
+## Dependency injection
+
+- Group operations by domain in `Repositories`, adding fields as callers migrate
+- The provider builds the container from existing handles
+
+```python
+@dataclass(frozen=True, slots=True)
+class Repositories:
+    banners: UserBannerRepository
+
+
+async def get_repositories(database: DatabaseDependency) -> Repositories | None:
+    return None if database is None else Repositories(banners=UserBannerRepository(database.prisma))
+
+
+RepositoriesDependency: TypeAlias = Annotated[Repositories | None, Depends(get_repositories)]
+```
+
+- FastAPI supplies the container to routes and caches the provider within each request
+- Ordinary helpers and jobs receive explicit arguments, never unresolved `Depends` defaults
+
+```python
+async def get_user_banner(repositories: RepositoriesDependency) -> UserBanner:
+    if repositories is None:
+        return UserBanner()
+    return await repositories.banners.get_banner()
+```
+
+## Lifespan state
+
+- Prisma startup owns connection setup, authentication refresh, health checks and migrations
+- `Database` wraps the initialized Prisma client, without opening another pool or holding fields per feature
+- FastAPI lifespan enters `litellm.repositories.database.lifespan` and publishes its handle in request state
+- Starlette shallow-copies lifespan state into request state, which is distinct from `app.state`
+- The database module never imports a repository, legacy globals borrow its handle without owning shutdown
+- Drain requests, jobs, logging, spend and audit work before database teardown
+- `litellm.repositories.database.lifespan` closes the handle on exit, including failures after entry
+- When adding native resources, initialize them in the serving process, respecting runtime fork restrictions
+- Await native shutdown instead of depending on Python garbage collection
+- Apply the [connection ownership rules](../SKILL.md#connection-ownership) when changing the owner
+
+## Incremental migration
+
+- Port one operation while other operations stay on Prisma, without waiting for unrelated routes to adopt dependency injection
+- Keep callers using `repositories.banners.get_banner()` with the same inputs and outputs
+- This example ports banner reads while writes remain Prisma
+  - `NativeBannerAdapter` and `database.native` are proposed interfaces, not implemented by the current banner PR
+  - The adapter converts typed PyO3 results to `UserBanner`, preserving defaults, invalid-data handling and storage errors
+  - SQL stays in `crates/db/queries/`, Python never sends SQL strings
+
+```diff
+ class UserBannerRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]):
+     table_name = "litellm_uisettings"
+
++    def __init__(self, prisma: PrismaClient, native: NativeBannerAdapter) -> None:
++        super().__init__(prisma)
++        self._native = native
++
+     async def get_banner(self) -> UserBanner:
+-        return parse_user_banner(await self.get_raw_settings())
++        return await self._native.get_banner()
+
+-    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
+```
+
+- Keep `save_banner` and its Prisma write unchanged
+- Preserve existing inheritance only while unmigrated methods need it, new repositories use composition
+- After the existing lifecycle owner initializes the native handle, update the provider
+
+```diff
+ async def get_repositories(database: DatabaseDependency) -> Repositories | None:
+-    return None if database is None else Repositories(banners=UserBannerRepository(database.prisma))
++    if database is None:
++        return None
++    return Repositories(
++        banners=UserBannerRepository(prisma=database.prisma, native=database.native),
++    )
+```
+
+- Both handles coexist for different operations, only the native `get_banner` implementation remains
+- Switch every caller, including helpers and jobs, and delete the old query in the same PR
+- Validate against real PostgreSQL and refresh SQLx metadata
+- Move a complete transaction together when the query belongs to one
+- Remove the Prisma handle and inherited table access after the repository's last Prisma operation moves
+- Roll back the release if needed, never retry an uncertain native write through Prisma
+
+## Timing and errors
+
+- FastAPI resolves dependencies before entering the handler
+- Handler exception handling cannot catch failures raised during dependency resolution
+- A universal required-database dependency can override authorization and validation behavior
+- Banner GET without a database returns the default disabled banner
+- Banner PATCH checks the admin role before returning the existing missing-database error
+- Invalid stored banner data falls back to the default, but storage failures still propagate
+- Banner writes finish before the success audit is queued as a FastAPI background task
+- Queue post-response work with `BackgroundTasks`, not bare `asyncio.create_task`, so the in-flight drain waits for it before the database closes
+- Do not move background work into an awaited provider merely to remove an import
+
+## Verification
+
+- Exercise real FastAPI handlers using `dependency_overrides` for external boundaries
+- Keep production parsing and persistence code real
+- Restore overrides after each test
+- Verify complete response bodies, status codes and OpenAPI against the base
+- Include combined failure conditions to pin error precedence
+- Check saved values, fresh revisions, audit contents and failed-write behavior
+- Check lifespan state with a running lifespan, not only a TestClient that skips startup
+- Verify providers do not connect, query, disconnect or audit eagerly
+- Use real PostgreSQL and live proxy requests when changing persistence behavior
+
+## Scope
+
+- Preserve direct Python call sites or migrate them explicitly
+- Inventory imports, state writers, side effects and error ordering before migrating a feature
+- Preserve configuration sync and serialization, and state which legacy reads remain
+- Do not migrate every route in a large module just to establish the pattern
+- Do not invent providers for resources without consumers
+
+## Framework references
+
+- [Dependency composition and caching](https://fastapi.tiangolo.com/tutorial/dependencies/sub-dependencies/)
+- [Dependency overrides](https://fastapi.tiangolo.com/advanced/testing-dependencies/)

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

@@ -0,0 +1,43 @@
+# Rust and Axum
+
+## Existing pattern
+
+- `crates/gateway` composes endpoint routers and owns server lifecycle
+- Endpoint crates inject state with `.with_state(...)`, such as `crates/gateway-inference` using `Arc<Gateway>`
+- Follow the owning crate's `AGENTS.md` before adding a database consumer
+
+## Database state scaffold
+
+- This is a proposed wiring shape, not an existing gateway database API
+- The runtime supplies an initialized pool
+
+```rust
+use axum::Router;
+use sqlx::PgPool;
+
+#[derive(Clone)]
+struct DatabaseState {
+    pool: PgPool,
+}
+
+fn with_database(routes: Router<DatabaseState>, pool: PgPool) -> Router {
+    routes.with_state(DatabaseState { pool })
+}
+```
+
+- Handlers extract `State<DatabaseState>` and call typed `litellm_db` operations with the pool or transaction
+- Extend existing endpoint state or use `FromRef` when other dependencies are already present
+- Cloning `PgPool` shares its handle, it does not create another pool
+- Keep SQL in `crates/db`, and response conversion and HTTP error mapping in the endpoint crate
+- Reuse the same query APIs as Python's bridge, without creating a second SQL implementation
+
+## Lifetime and verification
+
+- Apply the [connection ownership rules](../SKILL.md#connection-ownership), finish active requests and database jobs before closing the pool
+- Preserve authentication, optional-database behavior and transaction boundaries for the specific route
+- Verify HTTP behavior through the owning router and query behavior against real PostgreSQL
+
+## Framework references
+
+- [Axum state and substates](https://docs.rs/axum/latest/axum/extract/struct.State.html)
+- [SQLx pool ownership](https://docs.rs/sqlx/latest/sqlx/struct.Pool.html)

AGENTS.md

@@ -1,3 +1,5 @@
+For database operations, connection ownership and proxy dependency wiring, follow [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/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/deps.py

@@ -0,0 +1,31 @@
+from typing import Annotated, Final, TypeAlias
+
+from fastapi import Depends, Request
+
+from litellm.repositories.container import Repositories
+from litellm.repositories.database import Database
+from litellm.repositories.user_banner_repository import UserBannerRepository
+
+_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)]
+
+
+async def get_repositories(database: DatabaseDependency) -> Repositories | None:
+    return None if database is None else Repositories(banners=UserBannerRepository(database.prisma))
+
+
+RepositoriesDependency: TypeAlias = Annotated[Repositories | None, Depends(get_repositories)]

litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py

@@ -0,0 +1,31 @@
+from collections.abc import Coroutine
+from typing import Annotated, Protocol, TypeAlias
+
+from fastapi import Depends
+
+from litellm.models.user_banner import UserBanner
+from litellm.proxy._types import LitellmTableNames, UserAPIKeyAuth
+from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID
+
+
+class UserBannerAudit(Protocol):
+    def __call__(self, before: UserBanner, after: UserBanner, user: UserAPIKeyAuth) -> Coroutine[None, None, None]: ...
+
+
+async def get_user_banner_audit() -> UserBannerAudit:
+    from litellm.proxy.proxy_server import create_config_audit_log
+
+    async def audit(before: UserBanner, after: UserBanner, user: UserAPIKeyAuth) -> None:
+        await create_config_audit_log(
+            param_name=USER_BANNER_ROW_ID,
+            action="updated",
+            before_value=before.model_dump(),
+            after_value=after.model_dump(),
+            user_api_key_dict=user,
+            table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME,
+        )
+
+    return audit
+
+
+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.deps import RepositoriesDependency
+from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import (
+    UserBannerAuditDependency,
+)
 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(repositories: RepositoriesDependency) -> 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 repositories is None:
         return UserBanner()
 
-    raw_settings: Final = await UserBannerRepository(prisma_client).get_raw_settings()
-    return parse_user_banner(raw_settings)
+    return await repositories.banners.get_banner()
 
 
 @router.patch(
@@ -93,21 +48,21 @@ async def get_user_banner() -> UserBanner:
 async def update_user_banner(
     banner_update: UserBannerUpdate,
     user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
+    repositories: RepositoriesDependency,
+    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 repositories 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 repositories.banners.get_banner()
     banner: Final = UserBanner(
         enabled=banner_update.enabled,
         message=banner_update.message,
@@ -115,17 +70,8 @@ async def update_user_banner(
         revision=uuid4().hex,
     )
 
-    await repository.upsert_settings(json.dumps(banner.model_dump()))
+    await repositories.banners.save_banner(banner)
 
-    asyncio.create_task(
-        create_config_audit_log(
-            param_name=USER_BANNER_ROW_ID,
-            action="updated",
-            before_value=before.model_dump(),
-            after_value=banner.model_dump(),
-            user_api_key_dict=user_api_key_dict,
-            table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME,
-        )
-    )
+    background_tasks.add_task(audit, before, banner, user_api_key_dict)
 
     return UpdateUserBannerResponse(message="User banner updated successfully", banner=banner)

litellm/repositories/__init__.py

@@ -5,6 +5,7 @@ Repository classes for database operations.
 from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository
 from litellm.repositories.budget_repository import BudgetRepository
 from litellm.repositories.config_repository import ConfigRepository
+from litellm.repositories.container import Repositories
 from litellm.repositories.credentials_repository import CredentialsRepository
 from litellm.repositories.managed_file_repository import ManagedFileRepository
 from litellm.repositories.model_repository import ModelRepository
@@ -141,6 +142,7 @@ __all__ = [
     "ProjectRepository",
     "PromptRepository",
     "ReadOnlyTable",
+    "Repositories",
     "SSOConfigRepository",
     "SearchToolsRepository",
     "SkillsRepository",

litellm/repositories/container.py

@@ -0,0 +1,8 @@
+from dataclasses import dataclass
+
+from litellm.repositories.user_banner_repository import UserBannerRepository
+
+
+@dataclass(frozen=True, slots=True)
+class Repositories:
+    banners: UserBannerRepository

litellm/repositories/user_banner_repository.py

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

tests/unit/proxy/test_deps.py

@@ -0,0 +1,75 @@
+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.models.user_banner import UserBanner
+from litellm.proxy.deps import DatabaseStateMissingError, RepositoriesDependency, get_database
+from litellm.repositories import database as database_runtime
+
+
+def _app(database: database_runtime.Database | None) -> FastAPI:
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        async with database_runtime.lifespan(database) as owned_database:
+            yield {"database": owned_database}
+
+    app: Final = FastAPI(lifespan=lifespan)
+
+    @app.get("/banner")
+    async def banner(repositories: RepositoriesDependency) -> UserBanner:
+        return UserBanner() if repositories is None else await repositories.banners.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(database_runtime.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(database_runtime.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/ui_crud_endpoints/test_user_banner_endpoints.py

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