diff --git a/.agents/skills/proxy-dependencies/SKILL.md b/.agents/skills/proxy-dependencies/SKILL.md new file mode 100644 index 00000000000..7ff6c54aa86 --- /dev/null +++ b/.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 diff --git a/.agents/skills/proxy-dependencies/references/route-migration.md b/.agents/skills/proxy-dependencies/references/route-migration.md new file mode 100644 index 00000000000..c5e2c5da878 --- /dev/null +++ b/.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/) diff --git a/.agents/skills/proxy-dependencies/references/runtime-ownership.md b/.agents/skills/proxy-dependencies/references/runtime-ownership.md new file mode 100644 index 00000000000..e2d218b5de5 --- /dev/null +++ b/.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) diff --git a/AGENTS.md b/AGENTS.md index f4b1de57807..24aeb7b615e 100644 --- a/AGENTS.md +++ b/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: diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 63137949f38..42916002a9b 100644 --- a/litellm-rust/Cargo.lock +++ b/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", diff --git a/litellm-rust/crates/db/.sqlx/query-136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae.json b/litellm-rust/crates/db/.sqlx/query-136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae.json new file mode 100644 index 00000000000..8b49a9c5e82 --- /dev/null +++ b/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" +} diff --git a/litellm-rust/crates/db/.sqlx/query-c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc.json b/litellm-rust/crates/db/.sqlx/query-c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc.json new file mode 100644 index 00000000000..88f62dd8e06 --- /dev/null +++ b/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" +} diff --git a/litellm-rust/crates/db/queries/user_banner/find_user_banner.sql b/litellm-rust/crates/db/queries/user_banner/find_user_banner.sql new file mode 100644 index 00000000000..6833855a92a --- /dev/null +++ b/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' diff --git a/litellm-rust/crates/db/queries/user_banner/save_user_banner.sql b/litellm-rust/crates/db/queries/user_banner/save_user_banner.sql new file mode 100644 index 00000000000..193c08d409c --- /dev/null +++ b/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 diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index 541f97c7db9..501b890e568 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/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}; diff --git a/litellm-rust/crates/db/src/user_banner.rs b/litellm-rust/crates/db/src/user_banner.rs new file mode 100644 index 00000000000..a1f0207fdfd --- /dev/null +++ b/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, 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>, + ) -> Result<(), Box> { + 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, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + 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>, + ) -> Result<(), Box> { + 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(()) + } +} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 8afde8362f6..8bd85dd4034 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/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 diff --git a/litellm-rust/crates/python-bridge/src/database.rs b/litellm-rust/crates/python-bridge/src/database.rs new file mode 100644 index 00000000000..251e4270ebe --- /dev/null +++ b/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 = + 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> { + 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> { + 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> { + 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> { + let pool = self.pool.clone(); + crate::execution::run_async( + py, + async move { + pool.close().await; + Ok::<(), Error>(()) + }, + map_error, + ) + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 12827b76ff3..fc7b5545545 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/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", diff --git a/litellm/_db.py b/litellm/_db.py new file mode 100644 index 00000000000..c064cbf2fdb --- /dev/null +++ b/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() diff --git a/litellm/constants.py b/litellm/constants.py index 2901b7d8948..c870f2c8ec0 100644 --- a/litellm/constants.py +++ b/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 diff --git a/litellm/models/user_banner.py b/litellm/models/user_banner.py new file mode 100644 index 00000000000..5470ad4dafa --- /dev/null +++ b/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." + ), + ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 49d23d82aa9..83af47feb6a 100644 --- a/litellm/proxy/_types.py +++ b/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): diff --git a/litellm/proxy/deps.py b/litellm/proxy/deps.py new file mode 100644 index 00000000000..a10ddfedd69 --- /dev/null +++ b/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)] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bcba76ca1b3..c7ed0f244ea 100644 --- a/litellm/proxy/proxy_server.py +++ b/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() diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py new file mode 100644 index 00000000000..895ee54cba9 --- /dev/null +++ b/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)] diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py b/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py index 56339202d46..ac2c3be81e9 100644 --- a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py +++ b/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) diff --git a/litellm/repositories/user_banner_repository.py b/litellm/repositories/user_banner_repository.py index 4e113dff988..ffdaecdf14c 100644 --- a/litellm/repositories/user_banner_repository.py +++ b/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 diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 7364a1d72f2..f073620a310 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/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", diff --git a/litellm/rust_bridge/database.py b/litellm/rust_bridge/database.py new file mode 100644 index 00000000000..91c1f7d64bc --- /dev/null +++ b/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) diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 11a44a0ec16..707bd8ca237 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/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(") diff --git a/tests/unit/proxy/test_deps.py b/tests/unit/proxy/test_deps.py new file mode 100644 index 00000000000..0591064cfe2 --- /dev/null +++ b/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) diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index a26a64c6fb0..cb7e12720b5 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/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"}, diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 745488bf93f..a9b335c597c 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/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 diff --git a/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py index c014ffd04c6..a7e01252539 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py +++ b/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 diff --git a/tests/unit/test__db.py b/tests/unit/test__db.py new file mode 100644 index 00000000000..35b1ec9623a --- /dev/null +++ b/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()