From b50f283c04c24adf35914e1ebe2ff79f449d7f27 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:50:59 +0000 Subject: [PATCH] refactor: add lifespan-owned database dependency foundation --- AGENTS.md | 2 + .../skills/proxy-dependencies/SKILL.md | 39 ++++++++ .../references/route-migration.md | 44 +++++++++ .../references/runtime-ownership.md | 48 ++++++++++ litellm-rust/AGENTS.md | 2 + litellm/_db.py | 43 +++++++++ litellm/models/user_banner.py | 42 ++++++++ litellm/proxy/_types.py | 2 + litellm/proxy/deps.py | 18 ++++ litellm/proxy/proxy_server.py | 19 ++-- .../user_banner_dependencies.py | 37 ++++++++ .../user_banner_endpoints.py | 92 ++++-------------- .../repositories/user_banner_repository.py | 28 +++++- .../unit/proxy/proxy_server/test_lifecycle.py | 22 +++-- tests/unit/proxy/test_deps.py | 95 +++++++++++++++++++ .../test_user_banner_endpoints.py | 61 +++++++++++- tests/unit/test__db.py | 36 +++++++ 17 files changed, 538 insertions(+), 92 deletions(-) create mode 100644 litellm-rust/.agents/skills/proxy-dependencies/SKILL.md create mode 100644 litellm-rust/.agents/skills/proxy-dependencies/references/route-migration.md create mode 100644 litellm-rust/.agents/skills/proxy-dependencies/references/runtime-ownership.md create mode 100644 litellm/_db.py create mode 100644 litellm/models/user_banner.py create mode 100644 litellm/proxy/deps.py create mode 100644 litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py create mode 100644 tests/unit/proxy/test_deps.py create mode 100644 tests/unit/test__db.py diff --git a/AGENTS.md b/AGENTS.md index f4b1de57807..69e34be00fc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,3 +1,5 @@ +For proxy dependency wiring and database lifetime changes, follow [proxy-dependencies](litellm-rust/.agents/skills/proxy-dependencies/SKILL.md) alongside [database-access](.agents/skills/database-access/SKILL.md) + Do not write comments unless they are any of: - absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear) - used as an input for tools to read and act on. For example: diff --git a/litellm-rust/.agents/skills/proxy-dependencies/SKILL.md b/litellm-rust/.agents/skills/proxy-dependencies/SKILL.md new file mode 100644 index 00000000000..4d714adc48c --- /dev/null +++ b/litellm-rust/.agents/skills/proxy-dependencies/SKILL.md @@ -0,0 +1,39 @@ +--- +name: proxy-dependencies +description: Inject typed repositories and services into Python proxy handlers, or change database runtime ownership while preserving behavior for future Rust backends +--- + +# Proxy dependencies + +- Follow the root [database-access skill](../../../../.agents/skills/database-access/SKILL.md) for queries, repository extraction and backend migration +- This skill owns dependency wiring and resource lifetime, not SQL or schema ownership + +## Handler contracts + +- Inject named FastAPI providers through `Annotated[..., Depends(...)]` +- Expose narrow repository protocols and typed service callables, not Prisma clients, models, tables or query dictionaries +- Keep database I/O out of dependency construction +- Keep no-database behavior and authorization precedence specific to the route +- Ordinary helpers and background jobs receive explicit arguments, never unresolved `Depends` defaults +- Override named providers in route tests and keep production parsing and persistence code real + +## Database ownership + +- `litellm._db.Database` exposes typed repositories and an async close operation +- `from_prisma` adopts an already initialized client without opening a second pool +- FastAPI lifespan enters `_db.lifespan` and publishes its handle in request state +- Drain database consumers before exiting that context, which alone closes the handle +- Do not keep a parallel Prisma disconnect fallback in proxy shutdown +- Legacy globals remain compatibility access for unmigrated code, not a second pool or owner +- Never close a database from a request dependency or repository operation +- SQLx integration is a separate explicit change, preserve transaction boundaries and do not replay failed writes through Prisma + +## Migration process + +- Inventory imports, direct callers, state writers, side effects and error ordering first +- Migrate one complete feature and its helper calls before expanding +- Preserve response models, OpenAPI, audit scheduling, configuration sync and serialization +- Verify default behavior, authorization, invalid input, storage failure and shutdown ordering +- State which legacy reads and backend leaks remain +- Read [route migration](references/route-migration.md) for the banner example +- Read [runtime ownership](references/runtime-ownership.md) before changing resource lifetime or reload behavior diff --git a/litellm-rust/.agents/skills/proxy-dependencies/references/route-migration.md b/litellm-rust/.agents/skills/proxy-dependencies/references/route-migration.md new file mode 100644 index 00000000000..d3c2aad4204 --- /dev/null +++ b/litellm-rust/.agents/skills/proxy-dependencies/references/route-migration.md @@ -0,0 +1,44 @@ +# Route migration + +## Existing pattern + +- `get_database` reads the lifespan-owned handle from request state +- A feature provider selects the narrow repository, such as `database.user_banners` +- A handler receives `UserBannerStore | None`, never a Prisma client +- `UserBannerRepository` keeps the existing Prisma table, serialization and configuration-sync behavior +- The banner audit provider binds the legacy audit callable at request time, without performing audit I/O +- SQLx can later implement the same typed operations without changing handler signatures + +## Timing and errors + +- FastAPI resolves dependencies before entering the handler +- Handler exception handling cannot catch failures raised during dependency resolution +- A universal required-database dependency can override authorization and validation behavior +- Banner GET without a database returns the default disabled banner +- Banner PATCH checks the admin role before returning the existing missing-database error +- Invalid stored banner data falls back to the default, but storage failures still propagate +- Banner writes finish before the success audit is scheduled +- Do not move background work into an awaited provider merely to remove an import + +## Verification + +- Exercise real FastAPI handlers using `dependency_overrides` for external boundaries +- Restore overrides after each test +- Verify complete response bodies, status codes and OpenAPI against the base +- Include combined failure conditions to pin error precedence +- Check saved values, fresh revisions, audit contents and failed-write behavior +- Check lifespan state with a running lifespan, not only a TestClient that skips startup +- Verify providers do not connect, query, disconnect or audit eagerly +- Use real PostgreSQL and live proxy requests when changing persistence behavior + +## Scope + +- Preserve direct Python call sites or migrate them explicitly +- Keep repository extraction, backend switching and error normalization separate +- Do not migrate every route in a large module just to establish the pattern +- Do not invent providers for resources without consumers + +## Framework references + +- [Dependency composition and caching](https://fastapi.tiangolo.com/tutorial/dependencies/sub-dependencies/) +- [Dependency overrides](https://fastapi.tiangolo.com/advanced/testing-dependencies/) diff --git a/litellm-rust/.agents/skills/proxy-dependencies/references/runtime-ownership.md b/litellm-rust/.agents/skills/proxy-dependencies/references/runtime-ownership.md new file mode 100644 index 00000000000..de867addc36 --- /dev/null +++ b/litellm-rust/.agents/skills/proxy-dependencies/references/runtime-ownership.md @@ -0,0 +1,48 @@ +# Runtime ownership + +## Current foundation + +- Prisma startup still owns connection setup, authentication refresh, health checks and migrations +- Lifespan enters `_db.lifespan` with the initialized client adapted by `from_prisma` +- The `Database` handle provides typed repositories and `aclose`, without exposing Prisma to handlers +- The existing shutdown sequence drains requests, jobs, logging, spend and audit work before disconnect +- `_db.lifespan` alone closes the handle, including startup or shutdown failure +- No parallel disconnect fallback belongs in proxy shutdown +- This does not remove all proxy globals or implement a Rust database backend + +## Adding a backend + +- Keep domain operations in repositories and conversion at the native bridge +- Define who creates, borrows and closes each pool, including startup-failure cleanup +- Initialize native resources in the serving process, respecting the existing native runtime's fork restrictions +- Await native shutdown instead of depending on Python garbage collection +- Do not create a pool per request or close a shared pool from a borrowed handle +- Preserve error, cancellation, timeout and transaction semantics across the bridge +- Keep each atomic operation on one backend and transaction +- Follow the root database-access skill for backend selection and schema ownership + +## Changing ownership + +- Inventory initialization, reload, tests, background jobs and shutdown together +- Migrate every writer of a resource when moving its authoritative storage +- Starlette lifespan state is shallow-copied into request state +- `app.state` and yielded lifespan state are distinct mechanisms +- A frozen container does not make its database client, router or cache immutable +- Do not copy a replaceable global into state while leaving reload to update only the global + +## Runtime swapping + +- Backend selection at startup does not require hot swapping +- Define whether a request captures one runtime generation or reads again during processing +- Build and validate a replacement before publishing it, leaving the old generation usable on failure +- Serialize reloads and coordinate publication with shutdown +- Drain streams and background users before retiring old resources +- Retaining every old generation until shutdown is simple but can grow resource use without bounds +- Define multi-worker coordination separately from an in-process swap +- Axum `State` and `FromRef` can expose the same narrow dependencies to Rust handlers +- Do not prescribe `ArcSwap` before defining publication, capture and retirement semantics + +## Framework references + +- [Starlette lifespan](https://www.starlette.io/lifespan/) +- [Axum state](https://docs.rs/axum/latest/axum/extract/struct.State.html) diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index e62a9463cb5..c21837858e4 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,5 +1,7 @@ # Rust workspace rules +For proxy dependency wiring and database lifetime changes, follow [.agents/skills/proxy-dependencies/SKILL.md](.agents/skills/proxy-dependencies/SKILL.md) alongside the root database-access skill + For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.agents/skills/rust-tracing/SKILL.md) For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md) diff --git a/litellm/_db.py b/litellm/_db.py new file mode 100644 index 00000000000..587cbca54a5 --- /dev/null +++ b/litellm/_db.py @@ -0,0 +1,43 @@ +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +from litellm.repositories.user_banner_repository import UserBannerRepository, UserBannerStore + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +@runtime_checkable +class Database(Protocol): + @property + def user_banners(self) -> UserBannerStore: ... + + async def aclose(self) -> None: ... + + +@dataclass(frozen=True, slots=True) +class _PrismaDatabase: + client: "PrismaClient" + + @property + def user_banners(self) -> UserBannerStore: + return UserBannerRepository(self.client) + + async def aclose(self) -> None: + await self.client.disconnect() + + +def from_prisma(client: "PrismaClient") -> Database: + """Adapt an initialized client without opening another pool.""" + return _PrismaDatabase(client) + + +@asynccontextmanager +async def lifespan(database: Database | None) -> AsyncGenerator[Database | None, None]: + try: + yield database + finally: + if database is not None: + await database.aclose() 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 6280e2caf87..d238ff0ba91 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..4ee0aa15c90 --- /dev/null +++ b/litellm/proxy/deps.py @@ -0,0 +1,18 @@ +from typing import Annotated, Final, TypeAlias + +from fastapi import Depends, Request +from pydantic import ConfigDict, TypeAdapter + +from litellm._db import Database + +_DATABASE_ADAPTER: Final[TypeAdapter[Database | None]] = TypeAdapter( + Database | None, config=ConfigDict(arbitrary_types_allowed=True) +) +_MISSING_DATABASE: Final = object() + + +async def get_database(request: Request) -> Database | None: + return _DATABASE_ADAPTER.validate_python(getattr(request.state, "database", _MISSING_DATABASE)) + + +DatabaseDependency: TypeAlias = Annotated[Database | None, Depends(get_database)] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 129cd7e1589..d2667f5b052 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, @@ -1288,13 +1289,7 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N if worker_heartbeat is not None and prisma_client: await worker_heartbeat.deregister() if prisma_client: - # Drain the SGR fold first: it lives in memory, so an un-drained interval - # is lost, and a write attempted after disconnect raises - # ClientNotConnectedError rather than persisting anything. Ordering this - # inside the same guard is what keeps the two from drifting apart. await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client)) - verbose_proxy_logger.debug("Disconnecting from Prisma") - await prisma_client.disconnect() else: await _drain_audit_tasks_on_shutdown() @@ -1427,6 +1422,12 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries: @asynccontextmanager async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: + async with AsyncExitStack() as resources, _proxy_lifespan(app, resources) as state: + yield state + + +@asynccontextmanager +async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGenerator[ProxyLifespanState, None]: global \ prisma_client, \ master_key, \ @@ -1550,6 +1551,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState ) ) + database: Final = await resources.enter_async_context( + _db.lifespan(None if prisma_client is None else _db.from_prisma(prisma_client)) + ) + await migrate_if_requested( environ=os.environ, master_key=master_key, @@ -1784,7 +1789,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState enabled=tracing_enabled, settings=tracing_settings, ) as receiver: - state: Final[ProxyLifespanState] = {"tracing_receiver": receiver} + state: Final[ProxyLifespanState] = {"tracing_receiver": receiver, "database": database} from litellm.proxy.admin_mcp import admin_mcp_lifespan telemetry_settings: Final = load_telemetry_settings() 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..b4016d4a226 --- /dev/null +++ b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py @@ -0,0 +1,37 @@ +from collections.abc import Coroutine +from typing import Annotated, Protocol, TypeAlias + +from fastapi import Depends + +from litellm.models.user_banner import UserBanner +from litellm.proxy._types import LitellmTableNames, UserAPIKeyAuth +from litellm.proxy.deps import DatabaseDependency +from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerStore + + +class UserBannerAudit(Protocol): + def __call__(self, before: UserBanner, after: UserBanner, user: UserAPIKeyAuth) -> Coroutine[None, None, None]: ... + + +async def get_user_banner_repository(database: DatabaseDependency) -> UserBannerStore | None: + return None if database is None else database.user_banners + + +async def get_user_banner_audit() -> UserBannerAudit: + from litellm.proxy.proxy_server import create_config_audit_log + + async def audit(before: UserBanner, after: UserBanner, user: UserAPIKeyAuth) -> None: + await create_config_audit_log( + param_name=USER_BANNER_ROW_ID, + action="updated", + before_value=before.model_dump(), + after_value=after.model_dump(), + user_api_key_dict=user, + table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME, + ) + + return audit + + +UserBannerRepositoryDependency: TypeAlias = Annotated[UserBannerStore | None, Depends(get_user_banner_repository)] +UserBannerAuditDependency: TypeAlias = Annotated[UserBannerAudit, Depends(get_user_banner_audit)] diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py b/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py index 56339202d46..1ba6cfed3f0 100644 --- a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py @@ -1,88 +1,44 @@ import asyncio -import json -from typing import Annotated, Final, Literal +from typing import Annotated, Final from fastapi import APIRouter, Depends, HTTPException -from pydantic import Field, ValidationError, model_validator +from pydantic import ConfigDict from litellm._uuid import uuid4 -from litellm.proxy._types import LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth +from litellm.models.user_banner import UserBanner, UserBannerUpdate +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerRepository +from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import ( + UserBannerAuditDependency, + UserBannerRepositoryDependency, +) from litellm.types.llms.base import LiteLLMBaseModel router: Final = APIRouter() -USER_BANNER_MAX_MESSAGE_LENGTH: Final = 4000 - -UserBannerSeverity = Literal["info", "warning", "error"] - - -class UserBannerUpdate(LiteLLMBaseModel): - enabled: bool = Field( - default=False, - description="If true, the banner is shown to all authenticated dashboard users.", - ) - message: str = Field( - default="", - max_length=USER_BANNER_MAX_MESSAGE_LENGTH, - description="Banner text shown to dashboard users. Markdown is supported.", - ) - severity: UserBannerSeverity = Field( - default="info", - description="Visual style of the banner.", - ) - - @model_validator(mode="after") - def _require_message_when_enabled(self) -> "UserBannerUpdate": - if self.enabled and not self.message.strip(): - raise ValueError("message must be non-empty when the banner is enabled") - return self - - -class UserBanner(UserBannerUpdate): - revision: str = Field( - default="", - description=( - "Server-stamped opaque publish identity; a fresh value is generated on every " - "update so clients re-surface dismissed banners on republish." - ), - ) - class UpdateUserBannerResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + message: str banner: UserBanner -def parse_user_banner(raw_settings: object) -> UserBanner: - if raw_settings is None: - return UserBanner() - try: - parsed: Final = json.loads(raw_settings) if isinstance(raw_settings, str) else raw_settings - return UserBanner.model_validate(parsed) - except (json.JSONDecodeError, ValidationError): - return UserBanner() - - @router.get( "/get/user_banner", tags=["UI Settings"], dependencies=[Depends(user_api_key_auth)], response_model=UserBanner, ) -async def get_user_banner() -> UserBanner: +async def get_user_banner(repository: UserBannerRepositoryDependency) -> UserBanner: """ Get the admin-published dashboard banner. Readable by any authenticated user; rendered on every dashboard page. """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: + if repository is None: return UserBanner() - raw_settings: Final = await UserBannerRepository(prisma_client).get_raw_settings() - return parse_user_banner(raw_settings) + return await repository.get_banner() @router.patch( @@ -93,21 +49,20 @@ async def get_user_banner() -> UserBanner: async def update_user_banner( banner_update: UserBannerUpdate, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + repository: UserBannerRepositoryDependency, + audit: UserBannerAuditDependency, ) -> UpdateUserBannerResponse: """ Publish, edit, or unpublish the dashboard banner. Only proxy admins are allowed to modify it. """ - from litellm.proxy.proxy_server import create_config_audit_log, prisma_client - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException(status_code=403, detail="Only proxy admins can update the user banner.") - if prisma_client is None: + if repository is None: raise HTTPException(status_code=500, detail="Database not connected. Please connect a database.") - repository: Final = UserBannerRepository(prisma_client) - before: Final = parse_user_banner(await repository.get_raw_settings()) + before: Final = await repository.get_banner() banner: Final = UserBanner( enabled=banner_update.enabled, message=banner_update.message, @@ -115,17 +70,8 @@ async def update_user_banner( revision=uuid4().hex, ) - await repository.upsert_settings(json.dumps(banner.model_dump())) + await repository.save_banner(banner) - asyncio.create_task( - create_config_audit_log( - param_name=USER_BANNER_ROW_ID, - action="updated", - before_value=before.model_dump(), - after_value=banner.model_dump(), - user_api_key_dict=user_api_key_dict, - table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME, - ) - ) + asyncio.create_task(audit(before, banner, user_api_key_dict)) return UpdateUserBannerResponse(message="User banner updated successfully", banner=banner) diff --git a/litellm/repositories/user_banner_repository.py b/litellm/repositories/user_banner_repository.py index 4e113dff988..e8298f4aa0d 100644 --- a/litellm/repositories/user_banner_repository.py +++ b/litellm/repositories/user_banner_repository.py @@ -1,5 +1,9 @@ -from typing import TYPE_CHECKING, Final +import json +from typing import TYPE_CHECKING, Final, Protocol +from pydantic import ValidationError + +from litellm.models.user_banner import UserBanner from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: @@ -8,9 +12,31 @@ if TYPE_CHECKING: USER_BANNER_ROW_ID: Final = "user_banner" +class UserBannerStore(Protocol): + async def get_banner(self) -> UserBanner: ... + + async def save_banner(self, banner: UserBanner) -> None: ... + + +def parse_user_banner(raw_settings: object) -> UserBanner: + if raw_settings is None: + return UserBanner() + try: + parsed: Final = json.loads(raw_settings) if isinstance(raw_settings, str) else raw_settings + return UserBanner.model_validate(parsed) + except (json.JSONDecodeError, ValidationError): + return UserBanner() + + class UserBannerRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]): table_name = "litellm_uisettings" + async def get_banner(self) -> UserBanner: + return parse_user_banner(await self.get_raw_settings()) + + async def save_banner(self, banner: UserBanner) -> None: + await self.upsert_settings(json.dumps(banner.model_dump())) + async def get_raw_settings(self) -> object: db_record: Final = await self.table.find_unique(where={"id": USER_BANNER_ROW_ID}) return db_record.ui_settings if db_record is not None else None diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 11a44a0ec16..b182d3ed33f 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.from_prisma(fake_prisma)): + await proxy_shutdown_event() + fake_prisma.disconnect.assert_not_awaited() observed = { "disconnect_called": fake_prisma.disconnect.await_count == 1, @@ -222,7 +225,8 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - await proxy_shutdown_event() + async with _db.lifespan(_db.from_prisma(fake_prisma)): + await proxy_shutdown_event() assert calls == ["flush", "disconnect"] @@ -262,7 +266,8 @@ async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_befor monkeypatch.setattr(litellm, "success_callback", [], raising=False) track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts())) - await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) + async with _db.lifespan(_db.from_prisma(fake_prisma)): + await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) assert calls == ["flush", "audit", "disconnect"] @@ -318,7 +323,7 @@ async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_t @pytest.mark.asyncio -async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): +async def test_database_lifespan_propagates_disconnect_error_after_shutdown(monkeypatch): fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone")) monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) @@ -333,7 +338,8 @@ async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): monkeypatch.setattr(litellm, "success_callback", [], raising=False) with pytest.raises(RuntimeError, match="db gone"): - await proxy_shutdown_event() + async with _db.lifespan(_db.from_prisma(fake_prisma)): + await proxy_shutdown_event() # --------------------------------------------------------------------------- @@ -1060,7 +1066,7 @@ def test_otel_global_provider_published_after_callback_init(): span was orphaned. The publish (``publish_global_otel_v2_provider``) must therefore appear after ``_initialize_startup_logging`` in the lifespan source. """ - wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event) + wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan) source = inspect.getsource(wrapped) init_pos = source.find("ProxyStartupEvent.initialize_startup_logging(") publish_pos = source.find("publish_global_otel_v2_provider(") @@ -1146,7 +1152,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): The call must follow Prisma setup so DB-backed deployments do not false-positive. Direct ``_warn_budget_without_db`` tests cover the warning behavior itself. """ - wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event) + wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan) source = inspect.getsource(wrapped) budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:") warn_pos = source.find("warn_budget_without_db(") diff --git a/tests/unit/proxy/test_deps.py b/tests/unit/proxy/test_deps.py new file mode 100644 index 00000000000..712b206c8bb --- /dev/null +++ b/tests/unit/proxy/test_deps.py @@ -0,0 +1,95 @@ +from contextlib import asynccontextmanager +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from pydantic import ValidationError + +from litellm import _db +from litellm.models.user_banner import UserBanner +from litellm.proxy.deps import get_database +from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import ( + UserBannerRepositoryDependency, + get_user_banner_repository, +) + + +def _app(database: _db.Database | None) -> FastAPI: + @asynccontextmanager + async def lifespan(app: FastAPI): + async with _db.lifespan(database) as owned_database: + yield {"database": owned_database} + + app: Final = FastAPI(lifespan=lifespan) + + @app.get("/banner") + async def banner(repository: UserBannerRepositoryDependency) -> UserBanner: + return UserBanner() if repository is None else await repository.get_banner() + + return app + + +def test_lifespan_database_is_shared_until_shutdown(): + prisma: Final = MagicMock() + prisma.disconnect = AsyncMock() + prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None) + app: Final = _app(_db.from_prisma(prisma)) + with TestClient(app) as client: + assert client.get("/banner").json() == UserBanner().model_dump() + prisma.db.litellm_uisettings.find_unique.return_value = SimpleNamespace( + ui_settings={"enabled": True, "message": "published", "revision": "revision"} + ) + assert ( + client.get("/banner").json() + == UserBanner(enabled=True, message="published", revision="revision").model_dump() + ) + prisma.disconnect.assert_not_awaited() + prisma.disconnect.assert_awaited_once() + + +def test_no_database_state_preserves_optional_repository(): + with TestClient(_app(None)) as client: + response: Final = client.get("/banner") + assert response.status_code == 200 + assert response.json() == UserBanner().model_dump() + + +def test_dependency_override_changes_next_request_without_changing_lifespan_owner(monkeypatch): + prisma: Final = MagicMock() + prisma.disconnect = AsyncMock() + prisma.db.litellm_uisettings.find_unique = AsyncMock( + return_value=SimpleNamespace(ui_settings={"message": "stored"}) + ) + app: Final = _app(_db.from_prisma(prisma)) + with TestClient(app) as client: + assert client.get("/banner").json() == UserBanner(message="stored").model_dump() + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) + assert client.get("/banner").json() == UserBanner().model_dump() + prisma.disconnect.assert_not_awaited() + prisma.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_missing_lifespan_is_not_treated_as_a_disabled_database(): + from starlette.requests import Request + + request: Final = Request({"type": "http", "state": {}}) + with pytest.raises(ValidationError, match="Database"): + await get_database(request) + + +@pytest.mark.asyncio +async def test_constructing_repository_does_not_touch_database(): + prisma: Final = MagicMock() + repository: Final = await get_user_banner_repository(_db.from_prisma(prisma)) + assert repository is not None + assert prisma.mock_calls == [] + + +def test_dependencies_do_not_become_openapi_parameters(): + operation: Final = _app(None).openapi()["paths"]["/banner"]["get"] + assert "parameters" not in operation + assert "requestBody" not in operation 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..7fdf3222869 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.from_prisma(mock_prisma)) return mock_prisma @@ -94,13 +98,64 @@ class TestGetUserBanner: assert response.json() == DISABLED_BANNER def test_no_database_returns_disabled_banner(self, admin_auth, monkeypatch): - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) response = client.get("/get/user_banner") assert response.status_code == 200 assert response.json() == DISABLED_BANNER class TestUpdateUserBanner: + def test_storage_finishes_before_audit_runs(self, admin_auth, monkeypatch, mock_audit_log): + events = [] + prisma = _mock_prisma(monkeypatch) + + async def save(**kwargs): + await asyncio.sleep(0) + events.append("saved") + + async def audit(before, after, user): + events.append("audited") + + prisma.db.litellm_uisettings.upsert.side_effect = save + mock_audit_log.side_effect = audit + response = client.patch("/update/user_banner", json=PUBLISH_BODY) + assert response.status_code == 200 + assert events == ["saved", "audited"] + + @pytest.mark.parametrize( + ("role", "status", "detail"), + [ + (LitellmUserRoles.INTERNAL_USER, 403, "Only proxy admins can update the user banner."), + (LitellmUserRoles.PROXY_ADMIN, 500, "Database not connected. Please connect a database."), + ], + ) + def test_no_database_preserves_authorization_precedence(self, monkeypatch, role, status, detail): + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override(role)) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) + response = client.patch("/update/user_banner", json=PUBLISH_BODY) + assert response.status_code == status + assert response.json() == {"detail": detail} + + def test_failed_write_does_not_schedule_success_audit(self, admin_auth, monkeypatch, mock_audit_log): + prisma = _mock_prisma(monkeypatch) + prisma.db.litellm_uisettings.upsert.side_effect = RuntimeError("storage unavailable") + with pytest.raises(RuntimeError, match="storage unavailable"): + client.patch("/update/user_banner", json=PUBLISH_BODY) + mock_audit_log.assert_not_called() + + def test_audit_receives_previous_and_saved_banner(self, admin_auth, monkeypatch, mock_audit_log): + prisma = _mock_prisma(monkeypatch, record=SimpleNamespace(ui_settings=json.dumps(PUBLISHED_BANNER))) + response = client.patch("/update/user_banner", json=PUBLISH_BODY) + assert response.status_code == 200 + before, after, user = mock_audit_log.call_args.args + assert before.model_dump() == PUBLISHED_BANNER + assert after.model_dump() == response.json()["banner"] + assert user.user_id == "test-user" + assert ( + json.loads(prisma.db.litellm_uisettings.upsert.await_args.kwargs["data"]["create"]["ui_settings"]) + == after.model_dump() + ) + def test_rejects_non_admin(self, internal_user_auth, monkeypatch): mock_prisma = _mock_prisma(monkeypatch) response = client.patch("/update/user_banner", json=PUBLISH_BODY) diff --git a/tests/unit/test__db.py b/tests/unit/test__db.py new file mode 100644 index 00000000000..e6c18defd97 --- /dev/null +++ b/tests/unit/test__db.py @@ -0,0 +1,36 @@ +import asyncio +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm import _db + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [RuntimeError, asyncio.CancelledError]) +async def test_lifespan_closes_database_when_owner_fails(failure): + prisma: Final = MagicMock() + prisma.disconnect = AsyncMock() + database: Final = _db.from_prisma(prisma) + + async def fail(): + async with _db.lifespan(database) as current: + assert current is database + prisma.disconnect.assert_not_awaited() + raise failure("owner failed") + + with pytest.raises(failure): + await fail() + prisma.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_lifespan_without_database_preserves_owner_error(): + async def fail(): + async with _db.lifespan(None) as database: + assert database is None + raise RuntimeError("owner failed") + + with pytest.raises(RuntimeError, match="owner failed"): + await fail()